From 8257cd1f8abaa6504248b946f31c5173c0228b3d Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Thu, 17 Apr 2025 01:18:42 +0300
Subject: [PATCH] Fix TeaCache end step

---
 wanvideo/modules/model.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/wanvideo/modules/model.py b/wanvideo/modules/model.py
index c19be23..33e6b3d 100644
--- a/wanvideo/modules/model.py
+++ b/wanvideo/modules/model.py
@@ -1157,7 +1157,7 @@ class WanModel(ModelMixin, ConfigMixin):
                 if b <= self.blocks_to_swap and self.blocks_to_swap >= 0:
                     block.to(self.offload_device, non_blocking=self.use_non_blocking)
 
-            if self.enable_teacache and pred_id is not None:
+            if self.enable_teacache and (self.teacache_start_step <= current_step <= self.teacache_end_step) and pred_id is not None:
                 self.teacache_state.update(
                     pred_id,
                     previous_residual=(x.to(original_x.device) - original_x),