diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py index 21e38b54fe..c7f41d5d01 100644 --- a/studio/backend/core/inference/diffusion_precision.py +++ b/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) diff --git a/studio/backend/tests/test_diffusion_te_prequant.py b/studio/backend/tests/test_diffusion_te_prequant.py index 52aaeed83b..e8d3cbbe57 100644 --- a/studio/backend/tests/test_diffusion_te_prequant.py +++ b/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