From 6bfeb63c4ce7581c4b18ab7bf96eb77533e435b3 Mon Sep 17 00:00:00 2001 From: digital-garbage <133010426+digital-garbage@users.noreply.github.com> Date: Sat, 3 Oct 2026 10:24:43 +0300 Subject: [PATCH] =?UTF-8?q?feat:=20LTX=202.x=20model=20module=20=E2=80=94?= =?UTF-8?q?=20recognised=20from=20the=20checkpoint=20header,=20joint=20vid?= =?UTF-8?q?eo+audio=20latent=20sized=20from=20the=20audio=20VAE,=20decode,?= =?UTF-8?q?=20text/image-to-video=20pipelines?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-Authored-By: Claude Sonnet 5.5 --- app/legacy/store.js | 1 + core/tests/test_shipped_nodes.py | 2 +- modules/latent/empty/nodes.py | 15 ++- modules/models/ltx/__init__.py | 137 +++++++++++++++++++++++++ modules/models/ltx/pipeline.py | 88 ++++++++++++++++ modules/models/ltx/tests/test_ltx.py | 146 +++++++++++++++++++++++++++ 6 files changed, 386 insertions(+), 3 deletions(-) create mode 100644 modules/models/ltx/__init__.py create mode 100644 modules/models/ltx/pipeline.py create mode 100644 modules/models/ltx/tests/test_ltx.py diff --git a/app/legacy/store.js b/app/legacy/store.js index 5101eb19..2319c5a9 100644 --- a/app/legacy/store.js +++ b/app/legacy/store.js @@ -3740,6 +3740,7 @@ if (p.video && p.video[input] !== undefined) return p.video[input]; if (input === "width" || input === "target_width") return p.width; if (input === "height" || input === "target_height") return p.height; + if (input === "frame_rate" || input === "fps") return p.frame_rate; if (input === "length") { // The PROJECT's length per scene, never the scene's own: a scene cropped on the timeline // regenerates at the project length (the crop was a timeline decision; a regenerate is new). diff --git a/core/tests/test_shipped_nodes.py b/core/tests/test_shipped_nodes.py index cabdd80b..f63b1e9a 100644 --- a/core/tests/test_shipped_nodes.py +++ b/core/tests/test_shipped_nodes.py @@ -65,7 +65,7 @@ FROZEN = { "FunPackModifierSettings": ["settings"], "FunPackDecode": ["samples", "vae", "model", "audio_vae", "tile_size"], "FunPackUpscaleVideo": ["filename", "subfolder", "type", "upscale_model"], - "FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size"], + "FunPackEmptyLatent": ["model", "width", "height", "length", "batch_size", "audio_vae", "frame_rate"], "FunPackDiffusionModelLoader": [ "model_name", "weight_dtype", "compute_dtype", "attention", "fp16_accumulation", "sla", "sla_sparsity", "sla_block_size", "sla_protect_audio", "sla_min_seq_len", diff --git a/modules/latent/empty/nodes.py b/modules/latent/empty/nodes.py index 33238ded..1c5c2549 100644 --- a/modules/latent/empty/nodes.py +++ b/modules/latent/empty/nodes.py @@ -31,6 +31,8 @@ nothing here assumes video. A 2-D latent format yields an image latent because that is what the model says it wants. """ +import inspect + import comfy.model_management import torch from comfy_api.latest import io @@ -108,6 +110,12 @@ class FunPackEmptyLatent(io.ComfyNode): "time axis, and it is frames -- not seconds and " "not samples."), io.Int.Input("batch_size", default=1, min=1, max=4096, optional=True), + io.Vae.Input("audio_vae", optional=True, + tooltip="Only for models that generate sound alongside the picture " + "and size its latent from this VAE (LTX)."), + io.Float.Input("frame_rate", default=25.0, min=1.0, max=1000.0, optional=True, + tooltip="Only used with an audio VAE: how long the sound latent " + "is depends on the frames per second."), ], outputs=[ io.Latent.Output(display_name="latent"), @@ -117,11 +125,14 @@ class FunPackEmptyLatent(io.ComfyNode): @classmethod def execute(cls, model, width: int, height: int, length: int, - batch_size: int = 1) -> io.NodeOutput: + batch_size: int = 1, audio_vae=None, frame_rate: float = 25.0) -> io.NodeOutput: for spec, build in registry_mod.current().providers(CAPABILITY): try: + # What only some models need is offered to those that ask for it by name. + extra = {k: v for k, v in (("audio_vae", audio_vae), ("frame_rate", frame_rate)) + if k in inspect.signature(build).parameters} claimed = build(model, width=width, height=height, - length=length, batch_size=batch_size) + length=length, batch_size=batch_size, **extra) except Exception as exc: # noqa: BLE001 # It got past recognising the model, so this IS its model and it # broke. Carrying on to the derivation would answer a question diff --git a/modules/models/ltx/__init__.py b/modules/models/ltx/__init__.py new file mode 100644 index 00000000..31461f42 --- /dev/null +++ b/modules/models/ltx/__init__.py @@ -0,0 +1,137 @@ +"""LTX 2.x (audio+video) support. + +The second model-support module, built the way the first was: it teaches the system to recognise +this model, build its latent and decode it, and offers pipelines for it. Nothing outside this +folder mentions LTX. + +What it contributes, and what it deliberately does NOT: + +* **Traits** -- a joint audio+video model with a compressed time axis. It does NOT announce + `dit_block_hooks`: the block hooks other modules install (attention temperature, block repeat...) + are written for H3's packed block layout, and LTX's blocks take `(video, audio)` as a pair. Until + a module is written for that layout it is absent here, said, rather than half-working. +* **An empty latent** -- video 128 channels at /32 and /8, plus an audio latent whose length comes + from the AUDIO VAE and the frame rate. That is why this provider asks for `audio_vae` and + `frame_rate`: without them the audio half cannot be sized, and a video-only latent would make a + silent clip while reporting success, so it refuses instead. +* **Decode** -- the picture through the video VAE, the sound through the audio VAE. + +Not validated on a real LTX model yet: the shapes are tested against fakes and ComfyUI's own +stock nodes, and the pipelines against the graph builder. +""" + +import comfy.model_management +import torch +from comfy.nested_tensor import NestedTensor + +from ..._core import log, traits as _traits +from .pipeline import presets as _pipeline_presets + +has_block = _traits.has_block + +ID = "model_ltx" +TITLE = "LTX 2.x" +STAGE = "load" +CATEGORY = "system" +STATUS = "experimental" + +# Qualified: a bare class name matches anything built with a class of that name. +MODEL_CLASS = "comfy.ldm.lightricks.av_model.LTXAVModel" + +VIDEO_CHANNELS = 128 +VIDEO_SPATIAL_RATIO = 32 +VIDEO_TEMPORAL_RATIO = 8 + +# ComfyUI tells the audio+video architecture from the video-only one by this key +# (comfy/model_detection.py: `audio_adaln_single.linear.weight`). +_PREFIXES = ("", "model.diffusion_model.") +_SIGNATURE_KEYS = ("audio_adaln_single.linear.weight",) + + +def is_ltx(model) -> bool: + return has_block(model, MODEL_CLASS) + + +def detect(keys) -> bool: + """Whether a checkpoint's tensor names are an LTX audio+video model's, from its header alone.""" + keyset = set(keys) + return any(all(f"{prefix}{key}" in keyset for key in _SIGNATURE_KEYS) for prefix in _PREFIXES) + + +def probe_traits(keys) -> list: + """What `traits()` would say, read before anything loads.""" + if not detect(keys): + return [] + return ["audio_stream", "temporal_latent", "temporal_compression"] + + +def traits(model): + if not is_ltx(model): + return () + return ["audio_stream"] + + +def frames_for(length: int) -> int: + """LTX makes 8k+1 frames; a length off that grid is rounded UP to it.""" + length = max(1, int(length)) + return ((length - 1 + VIDEO_TEMPORAL_RATIO - 1) // VIDEO_TEMPORAL_RATIO) * VIDEO_TEMPORAL_RATIO + 1 + + +def empty_latent(model, width, height, length, batch_size=1, audio_vae=None, frame_rate=25.0): + """The joint latent, or None when this is not an LTX audio+video model.""" + if not is_ltx(model): + return None + if audio_vae is None: + raise RuntimeError( + "this model makes sound with the picture and sizes it from the audio VAE: wire the audio " + "VAE into the empty latent (a video-only latent would give a silent clip)") + + frames = frames_for(length) + device = comfy.model_management.intermediate_device() + video = torch.zeros( + [batch_size, VIDEO_CHANNELS, ((frames - 1) // VIDEO_TEMPORAL_RATIO) + 1, + max(1, height // VIDEO_SPATIAL_RATIO), max(1, width // VIDEO_SPATIAL_RATIO)], + device=device) + z_channels = audio_vae.latent_channels + bins = audio_vae.first_stage_model.latent_frequency_bins + steps = audio_vae.first_stage_model.num_of_latents_from_frames(frames, frame_rate) + audio = torch.zeros((batch_size, z_channels, steps, bins), device=device) + return {"samples": NestedTensor((video, audio)), + "downscale_ratio_spacial": VIDEO_SPATIAL_RATIO, + "downscale_ratio_temporal": VIDEO_TEMPORAL_RATIO} + + +def decode(latent, model=None, vae=None, audio_vae=None, tile_size=0): + """The picture and the sound, each through the VAE that understands it.""" + if model is None or not is_ltx(model): + return None + if not getattr(latent, "is_nested", False): + return None + parts = latent.unbind() + if len(parts) != 2: + raise RuntimeError(f"expected a video and an audio part, got {len(parts)}") + video_latent, audio_latent = parts + if audio_vae is None: + raise RuntimeError("this model generates sound and no audio VAE is wired into the decode") + + if tile_size and tile_size > 0: + try: + images = vae.decode_tiled(video_latent, tile_x=max(1, tile_size // VIDEO_SPATIAL_RATIO), + tile_y=max(1, tile_size // VIDEO_SPATIAL_RATIO)) + except Exception as exc: # noqa: BLE001 + log.once(f"ltx_decode_tiles:{type(exc).__name__}", log.ALERT, "FunPack LTX decode", + f"the tiled decode failed ({type(exc).__name__}: {exc}); decoded in one piece instead") + images = vae.decode(video_latent) + else: + images = vae.decode(video_latent) + if len(images.shape) == 5: + images = images.reshape(-1, *images.shape[-3:]) + + waveform = audio_vae.decode(audio_latent).movedim(-1, 1).to(audio_latent.device) + audio = {"waveform": waveform, "sample_rate": int(audio_vae.first_stage_model.output_sample_rate)} + return images, audio + + +TRAITS = traits +PROVIDES = {"empty_latent": empty_latent, "decode": decode, "detect": detect, + "probe_traits": probe_traits, "pipeline_presets": _pipeline_presets} diff --git a/modules/models/ltx/pipeline.py b/modules/models/ltx/pipeline.py new file mode 100644 index 00000000..096d50e5 --- /dev/null +++ b/modules/models/ltx/pipeline.py @@ -0,0 +1,88 @@ +"""LTX 2.3 pipelines, as loadable presets. + +Built from ComfyUI's own LTX nodes plus FunPack's, the chain the stock templates use: the +checkpoint carries the model and the picture VAE, the audio VAE and the text encoder are loaded +from the same file, conditioning carries the frame rate, and the sampler runs the distilled +schedule written out as numbers (`project_sigmas`: the official anchors, never replaced). + +The checkpoint is picked three times (model, audio VAE, text encoder): they are three nodes that +each read a file, and the first two are the same file. + +Every `frame_rate` and `fps` is a role at `project.video`, so the project's FPS is the one number +the conditioning, the sound's length and the saved video agree on. +""" + +#: The distilled 8-step schedule, from the official ComfyUI LTX 2.x templates. +DISTILLED_SIGMAS = "1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0" + + +def _base(image_to_video: bool): + video = {"at": "project.video", "input": "frame_rate", "label": "FPS"} + latent_source = ["latent", 0] + slots = [ + {"id": "model", "group": "Loaders", "node": "FunPackCheckpointLoader", "inputs": {"ckpt_name": ""}}, + {"id": "audio_vae", "group": "Loaders", "node": "LTXVAudioVAELoader", "inputs": {"ckpt_name": ""}}, + {"id": "clip", "group": "Loaders", "node": "LTXAVTextEncoderLoader", + "inputs": {"text_encoder": "", "ckpt_name": "", "device": "default"}}, + + {"id": "positive", "group": "Preparation", "node": "CLIPTextEncode", + "roles": [{"at": "generation.prompt", "input": "text", "label": "Prompt"}], + "inputs": {"clip": ["clip", 0], "text": ""}}, + {"id": "negative", "group": "Preparation", "node": "CLIPTextEncode", + "roles": [{"at": "project.negative", "input": "text", "label": "Negative prompt"}], + "inputs": {"clip": ["clip", 0], "text": ""}}, + {"id": "conditioning", "group": "Preparation", "node": "LTXVConditioning", + "roles": [video], + "inputs": {"positive": ["positive", 0], "negative": ["negative", 0], "frame_rate": 25.0}}, + + {"id": "latent", "group": "Preparation", "node": "FunPackEmptyLatent", + "roles": [{"at": "project.video", "input": "width", "label": "Width"}, + {"at": "project.video", "input": "height", "label": "Height"}, + {"at": "project.video", "input": "length", "label": "Length"}, + video], + "inputs": {"model": ["model", 0], "audio_vae": ["audio_vae", 0], + "width": 768, "height": 512, "length": 121, "frame_rate": 25.0, "batch_size": 1}}, + ] + if image_to_video: + slots += [ + {"id": "source_image", "group": "Reference media", "node": "FunPackLoadMedia", + "roles": [{"at": "assets.source_image", "input": "media_id"}], + "inputs": {"media_id": ""}}, + {"id": "split", "group": "Preparation", "node": "LTXVSeparateAVLatent", + "inputs": {"av_latent": ["latent", 0]}}, + {"id": "anchor", "group": "Preparation", "node": "LTXVImgToVideoInplace", + "inputs": {"vae": ["model", 2], "image": ["source_image", 0], "latent": ["split", 0], + "strength": 1.0, "bypass": False}}, + {"id": "join", "group": "Preparation", "node": "LTXVConcatAVLatent", + "inputs": {"video_latent": ["anchor", 0], "audio_latent": ["split", 1]}}, + ] + latent_source = ["join", 0] + slots += [ + {"id": "settings", "group": "Preparation", "node": "FunPackModifierSettings", "inputs": {"settings": "{}"}}, + {"id": "modifiers", "group": "Preparation", "node": "FunPackLoadModifiers", + "inputs": {"model": ["model", 0], "settings": ["settings", 0]}}, + {"id": "sigmas", "group": "Sampling", "node": "ManualSigmas", "inputs": {"sigmas": DISTILLED_SIGMAS}}, + {"id": "sampler", "group": "Sampling", "node": "FunPackSampler", + "roles": [{"at": "generation.sampling", "input": "sampler_name", "label": "Sampler"}], + "inputs": {"model": ["modifiers", 0], "positive": ["conditioning", 0], "negative": ["conditioning", 1], + "latent": latent_source, "settings": ["settings", 0], "sigmas": ["sigmas", 0], + "seed": 0, "steps": 8, "cfg": 1.0, "sampler_name": "euler", "scheduler": "normal", + "denoise": 1.0}}, + + {"id": "decode", "group": "Render", "node": "FunPackDecode", "inputs": { + "samples": ["sampler", 0], "vae": ["model", 2], "model": ["model", 0], "audio_vae": ["audio_vae", 0]}}, + {"id": "video", "group": "Render", "node": "CreateVideo", + "roles": [{"at": "project.video", "input": "fps", "label": "FPS"}], + "inputs": {"images": ["decode", 0], "fps": 25.0, "audio": ["decode", 1]}}, + {"id": "save", "group": "Render", "node": "SaveVideo", "inputs": { + "video": ["video", 0], "filename_prefix": "FunPack", "format": "auto"}}, + ] + return slots + + +def presets(): + return [{"id": "ltx23_text_to_video", "title": "LTX 2.3 · Text to Video", "slots": _base(False)}, + {"id": "ltx23_image_to_video", "title": "LTX 2.3 · Image to Video", "slots": _base(True)}] + + +PROVIDES = {"pipeline_presets": presets} diff --git a/modules/models/ltx/tests/test_ltx.py b/modules/models/ltx/tests/test_ltx.py new file mode 100644 index 00000000..0e938766 --- /dev/null +++ b/modules/models/ltx/tests/test_ltx.py @@ -0,0 +1,146 @@ +"""The LTX module: recognising the model, the joint latent, the decode, and its pipelines.""" + +import asyncio +import types +from pathlib import Path + +import pytest + + +@pytest.fixture(scope="module") +def registered(comfyui): + import nodes as comfy_nodes + + 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()) + return comfy_nodes + + +def _ltx(): + from modules.models import ltx + return ltx + + +def test_the_audio_video_architecture_is_recognised_by_its_key_names_alone(comfyui): + m = _ltx() + assert m.detect({"audio_adaln_single.linear.weight"}) + assert m.detect({"model.diffusion_model.audio_adaln_single.linear.weight"}) + assert not m.detect({"adaln_single.linear.weight"}) # video-only LTX + assert not m.detect({"diffusion_model.audio_adaln_single.linear.weight"}) # not a prefix ComfyUI reads + assert not m.detect({"video_patch_proj.weight", "audio_patch_proj.weight"}) # that is H3 + assert "audio_stream" in m.probe_traits({"audio_adaln_single.linear.weight"}) + assert m.probe_traits({"x"}) == [] + + +def test_it_does_not_claim_block_hooks_it_has_no_modules_for(comfyui, monkeypatch): + m = _ltx() + monkeypatch.setattr(m, "has_block", lambda model, cls: True) + assert "dit_block_hooks" not in m.traits(object()) + + +@pytest.mark.parametrize("asked, got", [(1, 1), (9, 9), (10, 17), (121, 121), (122, 129)]) +def test_a_length_off_the_8k_plus_1_grid_is_rounded_up(comfyui, asked, got): + assert _ltx().frames_for(asked) == got + + +class _AudioVae: + latent_channels = 8 + first_stage_model = types.SimpleNamespace( + latent_frequency_bins=16, output_sample_rate=24000, + num_of_latents_from_frames=lambda frames, rate: int(frames / rate * 25)) + + def decode(self, latent): + import torch + return torch.zeros(latent.shape[0], 100, 2) + + +def test_the_latent_is_video_and_audio_sized_from_the_audio_vae(comfyui, monkeypatch): + m = _ltx() + monkeypatch.setattr(m, "is_ltx", lambda model: True) + out = m.empty_latent(object(), width=768, height=512, length=121, batch_size=1, + audio_vae=_AudioVae(), frame_rate=25.0) + video, audio = out["samples"].unbind() + assert tuple(video.shape) == (1, 128, 16, 16, 24) + assert tuple(audio.shape) == (1, 8, 121, 16) + assert out["downscale_ratio_spacial"] == 32 and out["downscale_ratio_temporal"] == 8 + + +def test_no_audio_vae_is_refused_not_made_silent(comfyui, monkeypatch): + m = _ltx() + monkeypatch.setattr(m, "is_ltx", lambda model: True) + with pytest.raises(RuntimeError, match="audio VAE"): + m.empty_latent(object(), 768, 512, 121, 1) + + +def test_another_model_is_not_claimed(comfyui, monkeypatch): + m = _ltx() + monkeypatch.setattr(m, "is_ltx", lambda model: False) + assert m.empty_latent(object(), 768, 512, 121, 1) is None + assert m.decode(object(), model=object()) is None + + +def test_decode_gives_a_picture_and_a_waveform(comfyui, monkeypatch): + import torch + from comfy.nested_tensor import NestedTensor + m = _ltx() + monkeypatch.setattr(m, "is_ltx", lambda model: True) + + class Vae: + def decode(self, latent): + return torch.zeros(1, 9, 64, 64, 3) + + latent = NestedTensor((torch.zeros(1, 128, 2, 2, 2), torch.zeros(1, 8, 5, 16))) + images, audio = m.decode(latent, model=object(), vae=Vae(), audio_vae=_AudioVae()) + assert tuple(images.shape) == (9, 64, 64, 3) + assert audio["sample_rate"] == 24000 and audio["waveform"].shape[0] == 1 + with pytest.raises(RuntimeError, match="audio VAE"): + m.decode(latent, model=object(), vae=Vae(), audio_vae=None) + + +def test_the_empty_latent_node_passes_the_audio_vae_to_a_provider_that_asks_and_not_to_one_that_does_not(registered): + from core import registry + from modules.latent.empty.nodes import FunPackEmptyLatent + seen = {} + + def wants(model, width, height, length, batch_size, audio_vae=None, frame_rate=25.0): + seen["wants"] = (audio_vae, frame_rate) + return {"samples": 1} + + def plain(model, width, height, length, batch_size): + seen["plain"] = True + return None + + class Reg: + def providers(self, cap): + return [(types.SimpleNamespace(id="x"), plain), (types.SimpleNamespace(id="y"), wants)] + import unittest.mock as mock + with mock.patch.object(registry, "current", lambda rescan=False: Reg()): + FunPackEmptyLatent.execute(object(), 64, 64, 1, 1, audio_vae="A", frame_rate=30.0) + assert seen == {"plain": True, "wants": ("A", 30.0)} + + +@pytest.mark.parametrize("pid", ["ltx23_text_to_video", "ltx23_image_to_video"]) +def test_the_pipelines_build_a_real_graph(registered, pid): + from core import graph + from modules.models.ltx.pipeline import presets + + slots = next(p["slots"] for p in presets() if p["id"] == pid) + ids = [s["id"] for s in slots] + assert len(ids) == len(set(ids)) + for slot in slots: + for value in slot["inputs"].values(): + if isinstance(value, list) and len(value) == 2 and isinstance(value[0], str): + assert value[0] in ids + _prompt, problems = graph.build(slots) + unset = ("needs 'ckpt_name'", "needs 'text_encoder'", "is '', which is not one of") # a file nobody has picked yet + structural = [p for p in problems if not any(u in p for u in unset)] + assert structural == [], structural + + +def test_the_presets_are_offered_through_the_registry(registered): + from core import routes + found = {preset["id"] for _spec, make in routes.modules().providers("pipeline_presets") for preset in make()} + assert {"ltx23_text_to_video", "ltx23_image_to_video"} <= found