File size: 7,005 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 | #!/usr/bin/env python3
"""Unit test for H3 Latent Upscale, the node that resizes the VIDEO half of
an H3 nested AV latent and leaves the audio alone.
python tests/test_latent_upscale.py
Synthetic, no GPU, no comfy, no models.
What it checks:
1. SHAPES. The video's h/w become the requested size; batch, channels and
the TIME axis are untouched (the token clock is not this node's
business).
2. AUDIO PASSTHROUGH is bit-exact and the nested container survives.
3. BOTH MODES run and differ: nearest-exact replicates cell values
verbatim (a 2x nearest upscale is a 2x2 block copy), bilinear does not.
4. GRID SNAPPING. Targets land on the legal /32-pixel = /2-cell grid, from
scale factors AND from explicit width/height, and an odd request is
snapped rather than accepted.
5. PLAIN LATENTS. A bare VAEEncode video latent (no audio) works, because
that is where the node sits in the latent-resident graph.
6. NO-OP. scale 1.0 changes nothing and returns the video by identity.
7. NOISE MASK, if present, follows the video (nearest) so a freeze mask
does not silently apply at the wrong resolution.
Exit code 0 = pass.
"""
import importlib.util
import os
import sys
import torch
HERE = os.path.dirname(os.path.abspath(__file__))
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)
FAILS = []
def check(name, ok, detail=""):
print((" PASS " if ok else " FAIL ") + name + (" " + detail if detail else ""))
if not ok:
FAILS.append(name)
class FakeNested:
"""Stands in for comfy.nested_tensor.NestedTensor: same duck type."""
def __init__(self, tensors):
self.tensors = list(tensors)
self.is_nested = True
def unbind(self):
return self.tensors
def av_latent(h=30, w=54, t=32):
torch.manual_seed(3)
return {"samples": FakeNested([torch.randn(1, 24, t, h, w),
torch.randn(1, 32, 2, 178)])}
def main():
node = motion.H3LatentUpscale()
# 1 + 2. nested AV latent, 2x
lat = av_latent()
audio_in = lat["samples"].unbind()[1]
audio_ref = audio_in.clone()
out, rep = node.upscale(samples=lat, scale=2.0, mode="bilinear")
v, a = out["samples"].unbind()
check("2x resizes only h/w of the video",
tuple(v.shape) == (1, 24, 32, 60, 108), str(tuple(v.shape)))
check("the time axis is untouched", v.shape[2] == 32, str(v.shape[2]))
check("the container is still nested", getattr(out["samples"], "is_nested", False))
check("audio comes out bit-exact", a.shape == audio_ref.shape and
torch.equal(a, audio_ref), str(tuple(a.shape)))
check("audio is passed by reference (not copied)", a is audio_in)
check("the input latent's video is not mutated",
tuple(lat["samples"].unbind()[0].shape) == (1, 24, 32, 30, 54))
check("the report names both halves",
"30x54 cells -> 60x108 cells" in rep.replace(" -> ", " -> ") or
"video 30x54 -> 60x108 cells" in rep, rep.splitlines()[0])
check("the report says the audio was untouched",
"audio passed through untouched" in rep, rep.splitlines()[-1])
# 3. both modes; nearest-exact = 2x2 block copy of the source cells
src = av_latent()
v0 = src["samples"].unbind()[0]
near, _ = node.upscale(samples=av_latent(), scale=2.0, mode="nearest-exact")
vn = near["samples"].unbind()[0]
check("nearest-exact 2x is a verbatim 2x2 block copy",
torch.equal(vn[:, :, :, ::2, ::2], v0) and
torch.equal(vn[:, :, :, 1::2, 1::2], v0),
f"max|diff| {float((vn[:, :, :, ::2, ::2] - v0).abs().max()):.3e}")
bil, _ = node.upscale(samples=av_latent(), scale=2.0, mode="bilinear")
vb = bil["samples"].unbind()[0]
check("bilinear differs from nearest", not torch.equal(vb, vn),
f"mean|diff| {float((vb - vn).abs().mean()):.4f}")
check("bilinear stays in range of the source",
float(vb.min()) >= float(v0.min()) - 1e-5 and
float(vb.max()) <= float(v0.max()) + 1e-5,
f"[{float(vb.min()):.3f}, {float(vb.max()):.3f}] vs "
f"[{float(v0.min()):.3f}, {float(v0.max()):.3f}]")
# 4. grid snapping
check("cells snap to the /2 grid (=/32 px)",
[motion._snap_cells(px) for px in (864, 1728, 1000, 900, 33, 8)]
== [54, 108, 62, 56, 2, 2],
str([motion._snap_cells(px) for px in (864, 1728, 1000, 900, 33, 8)]))
odd, rep_odd = node.upscale(samples=av_latent(), scale=1.5, mode="bilinear")
vo = odd["samples"].unbind()[0]
check("a scale landing off-grid is snapped, not accepted",
tuple(vo.shape[3:]) == (46, 82), # 45->46, 81->82 cells
f"{tuple(vo.shape[3:])} cells = {vo.shape[3] * 16}x{vo.shape[4] * 16} px")
wh, _ = node.upscale(samples=av_latent(), scale=2.0, width=1728, height=960,
mode="bilinear")
vw = wh["samples"].unbind()[0]
check("explicit width/height override scale",
tuple(vw.shape[3:]) == (60, 108), str(tuple(vw.shape[3:])))
wh2, _ = node.upscale(samples=av_latent(), scale=2.0, width=1000,
mode="bilinear")
vw2 = wh2["samples"].unbind()[0]
check("a nonzero axis snaps and the zero axis still uses scale",
tuple(vw2.shape[3:]) == (60, 62), str(tuple(vw2.shape[3:])))
# 5. plain video latent (VAEEncode output, no audio)
plain = {"samples": torch.randn(1, 24, 32, 30, 54)}
pout, prep = node.upscale(samples=plain, scale=2.0, mode="bilinear")
check("a plain video latent upscales and stays plain",
isinstance(pout["samples"], torch.Tensor) and
tuple(pout["samples"].shape) == (1, 24, 32, 60, 108),
str(tuple(pout["samples"].shape)))
check("...and the report says there is no audio",
"no audio component" in prep, prep.splitlines()[-1])
# 6. no-op
same = av_latent()
sout, srep = node.upscale(samples=same, scale=1.0, mode="bilinear")
check("scale 1.0 is a no-op (video by identity)",
sout["samples"] is same["samples"], srep.splitlines()[0])
# 7. noise mask follows the video
masked = av_latent()
masked["noise_mask"] = FakeNested([torch.ones(1, 1, 32, 30, 54),
torch.ones(1, 32, 2, 178)])
mout, _ = node.upscale(samples=masked, scale=2.0, mode="bilinear")
mv, ma = mout["noise_mask"].unbind()
check("a nested noise_mask is resized with the video",
tuple(mv.shape) == (1, 1, 32, 60, 108), str(tuple(mv.shape)))
check("the audio mask is left alone", tuple(ma.shape) == (1, 32, 2, 178),
str(tuple(ma.shape)))
print()
if FAILS:
print(f"FAILED: {len(FAILS)} check(s): " + ", ".join(FAILS))
return 1
print("all checks passed")
return 0
if __name__ == "__main__":
sys.exit(main())
|