From 569220fc9a97a346b5b2f3f0405b7c355a520cd4 Mon Sep 17 00:00:00 2001
From: Daniel Han <danielhanchen@gmail.com>
Date: Sat, 18 Jul 2026 07:27:05 +0000
Subject: [PATCH] Key the fp8 cast idempotency on an explicit completion marker

Hook presence alone cannot distinguish a legitimately pre-cast text
encoder from leftover hooks after a cast that failed mid-pass, so the
early return now requires the completion marker _cast_fp8 sets once the
hooks are fully installed. Leftover partial state keeps failing closed.
Also tolerates non-Module encoder doubles in the hook probe and the
dtype override.
---
 .../core/inference/diffusion_precision.py     | 38 ++++++++++++-------
 1 file changed, 25 insertions(+), 13 deletions(-)

diff --git a/studio/backend/core/inference/diffusion_precision.py b/studio/backend/core/inference/diffusion_precision.py
index 7b19e9fa98..a97c7adb3e 100644
--- a/studio/backend/core/inference/diffusion_precision.py
+++ b/studio/backend/core/inference/diffusion_precision.py
@@ -248,8 +248,10 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
 
     # Idempotent: a pre-cast encoder (diffusion_te_prequant) arrives with the layerwise hooks
     # already installed, and re-registering the same hook name raises -- which would make
-    # quantize_text_encoders report the (actually engaged) cast as failed.
-    if _has_layerwise_hooks(encoder):
+    # quantize_text_encoders report the (actually engaged) cast as failed. Keyed on the explicit
+    # completion marker this function sets, NOT on hook presence alone: leftover hooks from a
+    # cast that failed mid-pass must still fail closed, not read as "already cast".
+    if getattr(encoder, "_unsloth_te_cast_complete", False) and _has_layerwise_hooks(encoder):
         return
 
     # Layerwise casting stores each leaf's weights in fp8 and upcasts per forward. Two things on a
@@ -293,21 +295,31 @@ def _cast_fp8(encoder: Any, target: Any) -> None:
     # 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.
     compute_dtype = getattr(target, "dtype", None)
-    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,
-            },
-        )
+    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,
+                },
+            )
+        # 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.
+        encoder._unsloth_te_cast_complete = True
+    except Exception:  # noqa: BLE001 — real HF encoders are heap-type nn.Modules; only doubles fail
+        pass
 
 
 def _has_layerwise_hooks(encoder: Any) -> bool:
     """True when any submodule already carries the diffusers layerwise-casting hook."""
-    for module in encoder.modules():
+    modules = getattr(encoder, "modules", None)
+    if not callable(modules):
+        return False
+    for module in modules():
         registry = getattr(module, "_diffusers_hook", None)
         get_hook = getattr(registry, "get_hook", None)
         if callable(get_hook) and get_hook("layerwise_casting") is not None: