update

videox_fun/utils/lora_utils.py

@@ -389,11 +389,28 @@ def merge_lora(pipeline, lora_path, multiplier, device='cpu', dtype=torch.float3
             key = key.replace(".self_attn.", "_self_attn_")
             key = key.replace(".cross_attn.", "_cross_attn_")
             key = key.replace(".ffn.", "_ffn_")
+            key = key.replace("text_embedding.", "text_embedding_")
+            key = key.replace("time_embedding.", "time_embedding_")
             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)
+        
+        if key.endswith(".lora_down.weight"):
+            layer = key[:-len(".lora_down.weight")]
+            elem = "lora_down.weight"
+        elif key.endswith(".lora_up.weight"):
+            layer = key[:-len(".lora_up.weight")]
+            elem = "lora_up.weight"
+        elif key.endswith(".alpha"):
+            layer = key[:-len(".alpha")]
+            elem = "alpha"
+        else:
+            continue
+
+        if layer.endswith("."):
+            layer = layer[:-1]
+            
         updates[layer][elem] = value
 
     sequential_cpu_offload_flag = False
@@ -512,11 +529,28 @@ def unmerge_lora(pipeline, lora_path, multiplier=1, device="cpu", dtype=torch.fl
             key = key.replace(".self_attn.", "_self_attn_")
             key = key.replace(".cross_attn.", "_cross_attn_")
             key = key.replace(".ffn.", "_ffn_")
+            key = key.replace("text_embedding.", "text_embedding_")
+            key = key.replace("time_embedding.", "time_embedding_")
             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)
+        
+        if key.endswith(".lora_down.weight"):
+            layer = key[:-len(".lora_down.weight")]
+            elem = "lora_down.weight"
+        elif key.endswith(".lora_up.weight"):
+            layer = key[:-len(".lora_up.weight")]
+            elem = "lora_up.weight"
+        elif key.endswith(".alpha"):
+            layer = key[:-len(".alpha")]
+            elem = "alpha"
+        else:
+            continue
+
+        if layer.endswith("."):
+            layer = layer[:-1]
+            
         updates[layer][elem] = value
 
     sequential_cpu_offload_flag = False