File size: 8,819 Bytes
e488ce7 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 | #!/usr/bin/env python3
"""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)
# multiples of 1/8 in [-4, 4): exactly representable in fp16
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()
|