Update app.py

app.py

@@ -460,39 +460,41 @@ def _rename_lora_keys_for_base_model(lora_sd: StateDict, base_keys: set[str]) ->
     """
     Rename LoRA state dict keys to match the base model's key format.
     
-    LoRA files from different sources use various key structures.
-    We need to normalize them to match the base model.
+    LoRA keys:    transformer_blocks.0.attn1.to_k.lora_A.weight
+    Base keys:    velocity_model.transformer_blocks.4.audio_attn1.to_out.0.weight
     
-    Common patterns:
-    - "diffusion_model.transformer.xxx" -> "xxx" (ComfyUI export)
-    - "diffusion_model.xxx" -> "xxx" (strip prefix)
-    - "transformer.xxx" -> "xxx" (strip module name)
+    We need to:
+    1. Add 'velocity_model.' prefix
+    2. Match block indices (LoRA block 0 might correspond to base block 4, etc.)
+    
+    Actually, LTX LoRA files typically have entries for ALL blocks.
+    The LoRA files from ComfyUI exports may have different structures.
+    We need to normalize to: velocity_model.transformer_blocks.X.attn.to_weight
     """
     renamed_sd = {}
+    
+    # First pass: identify the pattern
+    # LoRA keys like: transformer_blocks.0.attn1.to_k.lora_A.weight
+    # Need to become: velocity_model.transformer_blocks.0.attn1.to_k.weight
+    
     for key, tensor in lora_sd.sd.items():
         new_key = key
         
-        # Strip "diffusion_model." prefix
+        # Strip "diffusion_model." prefix if present
         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."):]
+        # Now add "velocity_model." prefix
+        if not new_key.startswith("velocity_model."):
+            new_key = "velocity_model." + new_key
         
-        # 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)
+        # Convert LoRA key suffix to base weight suffix
+        # .lora_A.weight -> .weight
+        # .lora_B.weight -> .weight
+        if new_key.endswith(".lora_A.weight"):
+            new_key = new_key[:-len(".lora_A.weight")] + ".weight"
+        elif new_key.endswith(".lora_B.weight"):
+            new_key = new_key[:-len(".lora_B.weight")] + ".weight"
         
         renamed_sd[new_key] = tensor
     
@@ -542,14 +544,18 @@ def build_fused_state_dict(
         # 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}")
+        print(f"[LoRA DEBUG] Before: {original_keys}")
+        print(f"[LoRA DEBUG] After:  {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]}")
+        # Check if renamed keys exist in base model
+        sample_renamed = list(sd_renamed.sd.keys())[0]
+        exists_in_base = sample_renamed in base_key_set
+        print(f"[LoRA DEBUG] Sample renamed key exists in base? {exists_in_base}")
+        if not exists_in_base:
+            print(f"[LoRA DEBUG] Checking for similar base keys...")
+            prefix = sample_renamed.rsplit(".", 1)[0]  # Remove .weight suffix
+            similar = [k for k in base_key_set if k.startswith(prefix[:30])]
+            print(f"[LoRA DEBUG] Similar base keys: {similar[:3]}")
     
     if progress_callback:
         progress_callback(0.3, "Extracting base transformer state dict")