| |
| """Unit test for H3TrajectoryLoad, the resume half of the trajectory bank. |
| |
| python tests/test_trajectory_load.py |
| CPU only, no GPU, no models. Needs comfy on sys.path for |
| comfy.utils.pack_latents / unpack_latents and comfy.nested_tensor. |
| |
| Why this exists: the bank records comfy's PACKED latent (every stream |
| flattened to (B, 1, N)), so for an AV model like H3 the file holds video |
| and audio in one vector, and the first Load handed that vector to the |
| model as a single stream (IndexError at model.py x[1]). The 2026-08-27 fix |
| unpacks it. These are the properties the fix promised, on a synthetic bank |
| whose values are exact in fp16 so equality is exact: |
| |
| 1. PACKED + latent_shapes in the file -> a NestedTensor of the original |
| streams; video bit-equal; audio bit-equal after / audio_scale. |
| 2. PACKED without latent_shapes + a `reference` LATENT -> the same. |
| 3. Single-stream packed + a plain reference -> reshaped to the reference. |
| 4. Legacy file with a separate "audio" key -> nested, audio / audio_scale. |
| 5. audio_scale 0 leaves audio untouched. |
| 6. Schedule: returns sigmas[k:] (x_step{k} is x ENTERING step k) and k. |
| 7. undo_const_scaling multiplies every stream by 1/(1-sigma_k). |
| 8. No shapes and no reference on a packed file -> a named ValueError |
| (handing the packed vector downstream is the exact pre-fix failure). |
| 9. IS_CHANGED accepts the optional inputs and tracks the file mtime. |
| 10. A reference whose streams do not sum to the banked vector -> a named |
| ValueError in both directions (shorter would silently misalign). |
| 11. audio_scale recorded in the bank file wins over the widget; 0 = as-is. |
| |
| Exit code 0 = pass. |
| """ |
| import os |
| import sys |
| import tempfile |
|
|
| 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) |
|
|
| import torch |
| import comfy.utils |
| from comfy.nested_tensor import NestedTensor |
|
|
| K = 17 |
| VSHAPE = (1, 32, 3, 4, 4) |
| ASHAPE = (1, 16, 8, 1) |
| FAILS = [] |
|
|
|
|
| def check(cond, msg): |
| print((" ok " if cond else " FAIL ") + msg) |
| if not cond: |
| FAILS.append(msg) |
|
|
|
|
| def exact_fp16(shape, seed): |
| g = torch.Generator().manual_seed(seed) |
| |
| return (torch.randint(-32, 32, shape, generator=g).float() / 8.0) |
|
|
|
|
| def make_bank(tmp, video, audio=None, packed=True, shapes_in_file=True, legacy_audio=False): |
| sigmas = torch.linspace(1.0, 0.0, 26) |
| sigmas[K] = 0.85 |
| torch.save(sigmas, os.path.join(tmp, "sigmas.pt")) |
| if legacy_audio: |
| payload = {"step": K, "total_steps": 25, "video": video.half(), "audio": audio.half(), "packed": False} |
| elif packed and audio is not None: |
| vec, shapes = comfy.utils.pack_latents([video, audio]) |
| payload = {"step": K, "total_steps": 25, "video": vec.half(), "packed": True} |
| if shapes_in_file: |
| payload["latent_shapes"] = [list(int(v) for v in s) for s in shapes] |
| elif packed: |
| vec, shapes = comfy.utils.pack_latents([video]) |
| payload = {"step": K, "total_steps": 25, "video": vec.half(), "packed": True} |
| if shapes_in_file: |
| payload["latent_shapes"] = [list(int(v) for v in s) for s in shapes] |
| else: |
| payload = {"step": K, "total_steps": 25, "video": video.half(), "packed": False} |
| torch.save(payload, os.path.join(tmp, f"x_step{K:03d}.pt")) |
| return sigmas |
|
|
|
|
| def main(): |
| node = motion.H3TrajectoryLoad() |
| video = exact_fp16(VSHAPE, 1) |
| audio = exact_fp16(ASHAPE, 2) |
|
|
| print("1. packed + latent_shapes in file") |
| with tempfile.TemporaryDirectory() as tmp: |
| sig = make_bank(tmp, video, audio) |
| out, rem, step = node.load(tmp, K) |
| s = out["samples"] |
| check(isinstance(s, NestedTensor) and len(s.tensors) == 2, "two streams come back nested") |
| check(tuple(s.tensors[0].shape) == VSHAPE and torch.equal(s.tensors[0], video), "video stream bit-equal") |
| check(tuple(s.tensors[1].shape) == ASHAPE and torch.equal(s.tensors[1], audio / 4.0), "audio stream = banked / 4") |
| print("6. schedule") |
| check(torch.equal(rem, sig[K:]), "remaining sigmas are sigmas[k:] (len %d)" % len(rem)) |
| check(step == K, "step returned") |
|
|
| print("2. packed, shapes only from `reference`") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, audio, shapes_in_file=False) |
| ref = {"samples": NestedTensor((torch.zeros(VSHAPE), torch.zeros(ASHAPE)))} |
| out, _, _ = node.load(tmp, K, reference=ref) |
| s = out["samples"] |
| check(isinstance(s, NestedTensor) and torch.equal(s.tensors[0], video) and torch.equal(s.tensors[1], audio / 4.0), |
| "reference supplies the stream shapes") |
|
|
| print("3. single stream packed + plain reference") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, None, shapes_in_file=False) |
| out, _, _ = node.load(tmp, K, reference={"samples": torch.zeros(VSHAPE)}) |
| s = out["samples"] |
| check(not getattr(s, "is_nested", False) and tuple(s.shape) == VSHAPE and torch.equal(s, video), |
| "reshaped to the reference shape") |
|
|
| print("4. legacy file with a separate audio key") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, audio, legacy_audio=True) |
| out, _, _ = node.load(tmp, K) |
| s = out["samples"] |
| check(isinstance(s, NestedTensor) and torch.equal(s.tensors[1], audio / 4.0), "nested, audio / 4") |
|
|
| print("5. audio_scale 0 leaves audio as-is") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, audio) |
| out, _, _ = node.load(tmp, K, audio_scale=0.0) |
| check(torch.equal(out["samples"].tensors[1], audio), "audio untouched") |
|
|
| print("7. undo_const_scaling") |
| with tempfile.TemporaryDirectory() as tmp: |
| sig = make_bank(tmp, video, audio) |
| out, _, _ = node.load(tmp, K, undo_const_scaling=True) |
| s = out["samples"] |
| f = 1.0 / (1.0 - float(sig[K])) |
| check(torch.allclose(s.tensors[0], video * f) and torch.allclose(s.tensors[1], (audio / 4.0) * f), |
| "every stream scaled by 1/(1-sigma_k) = %.4f" % f) |
|
|
| print("8. packed, no shapes, no reference: a named error, never the raw vector") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, audio, shapes_in_file=False) |
| try: |
| node.load(tmp, K); check(False, "raised") |
| except ValueError as e: |
| check("reference" in str(e), "ValueError names the reference input") |
|
|
| print("10. reference from a different clip: a named error, never silent misalignment") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, audio, shapes_in_file=False) |
| short = {"samples": NestedTensor((torch.zeros(1, 32, 2, 4, 4), torch.zeros(1, 16, 5, 1)))} |
| try: |
| node.load(tmp, K, reference=short); check(False, "raised on a shorter reference") |
| except ValueError as e: |
| check("misalign" in str(e), "shorter reference rejected with the sizes named") |
| longer = {"samples": NestedTensor((torch.zeros(1, 32, 4, 4, 4), torch.zeros(1, 16, 9, 1)))} |
| try: |
| node.load(tmp, K, reference=longer); check(False, "raised on a longer reference") |
| except ValueError as e: |
| check("misalign" in str(e), "longer reference rejected with the sizes named") |
|
|
| print("11. audio_scale recorded in the bank file wins over the widget") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, audio) |
| d = torch.load(os.path.join(tmp, f"x_step{K:03d}.pt"), weights_only=True); d["audio_scale"] = 2.0 |
| torch.save(d, os.path.join(tmp, f"x_step{K:03d}.pt")) |
| out, _, _ = node.load(tmp, K, audio_scale=4.0) |
| check(torch.equal(out["samples"].tensors[1], audio / 2.0), "file value 2.0 used, widget 4.0 ignored") |
| out, _, _ = node.load(tmp, K, audio_scale=0.0) |
| check(torch.equal(out["samples"].tensors[1], audio), "widget 0 still means leave as-is") |
|
|
| print("9. IS_CHANGED") |
| with tempfile.TemporaryDirectory() as tmp: |
| make_bank(tmp, video, audio) |
| a = motion.H3TrajectoryLoad.IS_CHANGED(tmp, K, reference=None, audio_scale=4.0, undo_const_scaling=True) |
| check(a == a, "returns a number for an existing file") |
| b = motion.H3TrajectoryLoad.IS_CHANGED(tmp, K + 1) |
| check(b != b, "NaN for a missing step file") |
|
|
| print("\nRESULT:", "PASS" if not FAILS else "FAIL %d" % len(FAILS)) |
| sys.exit(1 if FAILS else 0) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
