Fix out-of-range indices in strided uniform-standard context windows

This commit is contained in:
Sean Kim 2026-07-13 22:10:13 -07:00
parent 0aecac867d
commit 56974505cf
2 changed files with 29 additions and 1 deletions

View File

@ -1033,7 +1033,7 @@ def shift_window_to_start(window: list[int], num_frames: int):
def shift_window_to_end(window: list[int], num_frames: int):
# 1) shift window to start
shift_window_to_start(window, num_frames)
end_val = window[-1]
end_val = max(window)
end_delta = num_frames - end_val - 1
for i in range(len(window)):
# 2) add end_delta to each val to slide windows to end

View File

@ -0,0 +1,28 @@
from types import SimpleNamespace
from comfy.context_windows import create_windows_uniform_standard, shift_window_to_end
def test_shift_window_to_end_uses_window_max():
# A dilated (strided) window wraps around, so it is not monotonic and its last
# element is not its largest. shift_window_to_end must slide the window so its
# maximum lands on num_frames - 1 and every index stays in range -- shifting
# based on window[-1] instead over-shifts and produces out-of-bounds indices.
num_frames = 48
window = [0, 4, 8, 12, 16, 20, 24, 28, 32, 36, 40, 44, 0, 4, 8, 12]
shift_window_to_end(window, num_frames)
assert max(window) == num_frames - 1
assert all(0 <= i < num_frames for i in window)
def test_create_windows_uniform_standard_strided_indices_in_range():
# WAN/LTXV-style manual context windows with context_stride >= 2 produce dilated
# windows; every generated index must be a valid frame index. create_windows_
# uniform_standard only reads these handler attributes.
num_frames = 48
handler = SimpleNamespace(context_length=16, context_stride=3, context_overlap=0, _step=0)
windows = create_windows_uniform_standard(num_frames, handler, {})
assert windows
for window in windows:
for idx in window:
assert 0 <= idx < num_frames, f"index {idx} out of range in window {window}"