Update app.py

app.py

@@ -407,6 +407,60 @@ def status() -> str:
         f"conditioner `{CONDITIONER_SPACE}`"
     )
 
+def _convert_full_checkpoint(raw: dict, base_shapes: dict) -> dict:
+    """Rename a `diffusion_model.blocks.*`-family (original-checkpoint) full transformer state dict onto
+    `MiniMaxH3Transformer3DModel`'s (`transformer_blocks.*`) naming — the full-weight sibling of
+    `_convert_diffusion_model_lora`'s renaming: no A/B factors, no network_alphas, no fc1 swap, just every real
+    weight tensor renamed (and, for the fused `qkv_proj`, split) onto its diffusers-side target. `base_shapes` is
+    the target model's own shapes — safe to read straight off a `torch.device("meta")`-constructed instance,
+    since shape is metadata, not data, and costs nothing to have before any real weights are loaded.
+    """
+    out: dict = {}
+
+    def rename_base(name: str) -> str:
+        if name.startswith("token_refiner.blocks."):
+            name = "token_refiner.refiner_blocks." + name[len("token_refiner.blocks."):]
+        elif name.startswith("blocks."):
+            name = "transformer_blocks." + name[len("blocks."):]
+        name = name.replace("final_layer.adaln_proj.linear", "norm_out.linear")
+        name = name.replace("final_layer.norm", "norm_out.norm")
+        name = name.replace("final_layer.video_out", "proj_out")
+        name = name.replace("final_layer.audio_out", "audio_proj_out")
+        name = name.replace(".attn.out_proj", ".attn.to_out.0")
+        name = name.replace(".attn.q_norm", ".attn.norm_q")
+        name = name.replace(".attn.k_norm", ".attn.norm_k")
+        name = name.replace(".mlp.fc2", ".ff.net.2")
+        name = name.replace(".mlp.fc1", ".ff.net.0.proj")
+        name = name.replace("video_patch_proj", "proj_in")
+        name = name.replace("audio_patch_proj", "audio_proj_in")
+        name = name.replace("condition_proj", "context_embedder")
+        name = name.replace("time_embedder.proj_in", "time_embedder.linear_1")
+        name = name.replace("time_embedder.proj_out", "time_embedder.linear_2")
+        return name
+
+    for key, tensor in raw.items():
+        if key == "rope.inv_freq":
+            # A registered buffer computed from `rope_theta`/`rope_freq_dim` at construction, never loaded —
+            # its presence here isn't a sign anything else is wrong.
+            continue
+
+        if key.endswith(".attn.qkv_proj.weight"):
+            prefix = key[: -len(".attn.qkv_proj.weight")]
+            renamed_prefix = rename_base(prefix)
+            q_out = base_shapes[f"{renamed_prefix}.attn.to_q.weight"][0]
+            k_out = base_shapes[f"{renamed_prefix}.attn.to_k.weight"][0]
+            v_out = base_shapes[f"{renamed_prefix}.attn.to_v.weight"][0]
+            assert tensor.shape[0] == q_out + k_out + v_out, (
+                f"{key}: expected {q_out + k_out + v_out} rows (q{q_out}+k{k_out}+v{v_out}), got {tensor.shape[0]}"
+            )
+            out[f"{renamed_prefix}.attn.to_q.weight"] = tensor[:q_out].clone()
+            out[f"{renamed_prefix}.attn.to_k.weight"] = tensor[q_out:q_out + k_out].clone()
+            out[f"{renamed_prefix}.attn.to_v.weight"] = tensor[q_out + k_out:].clone()
+            continue
+
+        out[rename_base(key)] = tensor
+
+    return out
 
 def load_models() -> str | None:
     """Load the denoising half at startup, plus FILM.
@@ -437,10 +491,12 @@ def load_models() -> str | None:
         if CUSTOM_TRANSFORMER_REPO:
             # `load_config` fetches only `transformer/config.json` (a few KB) — not the 61.7 GiB of weights
             # `load_components` would otherwise pull from MODEL_REPO. Constructed on `torch.device("meta")` so
-            # the architecture exists with no real memory behind it, then `load_state_dict(assign=True)`
-            # materializes real tensors straight from the finetune's own file — the only real allocation here.
+            # the architecture exists with no real memory behind it; `base_shapes` is read off that meta
+            # instance (shape is metadata, not data) purely so `_convert_full_checkpoint` knows the real
+            # to_q/to_k/to_v split points before any real weights exist. `load_state_dict(assign=True)`
+            # materializes real tensors straight from the converted dict — the only real allocation here.
             from huggingface_hub import hf_hub_download
-            from safetensors.torch import load_file
+            from safetensors import safe_open
 
             from diffusers.models import MiniMaxH3Transformer3DModel
 
@@ -449,8 +505,13 @@ def load_models() -> str | None:
             )
             with torch.device("meta"):
                 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)
-            custom_transformer.load_state_dict(load_file(custom_path), strict=True, assign=True)
+            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)
+            custom_transformer.load_state_dict(converted, strict=True, assign=True)
             pipe.update_components(transformer=custom_transformer)
             print(f"[gen] transformer replaced with {CUSTOM_TRANSFORMER_REPO}/{CUSTOM_TRANSFORMER_FILE}", flush=True)