From fab50bd721d53ce9cfa306fdbb749091cd325bb4 Mon Sep 17 00:00:00 2001
From: "huangkunzhe.hkz" <huangkunzhe.hkz@alibaba-inc.com>
Date: Mon, 9 Jun 2025 11:17:22 +0800
Subject: [PATCH] Support loading Wan-Fun LoRAs

---
 videox_fun/utils/lora_utils.py | 70 ++++++++++++++++++----------------
 1 file changed, 38 insertions(+), 32 deletions(-)

diff --git a/videox_fun/utils/lora_utils.py b/videox_fun/utils/lora_utils.py
index 585f0a8..6ff52ba 100755
--- a/videox_fun/utils/lora_utils.py
+++ b/videox_fun/utils/lora_utils.py
@@ -399,24 +399,27 @@ def merge_lora(pipeline, lora_path, multiplier, device='cpu', dtype=torch.float3
             curr_layer = curr_layer.__getattr__("_".join(layer_infos[1:]))
         except Exception:
             temp_name = layer_infos.pop(0)
-            while len(layer_infos) > -1:
-                try:
-                    curr_layer = curr_layer.__getattr__(temp_name + "_" + "_".join(layer_infos))
-                    break
-                except Exception:
+            try:
+                while len(layer_infos) > -1:
                     try:
-                        curr_layer = curr_layer.__getattr__(temp_name)
-                        if len(layer_infos) > 0:
-                            temp_name = layer_infos.pop(0)
-                        elif len(layer_infos) == 0:
-                            break
+                        curr_layer = curr_layer.__getattr__(temp_name + "_" + "_".join(layer_infos))
+                        break
                     except Exception:
-                        if len(layer_infos) == 0:
-                            print('Error loading layer')
-                        if len(temp_name) > 0:
-                            temp_name += "_" + layer_infos.pop(0)
-                        else:
-                            temp_name = layer_infos.pop(0)
+                        try:
+                            curr_layer = curr_layer.__getattr__(temp_name)
+                            if len(layer_infos) > 0:
+                                temp_name = layer_infos.pop(0)
+                            elif len(layer_infos) == 0:
+                                break
+                        except Exception:
+                            if len(layer_infos) == 0:
+                                print(f'Error loading layer: {layer}')
+                            if len(temp_name) > 0:
+                                temp_name += "_" + layer_infos.pop(0)
+                            else:
+                                temp_name = layer_infos.pop(0)
+            except Exception:
+                continue
 
         origin_dtype = curr_layer.weight.data.dtype
         origin_device = curr_layer.weight.data.device
@@ -472,24 +475,27 @@ def unmerge_lora(pipeline, lora_path, multiplier=1, device="cpu", dtype=torch.fl
             curr_layer = curr_layer.__getattr__("_".join(layer_infos[1:]))
         except Exception:
             temp_name = layer_infos.pop(0)
-            while len(layer_infos) > -1:
-                try:
-                    curr_layer = curr_layer.__getattr__(temp_name + "_" + "_".join(layer_infos))
-                    break
-                except Exception:
+            try:
+                while len(layer_infos) > -1:
                     try:
-                        curr_layer = curr_layer.__getattr__(temp_name)
-                        if len(layer_infos) > 0:
-                            temp_name = layer_infos.pop(0)
-                        elif len(layer_infos) == 0:
-                            break
+                        curr_layer = curr_layer.__getattr__(temp_name + "_" + "_".join(layer_infos))
+                        break
                     except Exception:
-                        if len(layer_infos) == 0:
-                            print('Error loading layer')
-                        if len(temp_name) > 0:
-                            temp_name += "_" + layer_infos.pop(0)
-                        else:
-                            temp_name = layer_infos.pop(0)
+                        try:
+                            curr_layer = curr_layer.__getattr__(temp_name)
+                            if len(layer_infos) > 0:
+                                temp_name = layer_infos.pop(0)
+                            elif len(layer_infos) == 0:
+                                break
+                        except Exception:
+                            if len(layer_infos) == 0:
+                                print(f'Error loading layer: {layer}')
+                            if len(temp_name) > 0:
+                                temp_name += "_" + layer_infos.pop(0)
+                            else:
+                                temp_name = layer_infos.pop(0)
+            except Exception:
+                continue
 
         origin_dtype = curr_layer.weight.data.dtype
         origin_device = curr_layer.weight.data.device