Update app.py

app.py

@@ -109,15 +109,16 @@ 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:
+def _convert_diffusion_model_lora(raw: dict, base_shapes: dict) -> 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.
+    `raw` maps original key -> tensor. `base_shapes` maps the *unwrapped* base model's parameter names to their
+    shapes — captured once before any adapter is attached, since `load_lora_adapter` wraps each target Linear in
+    a PEFT layer and renames its weight to `<name>.base_layer.weight`, so a live `transformer.state_dict()` call
+    after the first adapter attaches would no longer have `to_q.weight` etc. under their original names.
     """
     import re
 
-    base = transformer.state_dict()
     out = {}
     pattern = re.compile(r"^diffusion_model\.blocks\.(\d+)\.(attn|mlp)\.(\w+)\.(lora_[AB])\.weight$")
 
@@ -136,9 +137,9 @@ def _convert_diffusion_model_lora(raw: dict, transformer) -> dict:
                 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]
+                q_out = base_shapes[f"{prefix}attn.to_q.weight"][0]
+                k_out = base_shapes[f"{prefix}attn.to_k.weight"][0]
+                v_out = base_shapes[f"{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}), "
                     f"got {tensor.shape[0]}"
@@ -260,12 +261,15 @@ def load_models() -> str | None:
             from safetensors import safe_open
 
             try:
+                # Snapshot once, before either adapter attaches and wraps the target Linears — see the
+                # docstring on `_convert_diffusion_model_lora` for why this can't be read fresh per-file.
+                base_shapes = {k: tuple(v.shape) for k, v in pipe.transformer.state_dict().items()}
                 counts = {}
                 for name, filename in LORA_FILES.items():
                     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)
+                    converted = _convert_diffusion_model_lora(raw, base_shapes)
                     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.