Update app.py

app.py

@@ -456,24 +456,44 @@ def load_lora_into_cache(lora_path: str) -> StateDict:
     print(f"[LoRA] Cached {len(tensors)} tensors from {os.path.basename(lora_path)}")
     return state_dict
 
-def _rename_lora_keys_for_base_model(lora_sd: StateDict) -> StateDict:
+def _rename_lora_keys_for_base_model(lora_sd: StateDict, base_keys: set[str]) -> StateDict:
     """
     Rename LoRA state dict keys to match the base model's key format.
     
-    LoRA files typically use ComfyUI-style keys like:
-        "diffusion_model.layer1.blocks.0.lora_A.weight"
+    LoRA files from different sources use various key structures.
+    We need to normalize them to match the base model.
     
-    But the base model uses keys like:
-        "layer1.blocks.0.lora_A.weight"
-    
-    This function strips the "diffusion_model." prefix to match them.
+    Common patterns:
+    - "diffusion_model.transformer.xxx" -> "xxx" (ComfyUI export)
+    - "diffusion_model.xxx" -> "xxx" (strip prefix)
+    - "transformer.xxx" -> "xxx" (strip module name)
     """
     renamed_sd = {}
     for key, tensor in lora_sd.sd.items():
-        if key.startswith("diffusion_model."):
-            new_key = key[len("diffusion_model."):]
-        else:
-            new_key = key
+        new_key = key
+        
+        # Strip "diffusion_model." prefix
+        if new_key.startswith("diffusion_model."):
+            new_key = new_key[len("diffusion_model."):]
+        
+        # If key still doesn't match any base keys, try stripping "transformer." prefix
+        if new_key not in base_keys and new_key.startswith("transformer."):
+            new_key = new_key[len("transformer."):]
+        
+        # If it's a LoRA key, try to match the base weight key
+        if ".lora_A.weight" in new_key or ".lora_B.weight" in new_key:
+            base_path = new_key.replace(".lora_A.weight", ".weight").replace(".lora_B.weight", ".weight")
+            if base_path not in base_keys:
+                parts = base_path.split(".")
+                if parts[0] == "transformer":
+                    alt_path = ".".join(parts[1:])
+                    if alt_path in base_keys:
+                        new_key = new_key.replace(base_path, alt_path)
+                elif parts[0] == "layer" and len(parts) > 1:
+                    alt_path = ".".join(parts[1:])
+                    if alt_path in base_keys:
+                        new_key = new_key.replace(base_path, alt_path)
+        
         renamed_sd[new_key] = tensor
     
     return StateDict(
@@ -507,14 +527,29 @@ def build_fused_state_dict(
     if progress_callback:
         progress_callback(0.1, "Loading LoRA state dicts into memory")
     
-    # Step 1: Load all LoRA state dicts (uses cache after first load)
-    # and rename keys to match base model's key format
+    # Step 1: Load all LoRA state dicts and rename keys to match base model
     lora_sd_with_strengths = []
+    
+    # Get base key set for matching
+    base_dict = base_transformer.state_dict()
+    base_key_set = set(base_dict.keys())
+    print(f"[LoRA DEBUG] Total base model keys: {len(base_key_set)}")
+    
     for lora_path, strength in lora_configs:
         sd = load_lora_into_cache(lora_path)
-        # Apply key renaming to strip "diffusion_model." prefix
-        sd_renamed = _rename_lora_keys_for_base_model(sd)
+        sd_renamed = _rename_lora_keys_for_base_model(sd, base_key_set)
+        
+        # Show before/after for first few keys
+        original_keys = list(sd.sd.keys())[:3]
+        renamed_keys = list(sd_renamed.sd.keys())[:3]
+        print(f"[LoRA DEBUG] Before rename: {original_keys}")
+        print(f"[LoRA DEBUG] After rename:  {renamed_keys}")
+        
         lora_sd_with_strengths.append(LoraStateDictWithStrength(sd_renamed, float(strength)))
+        print(f"[LoRA] Loaded and renamed {len(sd_renamed.sd)} keys from {os.path.basename(lora_path)}")
+
+    # Debug: Check LoRA key matching
+    print(f"[LoRA DEBUG] Sample base keys: {list(base_key_set)[:10]}")
     
     if progress_callback:
         progress_callback(0.3, "Extracting base transformer state dict")
@@ -577,6 +612,14 @@ def on_prepare_loras_click(
     Only runs on button click, NOT on slider change.
     """
     global current_lora_key, FUSED_CACHE
+
+    # Debug: Verify transformer consistency
+    ledger_transformer = ledger.transformer()
+    pipeline_transformer = pipeline.model_ledger.transformer()
+    print(f"[LoRA DEBUG] ledger.transformer() id: {id(ledger_transformer)}")
+    print(f"[LoRA DEBUG] pipeline.model_ledger.transformer() id: {id(pipeline_transformer)}")
+    print(f"[LoRA DEBUG] Same object? {ledger_transformer is pipeline_transformer}")
+    print(f"[LoRA DEBUG] _transformer id: {id(_transformer)}")
     
     # Compute the cache key for this combination of strengths
     key, _ = _make_lora_key(