From d9a28a9b3c08339d12fec77a704d45b62dbd7516 Mon Sep 17 00:00:00 2001
From: kijai <40791699+kijai@users.noreply.github.com>
Date: Tue, 2 Jun 2026 18:52:46 +0300
Subject: [PATCH] Adjust context window

---
 comfy/model_base.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/comfy/model_base.py b/comfy/model_base.py
index 3742062ce..36b3f0fc6 100644
--- a/comfy/model_base.py
+++ b/comfy/model_base.py
@@ -1531,8 +1531,7 @@ class WAN21(BaseModel):
             out = []
             for lat in cond_value.cond:
                 if lat.ndim > dim and lat.shape[dim] > 1 and lat.shape[dim] == x_in.shape[dim]:
-                    idx = tuple([slice(None)] * dim + [window.index_list])
-                    out.append(lat[idx].to(device))
+                    out.append(window.get_tensor(lat, device, dim=dim, retain_index_list=retain_index_list))
                 else:
                     out.append(lat.to(device))
             return cond_value._copy_with(out)