| |
| """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() |
|
|
| |
| 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]) |
|
|
| |
| 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}]") |
|
|
| |
| 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), |
| 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:]))) |
|
|
| |
| 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]) |
|
|
| |
| 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]) |
|
|
| |
| 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()) |
|
|
