fusion x lora

videox_fun/utils/lora_utils.py

@@ -391,6 +391,8 @@ def merge_lora(pipeline, lora_path, multiplier, device='cpu', dtype=torch.float3
             key = key.replace(".ffn.", "_ffn_")
             key = key.replace(".lora_A.default.", ".lora_down.")
             key = key.replace(".lora_B.default.", ".lora_up.")
+            key = key.replace(".lora_A.weight", ".lora_down.weight")
+            key = key.replace(".lora_B.weight", ".lora_up.weight")
         layer, elem = key.split('.', 1)
         updates[layer][elem] = value
 
@@ -512,6 +514,8 @@ def unmerge_lora(pipeline, lora_path, multiplier=1, device="cpu", dtype=torch.fl
             key = key.replace(".ffn.", "_ffn_")
             key = key.replace(".lora_A.default.", ".lora_down.")
             key = key.replace(".lora_B.default.", ".lora_up.")
+            key = key.replace(".lora_A.weight", ".lora_down.weight")
+            key = key.replace(".lora_B.weight", ".lora_up.weight")
         layer, elem = key.split('.', 1)
         updates[layer][elem] = value