| |
| """Unit test for H3 V2V Init's time-varying freeze mask (torch only, no GPU, |
| no ComfyUI process). |
| |
| python tests/test_v2v_time_mask.py |
| |
| Checks the pixel-frame -> latent-token map and the pooling that rides on it: |
| spans tile the clip exactly for legal lengths, every 5th token is a singleton |
| sitting on a 17-multiple frame, a mask that switches on AT a 17-multiple |
| lights that singleton and nothing earlier, a mixed token max-pools (any |
| regenerate frame wins), off-length masks resample nearest, and a T==1 mask |
| reproduces the static-union result. |
| |
| Exit code 0 = pass. |
| """ |
| import os |
| import sys |
|
|
| import torch |
|
|
| HERE = os.path.dirname(os.path.abspath(__file__)) |
| sys.path.insert(0, os.path.dirname(HERE)) |
| import importlib.util |
|
|
| spec = importlib.util.spec_from_file_location( |
| "mainodes_motion", os.path.join(os.path.dirname(HERE), "motion.py")) |
| motion = importlib.util.module_from_spec(spec) |
| spec.loader.exec_module(motion) |
|
|
| LEGAL = [39, 90, 141, 192] |
|
|
|
|
| def t_lat_of(length): |
| return (length - 5) // 17 * 5 + 2 |
|
|
|
|
| def test_spans(): |
| for length in LEGAL: |
| t_lat = t_lat_of(length) |
| spans = motion._token_frame_spans(t_lat) |
| |
| assert (t_lat - 2) // 5 * 17 + 5 == length, length |
| |
| assert spans[0][0] == 0 and spans[-1][1] == length - 1, spans[-1] |
| for (a, b), (c, _) in zip(spans, spans[1:]): |
| assert b == c - 1, (a, b, c) |
| sizes = [b - a + 1 for a, b in spans] |
| assert sum(sizes) == length, (sum(sizes), length) |
| |
| for t, (a, b) in enumerate(spans): |
| if t % 5 == 0: |
| assert (a, b) == (17 * (t // 5),) * 2, (t, a, b) |
| else: |
| assert b - a + 1 == 4, (t, a, b) |
| print(f"PASS spans len={length}: t_lat={t_lat} sizes={sizes} " |
| f"singleton frames={[a for t, (a, b) in enumerate(spans) if t % 5 == 0]}") |
|
|
|
|
| def test_switch_on_17(): |
| for length in (39, 141): |
| t_lat = t_lat_of(length) |
| for onset in range(17, length, 17): |
| m = torch.zeros(length, 2, 2) |
| m[onset:] = 1.0 |
| tok = motion._tokenize_mask_time(m, t_lat, length) |
| hot = [t for t in range(t_lat) if tok[t].max() > 0] |
| first = onset // 17 * 5 |
| assert hot == list(range(first, t_lat)), (length, onset, hot) |
| assert first % 5 == 0 |
| print(f"PASS onset len={length} frame={onset}: first hot token " |
| f"{first} (singleton), tokens 0..{first - 1} cold") |
|
|
|
|
| def test_mixed_token_maxpool(): |
| length, t_lat = 39, t_lat_of(39) |
| spans = motion._token_frame_spans(t_lat) |
| a, b = spans[3] |
| assert (a, b) == (9, 12), spans[3] |
| for f in range(a, b + 1): |
| m = torch.zeros(length, 1, 1) |
| m[f] = 1.0 |
| tok = motion._tokenize_mask_time(m, t_lat, length) |
| hot = [t for t in range(t_lat) if tok[t].max() > 0] |
| assert hot == [3], (f, hot) |
| |
| m = torch.zeros(length, 1, 4) |
| m[9, 0, 0] = 1.0 |
| m[12, 0, 3] = 1.0 |
| tok = motion._tokenize_mask_time(m, t_lat, length) |
| assert tok[3].flatten().tolist() == [1.0, 0.0, 0.0, 1.0], tok[3] |
| print("PASS mixed token: frames 9..12 all light token 3 only; " |
| "pixel-wise max within the token = [1,0,0,1]") |
|
|
|
|
| def test_resample(): |
| length, t_lat = 90, t_lat_of(90) |
| |
| src = torch.zeros(45, 1, 1) |
| src[9:] = 1.0 |
| tok = motion._tokenize_mask_time(src, t_lat, length) |
| hot = [t for t in range(t_lat) if tok[t].max() > 0] |
| |
| assert hot == list(range(6, t_lat)), hot |
| ref = torch.zeros(length, 1, 1) |
| ref[18:] = 1.0 |
| assert torch.equal(tok, motion._tokenize_mask_time(ref, t_lat, length)) |
| print(f"PASS resample: T=45 -> {length} nearest, onset frame 18 -> first " |
| f"hot token 6, identical to the already-length mask") |
|
|
|
|
| def test_t1_equals_static(): |
| for length in LEGAL: |
| t_lat = t_lat_of(length) |
| m = (torch.rand(1, 6, 7) > 0.5).float() |
| tok = motion._tokenize_mask_time(m, t_lat, length) |
| static = m.max(dim=0, keepdim=True).values.expand(t_lat, 6, 7) |
| assert torch.equal(tok, static), length |
| print("PASS T==1: pooled result equals the static-union expand for " |
| f"lengths {LEGAL}") |
|
|
|
|
| if __name__ == "__main__": |
| test_spans() |
| test_switch_on_17() |
| test_mixed_token_maxpool() |
| test_resample() |
| test_t1_equals_static() |
| print("ALL PASS") |
|
|
