Update app.py

app.py

@@ -17,6 +17,8 @@ import time
 import traceback
 from functools import cache
 
+import torch
+
 # Before anything that could initialize CUDA: `import spaces` patches `torch.cuda` so the 72 GiB load can happen at
 # startup rather than on GPU time.
 import spaces
@@ -103,6 +105,61 @@ def lower_duration_floor(seconds: float = MIN_UI_DURATION) -> None:
 
     MiniMaxH3ModularPipeline.min_duration = property(lambda self: float(seconds))
 
+def _convert_diffusion_model_lora(raw: dict, transformer) -> dict:
+    """Rename a `diffusion_model.blocks.*` (original-checkpoint) LoRA state dict onto
+    `MiniMaxH3Transformer3DModel`'s (`transformer_blocks.*`) naming, so `load_lora_adapter` can attach it.
+    `raw` maps original key -> tensor. `transformer` is the already-loaded base model, used only to read its
+    real to_q/to_k/to_v output sizes for the qkv split, rather than assuming an even three-way split.
+    """
+    import re
+
+    base = transformer.state_dict()
+    out = {}
+    pattern = re.compile(r"^diffusion_model\.blocks\.(\d+)\.(attn|mlp)\.(\w+)\.(lora_[AB])\.weight$")
+
+    for key, tensor in raw.items():
+        match = pattern.match(key)
+        if not match:
+            print(f"[lora-convert] skipping unrecognized key: {key}", flush=True)
+            continue
+        block, kind, leaf, ab = match.groups()
+        prefix = f"transformer_blocks.{block}."
+
+        if kind == "attn" and leaf == "qkv_proj":
+            if ab == "lora_A":
+                # Shared low-rank input side — identical for q, k, v.
+                out[f"{prefix}attn.to_q.{ab}.weight"] = tensor
+                out[f"{prefix}attn.to_k.{ab}.weight"] = tensor
+                out[f"{prefix}attn.to_v.{ab}.weight"] = tensor
+            else:
+                q_out = base[f"{prefix}attn.to_q.weight"].shape[0]
+                k_out = base[f"{prefix}attn.to_k.weight"].shape[0]
+                v_out = base[f"{prefix}attn.to_v.weight"].shape[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}), "
+                    f"got {tensor.shape[0]}"
+                )
+                out[f"{prefix}attn.to_q.{ab}.weight"] = tensor[:q_out].clone()
+                out[f"{prefix}attn.to_k.{ab}.weight"] = tensor[q_out:q_out + k_out].clone()
+                out[f"{prefix}attn.to_v.{ab}.weight"] = tensor[q_out + k_out:].clone()
+
+        elif kind == "attn" and leaf == "out_proj":
+            out[f"{prefix}attn.to_out.0.{ab}.weight"] = tensor
+
+        elif kind == "mlp" and leaf == "fc1":
+            if ab == "lora_B" and SWAP_FC1_HALVES:
+                # Some checkpoints store the SwiGLU gate/value halves in the opposite order diffusers expects.
+                half = tensor.shape[0] // 2
+                tensor = torch.cat([tensor[half:], tensor[:half]], dim=0)
+            out[f"{prefix}ff.net.0.proj.{ab}.weight"] = tensor
+
+        elif kind == "mlp" and leaf == "fc2":
+            out[f"{prefix}ff.net.2.{ab}.weight"] = tensor
+
+        else:
+            print(f"[lora-convert] no mapping for {key}, skipping", flush=True)
+
+    return out
 
 PIPE = None
 FILM = None
@@ -189,18 +246,23 @@ def load_models() -> str | None:
             for k, shape in sorted(norm_out.items()):
                 print(f"[lora-debug]   {k}  {shape}", flush=True)
 
-        # Approach B: attach both LoRA adapters as PEFT layers on the transformer, inactive (weight 0) until a
+        # Approach B: convert each LoRA from its original `diffusion_model.blocks.*` naming onto this
+        # transformer's `transformer_blocks.*` naming, then attach as PEFT layers, inactive (weight 0) until a
         # request asks for them. `load_lora_adapter` is the model-level loader (`PeftAdapterMixin`), used because
         # `MiniMaxH3ModularPipeline` has no pipeline-level `load_lora_weights` of its own.
         if LORA_REPO.lower() not in ("", "off", "none"):
+            from huggingface_hub import hf_hub_download
             from peft.tuners.tuners_utils import BaseTunerLayer
+            from safetensors import safe_open
 
             try:
                 counts = {}
                 for name, filename in LORA_FILES.items():
-                    pipe.transformer.load_lora_adapter(
-                        LORA_REPO, weight_name=filename, adapter_name=name, prefix=None
-                    )
+                    path = hf_hub_download(LORA_REPO, filename)
+                    with safe_open(path, framework="pt") as handle:
+                        raw = {k: handle.get_tensor(k) for k in handle.keys()}
+                    converted = _convert_diffusion_model_lora(raw, pipe.transformer)
+                    pipe.transformer.load_lora_adapter(converted, adapter_name=name, prefix=None)
                     # `load_lora_adapter` warns-and-continues on a zero-key match instead of raising, so count
                     # matched layers ourselves and fail loudly if a file attached nothing.
                     counts[name] = sum(
@@ -210,8 +272,8 @@ def load_models() -> str | None:
                     )
                     if counts[name] == 0:
                         raise RuntimeError(
-                            f"'{filename}' matched 0 target modules on MiniMaxH3Transformer3DModel — its LoRA "
-                            f"keys don't line up with this transformer's naming. Inspect its safetensors keys."
+                            f"'{filename}' converted but matched 0 target modules on "
+                            f"MiniMaxH3Transformer3DModel — the rename map needs another look."
                         )
                 pipe.transformer.set_adapters(list(LORA_FILES), weights=[0.0] * len(LORA_FILES))
                 LORA_STATUS = "LoRAs loaded: " + ", ".join(