Update app.py

app.py

@@ -526,20 +526,27 @@ def load_models() -> str | None:
             custom_transformer.rope.inv_freq = 1.0 / (
                 rope_theta ** (torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32) / (2 * rope_freq_dim))
             )
+            pipe.update_components(transformer=custom_transformer)
+            print(f"[gen] transformer replaced with {CUSTOM_TRANSFORMER_REPO}/{CUSTOM_TRANSFORMER_FILE}", flush=True)
 
+        pipe.load_components(dtype=torch.bfloat16)
+
+        if CUSTOM_TRANSFORMER_REPO:
+            # Moved to *after* `load_components(dtype=torch.bfloat16)` above, not before it: that call runs
+            # unconditionally and its `dtype=` argument applies to every component regardless of whether it was
+            # freshly fetched or already installed via `update_components` — doing this upcast beforehand had it
+            # silently re-cast straight back to bf16 one line later, which is why the first attempt at this fix
+            # had no visible effect at all despite being otherwise correct.
+            #
             # `_keep_in_fp32_modules` is normally enforced by `from_pretrained`'s own post-load dtype pass — a
             # step this manual meta-device + `load_state_dict` path never goes through. The official checkpoint
             # genuinely ships these modules as float32 while everything else is bfloat16; if the finetune's file
-            # is uniformly bf16 (common for a community single-file export), `assign=True` just adopted that
-            # dtype for these layers too, silently dropping precision the model actually needs to run correctly.
-            for name, tensor in list(custom_transformer.named_parameters()) + list(custom_transformer.named_buffers()):
+            # is uniformly bf16 (common for a community single-file export), it adopted that dtype for these
+            # layers too, silently dropping precision the model actually needs to run correctly.
+            for name, tensor in list(pipe.transformer.named_parameters()) + list(pipe.transformer.named_buffers()):
                 if any(keep in name for keep in MiniMaxH3Transformer3DModel._keep_in_fp32_modules):
                     tensor.data = tensor.data.to(torch.float32)
 
-            pipe.update_components(transformer=custom_transformer)
-            print(f"[gen] transformer replaced with {CUSTOM_TRANSFORMER_REPO}/{CUSTOM_TRANSFORMER_FILE}", flush=True)
-
-        pipe.load_components(dtype=torch.bfloat16)
         pipe.transformer.set_attention_backend(ATTENTION)
 
         # --- Diagnostic: dump the LoRA files' key names/shapes and the transformer's own shapes to the Space