Report the fp8-cast compute dtype without swapping the encoder class

The dtype override swapped encoder.__class__ to a dynamic subclass, which
breaks transformers' kwargs-based output recording: a fp8-cast
Qwen3VLModel stopped returning hidden_states and every krea-2 generation
with text_encoder_quant=fp8 crashed at encode_prompt (regression from the
HiDream TE4 change; caught by the krea hosted-TE live smoke). The
override is now a property shadowed on the ORIGINAL class that prefers a
per-instance compute-dtype attribute, so class identity is preserved and
uncast instances keep the stock behaviour. The idempotency test now pins
exact class identity and the uncast-sibling fallback.

studio/backend/core/inference/diffusion_precision.py

@@ -497,19 +497,15 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
     # Module.dtype reports the first floating parameter, which is now fp8 STORAGE; pipelines
     # derive tensor dtypes from encoder.dtype (Flux2 casts prompt embeds to it and feeds the
     # result to randn_tensor, which has no fp8 kernel; VLM pipelines cast pixel_values to it,
-    # racing the upcast hooks). The encoder computes in target.dtype, so report that.
+    # racing the upcast hooks). The encoder computes in target.dtype, so report that -- via a
+    # property shadowed on the ORIGINAL class reading a per-instance override. Swapping
+    # __class__ to a dynamic subclass instead breaks transformers' kwargs-based output
+    # recording (Qwen3VLModel returned hidden_states=None and krea-2 crashed at encode).
     compute_dtype = getattr(target, "dtype", None)
     try:
-        if compute_dtype is not None and not getattr(encoder, "_unsloth_te_dtype_override", False):
-            cls = type(encoder)
-            encoder.__class__ = type(
-                cls.__name__,
-                (cls,),
-                {
-                    "dtype": property(lambda self, _d = compute_dtype: _d),
-                    "_unsloth_te_dtype_override": True,
-                },
-            )
+        if compute_dtype is not None:
+            _install_dtype_override(type(encoder))
+            encoder._unsloth_te_compute_dtype = compute_dtype
         # Marks the cast COMPLETE (hooks fully installed), enabling the idempotent early return
         # above. Best-effort like the dtype override: a non-Module double without settable
         # attributes still counts as cast, it just re-casts on a repeat call.
@@ -518,6 +514,28 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
         pass
 
 
+def _install_dtype_override(cls: type) -> None:
+    """Shadow ``cls.dtype`` with a property preferring the per-instance compute-dtype
+    override ``_cast_fp8`` sets; instances without it keep the original behaviour. Class
+    identity is untouched, applied once per class."""
+    existing = cls.__dict__.get("dtype")
+    if getattr(getattr(existing, "fget", None), "_unsloth_te_dtype_override", False):
+        return
+    # The property object itself when accessed through the class (property.__get__(None, cls)).
+    original_fget = getattr(getattr(cls, "dtype", None), "fget", None)
+
+    def _dtype(self):
+        override = self.__dict__.get("_unsloth_te_compute_dtype")
+        if override is not None:
+            return override
+        if original_fget is not None:
+            return original_fget(self)
+        raise AttributeError("dtype")
+
+    _dtype._unsloth_te_dtype_override = True
+    cls.dtype = property(_dtype)
+
+
 def _has_layerwise_hooks(encoder: Any) -> bool:
     """True when any submodule already carries the diffusers layerwise-casting hook."""
     return _has_layerwise_casting(encoder)

studio/backend/tests/test_diffusion_te_prequant.py

@@ -441,7 +441,13 @@ def test_cast_fp8_is_idempotent_on_precast_encoder():
     # Module.dtype must report the COMPUTE dtype: pipelines derive tensor dtypes from it
     # (Flux2 feeds it to randn_tensor, which has no fp8 kernel).
     assert enc.dtype == torch.bfloat16
-    assert isinstance(enc, torch.nn.Sequential)
+    # EXACT class identity: a dynamic-subclass swap here broke transformers' kwargs-based
+    # output recording (Qwen3VLModel returned hidden_states=None; krea-2 crashed at encode).
+    assert type(enc) is torch.nn.Sequential
+    # An uncast sibling of the same (now property-patched) class keeps original behaviour.
+    sibling = torch.nn.Sequential(torch.nn.Linear(8, 8))
+    with pytest.raises(AttributeError):
+        sibling.dtype
     _cast_fp8(enc, target)  # must not raise
     assert enc[0].weight.dtype == torch.float8_e4m3fn
     assert enc.dtype == torch.bfloat16