Update app.py

app.py

@@ -42,9 +42,9 @@ GPU_SIZE = os.environ.get("H3_GPU_SIZE", "xlarge")
 # default, which reproduces the official weights exactly. Confirmed diffusers-native key naming (not ComfyUI,
 # not pruned-AdaLN) via `xal2077/PinkCherry_MiniMax-H3-Demo`'s own working `load_state_dict(strict=True)` call
 # against an earlier release of the same lineage.
-CUSTOM_TRANSFORMER_REPO = os.environ.get("H3_CUSTOM_TRANSFORMER_REPO", "SexGod1979/PinkCherry_MiniMax-H3")
+CUSTOM_TRANSFORMER_REPO = os.environ.get("H3_CUSTOM_TRANSFORMER_REPO", "Comfy-Org/MiniMax-H3")
 CUSTOM_TRANSFORMER_FILE = os.environ.get(
-    "H3_CUSTOM_TRANSFORMER_FILE", "v1-final-fl2va/PinkCherry_v1_bf16_fla2va_H3.safetensors"
+    "H3_CUSTOM_TRANSFORMER_FILE", "diffusion_models/minimax_h3_fl2va_bf16.safetensors"
 )
 
 LORA_REPO = os.environ.get("H3_LORA_REPO", "dagloop5/LoRA")
@@ -507,15 +507,18 @@ 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()}
 
-            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.
+            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_ALL", "0") == "1":
+                # The meta+assign mechanism is proven correct (H3_DEBUG_OFFICIAL_VIA_META produced coherent,
+                # matching output using the official weights through this exact path), and two individual
+                # tensors (`norm1.weight`, `attn.to_q.weight`) are already proven bit-exact — so the remaining
+                # bug has to be in some *other* tensor `_convert_full_checkpoint` handles differently, not yet
+                # individually checked. This downloads the official transformer once (same cost as the mechanism
+                # test) and diffs every key against `converted`, rather than guessing which one to spot-check.
                 import json as _json
 
                 index_path = hf_hub_download(
@@ -530,34 +533,24 @@ def load_models() -> str | None:
                     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)
+
+                mismatches = []
+                for key, official_tensor in official_state_dict.items():
+                    converted_tensor = converted.get(key)
+                    if converted_tensor is None:
+                        mismatches.append((key, "missing from converted"))
+                        continue
+                    if converted_tensor.shape != official_tensor.shape:
+                        mismatches.append((key, f"shape {tuple(converted_tensor.shape)} != {tuple(official_tensor.shape)}"))
+                        continue
+                    if not torch.allclose(converted_tensor.float(), official_tensor.float(), atol=1e-3):
+                        diff = (converted_tensor.float() - official_tensor.float()).abs().max().item()
+                        mismatches.append((key, f"values differ, max abs diff {diff:.6f}"))
+                print(f"[debug-compare-all] checked {len(official_state_dict)} keys, {len(mismatches)} mismatches", flush=True)
+                for key, reason in mismatches[:30]:
+                    print(f"[debug-compare-all]   {key}: {reason}", 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