Update app.py

app.py

@@ -527,6 +527,15 @@ def load_models() -> str | None:
                 rope_theta ** (torch.arange(0, 2 * rope_freq_dim, 2, dtype=torch.float32) / (2 * rope_freq_dim))
             )
 
+            # `_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()):
+                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)