From 11adad2fb4459d3e055481fa11fa562b48129401 Mon Sep 17 00:00:00 2001 From: Claude Date: Tue, 27 Oct 2026 08:21:45 +0000 Subject: [PATCH] LTX direct v2v: end-to-end test (real mp4, saved-pipeline upgrade, Load Media, encode, sampler on a real tiny LTXAV) Co-Authored-By: Claude Sonnet 5.5 Claude-Session: https://claude.ai/code/session_01BzJLJWVu7RZnbRqm3nYAnX --- .../models/ltx/tests/test_v2v_end_to_end.py | 148 ++++++++++++++++++ 1 file changed, 148 insertions(+) create mode 100644 modules/models/ltx/tests/test_v2v_end_to_end.py diff --git a/modules/models/ltx/tests/test_v2v_end_to_end.py b/modules/models/ltx/tests/test_v2v_end_to_end.py new file mode 100644 index 00000000..184a9cd4 --- /dev/null +++ b/modules/models/ltx/tests/test_v2v_end_to_end.py @@ -0,0 +1,148 @@ +"""Direct LTX video-to-video, run for real on CPU: a multi-frame mp4 in the media library, a pipeline saved with the old +first-frame wiring, the per-run values the app sets, the graph FunUI builds from them, then the nodes that make the +latent run in order (Load Media, size, VAE encode, empty sound, concat) and the sampler on a real tiny LTX model. + +Only the two VAEs are stand-ins (their weights are not in the repo): they have the real shapes (8x time, 32x space).""" + +import json + +import pytest +import torch + + +@pytest.fixture +def clip_id(comfyui, tmp_path, monkeypatch): + import av + from core import config, media + monkeypatch.setattr(config, "MEDIA_DIR", tmp_path / "media") + path = tmp_path / "clip.mp4" + with av.open(str(path), "w") as out: + stream = out.add_stream("libx264", rate=25) + stream.width, stream.height, stream.pix_fmt = 128, 64, "yuv420p" + for i in range(41): + frame = av.VideoFrame.from_ndarray(torch.full((64, 128, 3), i * 6, dtype=torch.uint8).numpy(), format="rgb24") + for packet in stream.encode(frame): + out.mux(packet) + for packet in stream.encode(): + out.mux(packet) + return media.save_upload("clip.mp4", path.read_bytes())["id"] + + +class StubVae: + """The LTX video VAE's shape contract: 8x in time (first frame alone), 32x in space, 128 channels.""" + def encode(self, pixels): + b, f, h, w, _ = pixels.shape if pixels.ndim == 5 else (1, *pixels.shape) + return torch.randn(1, 128, (f - 1) // 8 + 1, h // 32, w // 32) + + +class StubAudioVae: + latent_channels = 8 + + class first_stage_model: + latent_frequency_bins = 16 + + @staticmethod + def num_of_latents_from_frames(frames, rate): + return max(1, round(frames / rate * 25)) + + +@pytest.fixture +def real_tiny_ltx(comfyui): + """ComfyUI's own model_base.LTXAV at a few thousand parameters: the real sampling object and packing.""" + import comfy.model_patcher as mp + import comfy.supported_models as sm + torch.manual_seed(0) + config = sm.LTXAV(dict( + image_model="ltxav", in_channels=128, audio_in_channels=128, cross_attention_dim=32, audio_cross_attention_dim=16, + attention_head_dim=16, audio_attention_head_dim=8, num_attention_heads=2, audio_num_attention_heads=2, + caption_channels=48, num_layers=2, caption_proj_before_connector=True)) + config.set_inference_dtype(torch.float32, None) + base = config.get_model({}, device=torch.device("cpu")) + for p in base.diffusion_model.parameters(): + torch.nn.init.normal_(p, std=0.05) + p.requires_grad_(False) + return mp.ModelPatcher(base, load_device=torch.device("cpu"), offload_device=torch.device("cpu")) + + +def _run(prompt, stubs): + """Run the graph's nodes in order. Loaders are replaced by `stubs` {slot id: outputs tuple}.""" + import nodes as comfy_nodes + done = {} + + def value(v): + return done[v[0]][v[1]] if isinstance(v, list) and len(v) == 2 and isinstance(v[0], str) and v[0] in done else v + + def run(node_id): + if node_id in done: + return + spec = prompt[node_id] + for v in spec["inputs"].values(): + if isinstance(v, list) and len(v) == 2 and v[0] in prompt: + run(v[0]) + if node_id in stubs: + done[node_id] = stubs[node_id] + return + cls = comfy_nodes.NODE_CLASS_MAPPINGS[spec["class_type"]] + kwargs = {k: value(v) for k, v in spec["inputs"].items()} + if hasattr(cls, "GET_SCHEMA"): + out = cls.execute(**kwargs) + done[node_id] = tuple(out.result) + else: + done[node_id] = getattr(cls(), cls.FUNCTION)(**kwargs) + return run, done + + +def test_direct_video_to_video_end_to_end(clip_id, real_tiny_ltx, monkeypatch): + import asyncio + from pathlib import Path + import nodes as comfy_nodes + from core import graph + from modules.models.ltx.pipeline import presets, upgrade + + async def load(): + await comfy_nodes.init_extra_nodes(init_custom_nodes=False) + await comfy_nodes.load_custom_node(str(Path(__file__).resolve().parents[4]), module_parent="custom_nodes") + asyncio.run(load()) + + import types + monkeypatch.setattr(comfy_nodes.NODE_CLASS_MAPPINGS["GetImageSize"], "hidden", types.SimpleNamespace(unique_id=None), raising=False) # no executor here + slots = next(p["slots"] for p in presets() if p["id"] == "ltx23_video_to_video_direct") + # saved by the build that wired the first-frame output: brought up to date as a run does + old = [dict(s, inputs={k: (["source_image", 0] if v == ["source_image", 1] else v) for k, v in s["inputs"].items()}) for s in slots] + # what the old wiring did, which is what a run showed as "1 frames": the whole-video flag cannot reach output 0 + old_slots, _ = graph.override(old, {"source_image": {"media_id": clip_id, "full": True}}) + old_prompt, _ = graph.build(old_slots) + run_old, done_old = _run(old_prompt, {"vae": (StubVae(), ""), "audio_vae": (StubAudioVae(), "")}) + run_old("size") + assert done_old["size"][2] == 1 + slots = upgrade(old) + + prepare = json.dumps({"skip_first": 1, "cut_last": 0, "every_nth": 1, "multiple_of": 32, "fps": 25}) + slots, problems = graph.override(slots, {"source_image": {"media_id": clip_id, "full": True, "prepare": prepare}}) + assert problems == [] + prompt, problems = graph.build(slots) + unset = ("needs 'ckpt_name'", "needs 'model_name'", "needs 'vae_name'", "needs 'clip_name1'", "is '', which is not one of") + assert [p for p in problems if not any(u in p for u in unset)] == [] + + vae, audio_vae = StubVae(), StubAudioVae() + run, done = _run(prompt, {"model": (real_tiny_ltx, ""), "vae": (vae, ""), "audio_vae": (audio_vae, ""), "clip": (None, "")}) + run("latent") + + frames = done["source_image"][1] + assert frames.shape[0] == 40 and frames.shape[1:3] == (64, 128) # the whole video, first frame skipped + assert done["size"][2] == 40 # the sound is sized from the real count + video, audio = done["latent"][0]["samples"].unbind() + assert video.shape == (1, 128, 5, 2, 4) # 40 frames -> 1 + (40-1)//8 latent frames + assert audio.shape[-1] == 16 + + # the sampler runs on it from part-way down the schedule, with the real tiny model + ctx = torch.randn(1, 5, 48) + from modules.sampling.sampler.nodes import FunUISampler + run("sigmas") + sig = done["sigmas"][0] + assert [round(float(x), 3) for x in sig] == [0.725, 0.422, 0.0] # the tail of the distilled schedule + start = done["latent"][0] + out = FunUISampler.execute(real_tiny_ltx, [[ctx, {}]], [[ctx, {}]], start, 1, 2, 1.0, "euler", "normal", sigmas=sig).result[0] + got_video, got_audio = out["samples"].unbind() + assert got_video.shape == video.shape and got_audio.shape == audio.shape + assert torch.isfinite(got_video).all() and not torch.allclose(got_video, video) # it was denoised, not passed through