ComfyUI/tests-unit/comfy_test/context_windows_test.py

29 lines
1.3 KiB
Python

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}"