From 164a6bbebd1b2259dbfa752a93f8a1a1307dcbd1 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Sat, 13 Dec 2025 00:54:22 +0200
Subject: [PATCH] Restore LoRA step scheduling functionality

---
 custom_linear.py          | 5 +++++
 wanvideo/modules/model.py | 6 ++----
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/custom_linear.py b/custom_linear.py
index eb6e8b0..b80be20 100644
--- a/custom_linear.py
+++ b/custom_linear.py
@@ -264,6 +264,11 @@ class CustomLinear(nn.Linear):
         del weight, input, bias
         return out
 
+def update_lora_step(module, step):
+    for name, submodule in module.named_modules():
+        if isinstance(submodule, CustomLinear) and hasattr(submodule, "_step"):
+            submodule._step.fill_(step)
+
 def remove_lora_from_module(module):
     for name, submodule in module.named_modules():
         if hasattr(submodule, "lora_diffs"):
diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index f0b5216..90c5628 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -21,6 +21,7 @@ from ...utils import log, get_module_memory_mb
 from ...cache_methods.cache_methods import TeaCacheState, MagCacheState, EasyCacheState, relative_l1_distance
 from ...multitalk.multitalk import get_attn_map_with_target
 from ...echoshot.echoshot import rope_apply_z, rope_apply_c, rope_apply_echoshot
+from ...custom_linear import update_lora_step
 
 from ...MTV.mtv import apply_rotary_emb
 from comfy.ldm.flux.math import apply_rope1 as apply_rope_comfy1
@@ -2251,10 +2252,7 @@ class WanModel(torch.nn.Module):
                 ip_scale = fantasy_portrait_input.get("strength", 1.0)
 
         if self.lora_scheduling_enabled:
-            for name, submodule in self.named_modules():
-                if isinstance(submodule, nn.Linear):
-                    if hasattr(submodule, 'step'):
-                        submodule.step = current_step
+            update_lora_step(self, current_step)
 
         # lynx
         lynx_x_ip = lynx_ref_feature = lynx_ref_buffer = lynx_ref_feature_extractor = None