Update app.py

app.py

@@ -154,6 +154,52 @@ def load_models() -> str | None:
         pipe.load_components(dtype=torch.bfloat16)
         pipe.transformer.set_attention_backend(ATTENTION)
 
+        pipe.load_components(dtype=torch.bfloat16)
+        pipe.transformer.set_attention_backend(ATTENTION)
+
+        # --- Diagnostic: dump the LoRA files' key names/shapes and the transformer's own shapes to the Space
+        # logs, so the exact rename map can be worked out without a notebook or shell. Set H3_LORA_DEBUG=0 in
+        # the Space's env vars to silence this once you're done, or just delete this block later.
+        if os.environ.get("H3_LORA_DEBUG", "1") != "0" and LORA_REPO.lower() not in ("", "off", "none"):
+            from huggingface_hub import hf_hub_download
+            from safetensors import safe_open
+
+            for filename in LORA_FILES.values():
+                try:
+                    path = hf_hub_download(LORA_REPO, filename)
+                    with safe_open(path, framework="pt") as handle:
+                        keys = sorted(handle.keys())
+                    print(f"[lora-debug] {filename}: {len(keys)} keys", flush=True)
+                    for k in keys[:40]:
+                        print(f"[lora-debug]   {k}  {tuple(handle.get_slice(k).get_shape())}", flush=True)
+                    if len(keys) > 40:
+                        print(f"[lora-debug]   ... and {len(keys) - 40} more", flush=True)
+                except Exception as error:
+                    print(f"[lora-debug] failed to inspect {filename}: {error}", flush=True)
+
+            block0 = {
+                k: tuple(v.shape)
+                for k, v in pipe.transformer.state_dict().items()
+                if k.startswith("transformer_blocks.0.")
+            }
+            print(f"[lora-debug] transformer_blocks.0.* ({len(block0)} keys):", flush=True)
+            for k, shape in sorted(block0.items()):
+                print(f"[lora-debug]   {k}  {shape}", flush=True)
+            norm_out = {
+                k: tuple(v.shape) for k, v in pipe.transformer.state_dict().items() if k.startswith("norm_out.")
+            }
+            print(f"[lora-debug] norm_out.* ({len(norm_out)} keys):", flush=True)
+            for k, shape in sorted(norm_out.items()):
+                print(f"[lora-debug]   {k}  {shape}", flush=True)
+
+        if LORA_REPO.lower() not in ("", "off", "none"):
+            from peft.tuners.tuners_utils import BaseTunerLayer
+
+            try:
+                counts = {}
+                for name, filename in LORA_FILES.items():
+                    ...
+
         # Approach B: attach both LoRA adapters as PEFT layers on the transformer, 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.