30 lines
1.3 KiB
Python
30 lines
1.3 KiB
Python
import sys
|
|
from unittest.mock import MagicMock, patch
|
|
|
|
import torch
|
|
|
|
# comfy.model_management initializes the torch device at import time and requires
|
|
# CUDA (or --cpu), which the unit-test environment does not provide. Stub it so
|
|
# importing comfy.context_windows does not trigger device initialization.
|
|
with patch.dict(sys.modules, {"comfy.model_management": MagicMock()}):
|
|
from comfy.context_windows import create_weights_overlap_linear
|
|
|
|
|
|
class TestCreateWeightsOverlapLinear:
|
|
def test_zero_overlap_returns_ones(self):
|
|
# context_overlap == 0 is reachable (the context-window nodes allow min=0,
|
|
# and the WAN/LTXV nodes derive it via max(overlap // 4, 0)). With no overlap
|
|
# band there is nothing to blend, so every weight is 1.
|
|
w = create_weights_overlap_linear(16, 100, list(range(20, 36)), 0)
|
|
assert w.shape == (16,)
|
|
assert torch.equal(w, torch.ones(16))
|
|
|
|
def test_overlap_ramps_edges(self):
|
|
w = create_weights_overlap_linear(16, 100, list(range(20, 36)), 4)
|
|
assert w.shape == (16,)
|
|
# middle (non-first) window: left edge ramps up to 1 over the overlap width
|
|
assert w[0] < w[3]
|
|
assert torch.isclose(w[3], torch.tensor(1.0))
|
|
# right edge ramps back down
|
|
assert w[-1] < w[-4]
|