Update app.py

app.py

@@ -456,6 +456,32 @@ 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:
+    """
+    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"
+    
+    But the base model uses keys like:
+        "layer1.blocks.0.lora_A.weight"
+    
+    This function strips the "diffusion_model." prefix to match them.
+    """
+    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
+        renamed_sd[new_key] = tensor
+    
+    return StateDict(
+        sd=renamed_sd,
+        device=lora_sd.device,
+        size=lora_sd.size,
+        dtype=lora_sd.dtype
+    )
 
 def build_fused_state_dict(
     base_transformer,
@@ -482,10 +508,13 @@ def build_fused_state_dict(
         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
     lora_sd_with_strengths = []
     for lora_path, strength in lora_configs:
         sd = load_lora_into_cache(lora_path)
-        lora_sd_with_strengths.append(LoraStateDictWithStrength(sd, float(strength)))
+        # Apply key renaming to strip "diffusion_model." prefix
+        sd_renamed = _rename_lora_keys_for_base_model(sd)
+        lora_sd_with_strengths.append(LoraStateDictWithStrength(sd_renamed, float(strength)))
     
     if progress_callback:
         progress_callback(0.3, "Extracting base transformer state dict")