feat: LTX 2.x model module — recognised from the checkpoint header, joint video+audio latent sized from the audio VAE, decode, text/image-to-video pipelines

Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com>

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).

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",

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

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}

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}

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