Update app.py

app.py

@@ -44,7 +44,7 @@ GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
 # against an earlier release of the same lineage.
 CUSTOM_TRANSFORMER_REPO = os.environ.get("H3_CUSTOM_TRANSFORMER_REPO", "SexGod1979/PinkCherry_MiniMax-H3")
 CUSTOM_TRANSFORMER_FILE = os.environ.get(
-    "H3_CUSTOM_TRANSFORMER_FILE", "v1-final-fl2va/PinkCherry_v1_fp16_fla2va_H3.safetensors"
+    "H3_CUSTOM_TRANSFORMER_FILE", "v1-final-fl2va/PinkCherry_v1_bf16_fla2va_H3.safetensors"
 )
 
 LORA_REPO = os.environ.get("H3_LORA_REPO", "dagloop5/LoRA")
@@ -507,27 +507,57 @@ def load_models() -> str | None:
                 custom_transformer = MiniMaxH3Transformer3DModel.from_config(config)
             base_shapes = {k: tuple(v.shape) for k, v in custom_transformer.state_dict().items()}
 
-            custom_path = hf_hub_download(CUSTOM_TRANSFORMER_REPO, CUSTOM_TRANSFORMER_FILE)
-            with safe_open(custom_path, framework="pt") as handle:
-                raw = {k: handle.get_tensor(k) for k in handle.keys()}
-            converted = _convert_full_checkpoint(raw, base_shapes)
-            if os.environ.get("H3_DEBUG_COMPARE", "0") == "1":
-                from huggingface_hub import hf_hub_download as _dl
+            if os.environ.get("H3_DEBUG_OFFICIAL_VIA_META", "0") == "1":
+                # Isolates the meta-device + `assign=True` loading mechanism from `_convert_full_checkpoint`'s
+                # own conversion logic: loads the official, diffusers-native shards — no renaming, no splitting,
+                # nothing converted at all — through the exact same path the custom-checkpoint loader uses. If
+                # this *also* produces the same corrupted output, the bug is in the mechanism itself (meta
+                # construction, `assign=True`, the fp32 upcast, or something in between), independent of
+                # anything the conversion does, since no conversion happens here at all. If this loads clean,
+                # the bug is narrowed back to `_convert_full_checkpoint` specifically, despite everything
+                # checked against it so far.
                 import json as _json
-            
-                index_path = _dl(MODEL_REPO, "transformer/diffusion_pytorch_model.safetensors.index.json")
+
+                index_path = hf_hub_download(
+                    MODEL_REPO, "transformer/diffusion_pytorch_model.safetensors.index.json"
+                )
                 with open(index_path) as handle:
                     index = _json.load(handle)
-                probe_key = "transformer_blocks.0.attn.to_q.weight"
-                shard_name = index["weight_map"][probe_key]
-                shard_path = _dl(MODEL_REPO, f"transformer/{shard_name}")
-                with safe_open(shard_path, framework="pt") as official_handle:
-                    official_tensor = official_handle.get_tensor(probe_key)
-                converted_tensor = converted[probe_key]
-                print(f"[debug-compare] official {probe_key}: mean={official_tensor.float().mean():.6f} std={official_tensor.float().std():.6f} first5={official_tensor.flatten()[:5].tolist()}", flush=True)
-                print(f"[debug-compare] converted {probe_key}: mean={converted_tensor.float().mean():.6f} std={converted_tensor.float().std():.6f} first5={converted_tensor.flatten()[:5].tolist()}", flush=True)
-                print(f"[debug-compare] allclose: {torch.allclose(official_tensor, converted_tensor)}", flush=True)
-            custom_transformer.load_state_dict(converted, strict=True, assign=True)
+                shard_names = sorted(set(index["weight_map"].values()))
+                official_state_dict = {}
+                for shard_name in shard_names:
+                    shard_path = hf_hub_download(MODEL_REPO, f"transformer/{shard_name}")
+                    with safe_open(shard_path, framework="pt") as shard_handle:
+                        for key in shard_handle.keys():
+                            official_state_dict[key] = shard_handle.get_tensor(key)
+                custom_transformer.load_state_dict(official_state_dict, strict=True, assign=True)
+                print(
+                    "[gen] loaded the OFFICIAL weights via the meta+assign path, bypassing "
+                    "_convert_full_checkpoint entirely — diagnostic only",
+                    flush=True,
+                )
+            else:
+                custom_path = hf_hub_download(CUSTOM_TRANSFORMER_REPO, CUSTOM_TRANSFORMER_FILE)
+                with safe_open(custom_path, framework="pt") as handle:
+                    raw = {k: handle.get_tensor(k) for k in handle.keys()}
+                converted = _convert_full_checkpoint(raw, base_shapes)
+                if os.environ.get("H3_DEBUG_COMPARE", "0") == "1":
+                    from huggingface_hub import hf_hub_download as _dl
+                    import json as _json
+
+                    index_path = _dl(MODEL_REPO, "transformer/diffusion_pytorch_model.safetensors.index.json")
+                    with open(index_path) as handle:
+                        index = _json.load(handle)
+                    probe_key = "transformer_blocks.0.attn.to_q.weight"
+                    shard_name = index["weight_map"][probe_key]
+                    shard_path = _dl(MODEL_REPO, f"transformer/{shard_name}")
+                    with safe_open(shard_path, framework="pt") as official_handle:
+                        official_tensor = official_handle.get_tensor(probe_key)
+                    converted_tensor = converted[probe_key]
+                    print(f"[debug-compare] official {probe_key}: mean={official_tensor.float().mean():.6f} std={official_tensor.float().std():.6f} first5={official_tensor.flatten()[:5].tolist()}", flush=True)
+                    print(f"[debug-compare] converted {probe_key}: mean={converted_tensor.float().mean():.6f} std={converted_tensor.float().std():.6f} first5={converted_tensor.flatten()[:5].tolist()}", flush=True)
+                    print(f"[debug-compare] allclose: {torch.allclose(official_tensor, converted_tensor)}", flush=True)
+                custom_transformer.load_state_dict(converted, strict=True, assign=True)
 
             # `rope.inv_freq` is a *non-persistent* buffer (`persistent=False` in `MiniMaxH3RotaryPosEmbed`) —
             # excluded from `state_dict()` entirely, which is exactly why `strict=True` above never complained