From e96cf80f5e04d078b7f5213b90fe8c043e4a01d1 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 27 Aug 2026 02:05:11 +0200 Subject: [PATCH] [modular] LTX-2.5: two-stage generation as one pipeline Add `LTX25TwoStageBlocks` / `LTX25TwoStageModularPipeline`: the distilled two-stage recipe (first pass, 2x latent upsample, second pass, diffusion decode) as a single modular pipeline for every workflow `LTX25AutoBlocks` supports. The stages are ordinary blocks that can be popped and run on their own. - split the shared LTX-2 leaves into first-pass / second-pass blocks (`LTX2Stage2PrepareLatentsStep`, `LTX2Stage2PrepareAudioLatentsStep`, `LTX2ConditionStage2PrepareLatentsStep`) with `sigmas_name` / `sigmas_default` init arguments instead of branching on `latents` inside one block - every core-denoise group takes and leaves latents in the VAE form: `LTX2UnpackLatentsStep` closes each group, encoders normalize and decoders denormalize, `LTX2LatentUpsampleStep` bridges the passes - `modular_blocks_ltx25.py` is self-contained (no imports from the LTX-2 preset), with the distilled schedules as defaults and no `num_inference_steps`; `LTX25ModularPipeline` carries the LTX-2.5 latent statistics as the fallback for stages run without an autoencoder - geometry and statistics come from pipeline properties instead of declaring `vae` / `audio_vae` in denoise-side blocks; `use_cross_timestep` is a pipeline property; `batch_size` / `dtype` come from the text input step; the in-context attention mask is built inside the prepare-latents block - agent guide: gotcha on latent form across block boundaries Co-Authored-By: Claude Fable 5 --- .ai/references/modular.md | 12 +- docs/source/en/api/pipelines/ltx2.md | 118 +- src/diffusers/__init__.py | 4 + src/diffusers/modular_pipelines/__init__.py | 11 +- .../modular_pipelines/ltx2/__init__.py | 12 +- .../modular_pipelines/ltx2/before_denoise.py | 1348 ++++++++----- .../modular_pipelines/ltx2/decoders.py | 221 ++- .../modular_pipelines/ltx2/denoise.py | 124 +- .../modular_pipelines/ltx2/encoders.py | 358 +--- .../ltx2/modular_blocks_ltx2.py | 740 +++---- .../ltx2/modular_blocks_ltx25.py | 1715 +++++++++++++++-- .../ltx2/modular_pipeline.py | 106 +- src/diffusers/modular_pipelines/ltx2/utils.py | 523 +++++ .../modular_pipelines/modular_pipeline.py | 1 + .../dummy_torch_and_transformers_objects.py | 30 + .../ltx2/test_modular_pipeline_ltx2.py | 10 +- .../ltx2/test_modular_pipeline_ltx25.py | 40 +- .../test_modular_pipeline_ltx25_two_stage.py | 250 +++ 18 files changed, 3962 insertions(+), 1661 deletions(-) create mode 100644 tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25_two_stage.py diff --git a/.ai/references/modular.md b/.ai/references/modular.md index 8bfb39ebb999..e32a9756f79f 100644 --- a/.ai/references/modular.md +++ b/.ai/references/modular.md @@ -315,7 +315,17 @@ ComponentSpec( 9. **Serving a checkpoint variant through a config flag in a shared block.** `ConfigSpec(name="is_distilled")` plus `if components.config.is_distilled:` bundles two checkpoints' behavior into one blockset — and it can't change the input surface at all (the distilled variant would still accept `negative_prompt`). Suggest a separate blockset for the variant instead (see Key pattern: Checkpoint variants). -10. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. +10. **Declaring a pretrained model component just to read a config value from it.** Everything in `expected_components` gets loaded, so an encoder or decoder block should not declare the `transformer` just to read its patch size, and a denoise block should not declare the `vae` just to read its compression ratio: a block run on its own would then load a model it never calls. Put such values on the `ModularPipeline` subclass as a property that reads the component when it is loaded and falls back to a constant otherwise (`vae_spatial_compression_ratio`, `latents_mean` in `ltx2/modular_pipeline.py`); the fallback lets a block run on its own, and the loaded component wins whenever it is there. + +11. **Raw `torch.randn(device=...)` for noise.** Use `randn_tensor(...)` from `utils/torch_utils`: it draws on the generator's device and moves the result, so CPU generators (what the test mixins pass) work, and the CUDA-generator path is bit-identical to `torch.randn`. + +12. **Latent form drifting across block boundaries.** Two transforms sit between a VAE and a transformer: *normalizing* (the VAE's latent statistics / `scaling_factor`) and *packing* (`[B, C, F, H, W]` → a token sequence `[B, S, D]`). Whoever applies one must have a mirror block that undoes it, and a core denoise group must hand back `latents` in the same form it took them -- otherwise latents get normalized twice, unpacked against the wrong geometry, or reach a decoder in a form it cannot read. The convention across the modular pipelines: + - **Normalize / denormalize live on the VAE blocks.** The VAE encoder emits normalized `image_latents`; the decoder denormalizes right before `vae.decode`. Nothing inside the denoise group applies or removes latent statistics -- that block would need the VAE's stats (gotcha 10) and the encoder's output would no longer be usable as-is. + - **Pack / unpack live inside the core denoise group.** The prepare-latents / input step packs, and a dedicated after-denoise step at the end of the group unpacks (`QwenImageAfterDenoiseStep`, `Flux2UnpackLatentsStep`, `MiniMaxH3AfterDenoiseStep`, `LTX2UnpackLatentsStep`). The group's `latents` output is then the VAE form its input had, so decoders, upsamplers and a second denoise pass take it as-is and need no `height` / `width` / `num_frames` just to unpack. Unpacking in the decoder instead (`flux`, `krea2`, `ltx`) leaves packed latents in state that no other block can consume and makes the decoder carry geometry inputs it does not otherwise need. + - If the transformer patchifies internally (SD3, SDXL, Wan, HunyuanVideo, Helios, Cosmos, Anima, Z-Image), don't pack at block level at all. + - Say which form a tensor is in wherever it crosses a boundary: `"packed, normalized [B, S, D]"` / `"[B, C, F, H, W], normalized"` in the `InputParam` / `OutputParam` descriptions. + + Known reasons to deviate, worth a comment where they apply: the statistics are stored over the *packed* channels, so denormalizing has to happen before unpacking or tile the stats (`ernie_image`, `ideogram4`); unpacking needs per-token ids rather than `height` / `width` (`flux2`, whose unpack step consumes `latent_ids`); or a later step conditions on decoded *pixels*, so decode has to run inside the loop (`wan_animate_2`, Cosmos transfer). ## Conversion checklist diff --git a/docs/source/en/api/pipelines/ltx2.md b/docs/source/en/api/pipelines/ltx2.md index f45a21014a4a..60684880431e 100644 --- a/docs/source/en/api/pipelines/ltx2.md +++ b/docs/source/en/api/pipelines/ltx2.md @@ -923,13 +923,12 @@ LTX-2.5 is also available as a modular pipeline. The default blockset uses the d import torch from diffusers import ModularPipeline, ComponentsManager from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor -from diffusers.pipelines.ltx2.utils import DEFAULT_NEGATIVE_PROMPT -from diffusers.utils import encode_video +from diffusers.utils import encode_video, load_image device = "cuda" # or "mps", "xpu", "cpu" -frame_rate = 24.0 random_seed = 42 generator = torch.Generator(device).manual_seed(random_seed) +frame_rate = 24.0 model_path = "Lightricks/LTX-2.5-Diffusers" @@ -949,13 +948,6 @@ prompt = ( output_state = pipe( prompt=prompt, - negative_prompt=DEFAULT_NEGATIVE_PROMPT, - width=768, - height=512, - num_frames=None, # Set to an int (e.g. 121) to specify a fixed video length - frame_rate=frame_rate, - num_inference_steps=30, - use_cross_timestep=True, enable_prompt_enhancement=True, generator=generator, output_type="np", @@ -975,43 +967,12 @@ encode_video( The modular pipeline will automatically switch workflows based on the supplied inputs. For example, if `image` is supplied, an I2V workflow will be used: ```py -import torch -from diffusers import ModularPipeline, ComponentsManager -from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor -from diffusers.pipelines.ltx2.utils import DEFAULT_NEGATIVE_PROMPT -from diffusers.utils import encode_video, load_image - -device = "cuda" # or "mps", "xpu", "cpu" -frame_rate = 24.0 -random_seed = 42 -generator = torch.Generator(device).manual_seed(random_seed) - -model_path = "Lightricks/LTX-2.5-Diffusers" - -cm = ComponentsManager() -pipe = ModularPipeline.from_pretrained(model_path, components_manager=cm) -pipe.load_components(dtype=torch.bfloat16) -cm.enable_auto_cpu_offload(device=device, memory_reserve_margin="20GB") -pipe.diffusion_decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor()) -pipe.diffusion_decoder.enable_tiling() - -prompt = ( - "An astronaut hatches from a fragile egg on the surface of the Moon, the shell cracking and peeling apart in " - "gentle low-gravity motion." -) image_path = "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg" image = load_image(image_path) output_state = pipe( image=image, prompt=prompt, - negative_prompt=DEFAULT_NEGATIVE_PROMPT, - width=768, - height=512, - num_frames=None, # Set to an int (e.g. 121) to specify a fixed video length - frame_rate=frame_rate, - num_inference_steps=30, - use_cross_timestep=True, enable_prompt_enhancement=True, generator=generator, output_type="np", @@ -1028,6 +989,73 @@ encode_video( ) ``` +#### Two-stage generation (modular) + +`LTX25TwoStageBlocks` runs the [distilled two-stage recipe](#two-stage-generation-for-ltx-25) in one call, for every workflow `LTX25AutoBlocks` supports (image and frame conditions are re-encoded at the upsampled resolution for the second pass): a first pass at the requested `height` / `width`, a 2x latent upsample, and a second pass that refines at the upsampled resolution. As with the standard pipelines, the resolution you pass is the first pass's, and the output is twice that size. `stage_1` is the same auto denoise step as `LTX25AutoBlocks`; `stage_2` selects the workflow's second-pass group, which re-noises the upsampled latents on `stage_2_sigmas` (the distilled stage-2 schedule by default) instead of sampling fresh noise, and the upsample step doubles `height` / `width` in between. + +The `latent_upsampler` is a component of the blockset like any other. Load it explicitly if the repository's `modular_model_index.json` does not list it: + +```py +import torch +from diffusers import ComponentsManager +from diffusers.modular_pipelines import LTX25TwoStageBlocks +from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.utils import encode_video + +device = "cuda" +model_path = "Lightricks/LTX-2.5-Diffusers" +prompt = "A cinematic shot of a red fox walking through a snowy forest at dawn, golden light filtering through pine trees." +frame_rate = 24.0 + +cm = ComponentsManager() +pipe = LTX25TwoStageBlocks().init_pipeline(model_path, components_manager=cm) +pipe.load_components(dtype=torch.bfloat16) +pipe.update_components( + latent_upsampler=LTX2LatentUpsamplerModel.from_pretrained( + model_path, subfolder="latent_upsampler", dtype=torch.bfloat16 + ) +) +cm.enable_auto_cpu_offload(device=device, memory_reserve_margin="20GB") + +# First pass at the default 704x512, output at 1408x1024; `num_frames` is predicted by the duration head. +output = pipe( + prompt=prompt, + generator=torch.Generator(device).manual_seed(42), + output_type="np", +) +video, audio = output.get("videos"), output.get("audio") + +encode_video( + video[0], + fps=frame_rate, + audio=audio[0].float().cpu(), + audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + output_path="ltx2_5_modular_two_stage.mp4", +) +``` + +The stages are ordinary blocks, so the same blockset splits into separate pipelines -- to preview the first pass, swap in a different upsampler, or load a LoRA for the second pass only. Every core denoise group leaves `[B, C, F, H, W]` video and `[B, C, L, M]` audio latents in state -- normalized, in the same form the VAE encoder blocks emit -- so `stage_1` followed by `decode` is a first-pass preview, and `upsample` and `stage_2` take exactly what `stage_1` leaves. Chained by hand with one generator threaded through, the result matches the single call: + +```py +blocks = LTX25TwoStageBlocks() +stage_2 = blocks.sub_blocks.pop("stage_2") +upsample = blocks.sub_blocks.pop("upsample") +decode = blocks.sub_blocks.pop("decode") + +# `blocks` now ends with `stage_1`; the four pipelines share components through the manager. +stage_1_pipe = blocks.init_pipeline(model_path, components_manager=cm) +upsample_pipe = upsample.init_pipeline(model_path, components_manager=cm) +stage_2_pipe = stage_2.init_pipeline(model_path, components_manager=cm) +decode_pipe = decode.init_pipeline(model_path, components_manager=cm) + +# Each pipeline reads what it needs from the state the previous one leaves. +generator = torch.Generator(device).manual_seed(42) +state = stage_1_pipe(prompt=prompt, generator=generator) +state = upsample_pipe(state=state) +state = stage_2_pipe(state=state) +video = decode_pipe(state=state, output_type="np", output="videos") +``` + You can see the supported workflows in the docs for each blockset (e.g. [`LTX2AutoBlocks`], [`LTX25AutoBlocks`]). ## LTX2Pipeline @@ -1085,3 +1113,11 @@ You can see the supported workflows in the docs for each blockset (e.g. [`LTX2Au ## LTX25AutoBlocks [[autodoc]] LTX25AutoBlocks + +## LTX25TwoStageModularPipeline + +[[autodoc]] LTX25TwoStageModularPipeline + +## LTX25TwoStageBlocks + +[[autodoc]] LTX25TwoStageBlocks diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 747bec2bdf84..750b76bb82ef 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -545,6 +545,8 @@ "Krea2TurboModularPipeline", "LTX25AutoBlocks", "LTX25ModularPipeline", + "LTX25TwoStageBlocks", + "LTX25TwoStageModularPipeline", "LTX2AutoBlocks", "LTX2ModularPipeline", "LTXAutoBlocks", @@ -1398,6 +1400,8 @@ LTX2ModularPipeline, LTX25AutoBlocks, LTX25ModularPipeline, + LTX25TwoStageBlocks, + LTX25TwoStageModularPipeline, LTXAutoBlocks, LTXModularPipeline, MiniMaxH3Blocks, diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 81b93f88f515..db6e92096509 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -124,8 +124,10 @@ _import_structure["ltx2"] = [ "LTX2AutoBlocks", "LTX25AutoBlocks", + "LTX25TwoStageBlocks", "LTX2ModularPipeline", "LTX25ModularPipeline", + "LTX25TwoStageModularPipeline", ] _import_structure["minimax_h3"] = [ "MiniMaxH3Blocks", @@ -189,7 +191,14 @@ Krea2TurboModularPipeline, ) from .ltx import LTXAutoBlocks, LTXModularPipeline - from .ltx2 import LTX2AutoBlocks, LTX2ModularPipeline, LTX25AutoBlocks, LTX25ModularPipeline + from .ltx2 import ( + LTX2AutoBlocks, + LTX2ModularPipeline, + LTX25AutoBlocks, + LTX25ModularPipeline, + LTX25TwoStageBlocks, + LTX25TwoStageModularPipeline, + ) from .minimax_h3 import ( MiniMaxH3Blocks, MiniMaxH3ModularPipeline, diff --git a/src/diffusers/modular_pipelines/ltx2/__init__.py b/src/diffusers/modular_pipelines/ltx2/__init__.py index caf44b9a6179..519e55558006 100644 --- a/src/diffusers/modular_pipelines/ltx2/__init__.py +++ b/src/diffusers/modular_pipelines/ltx2/__init__.py @@ -28,8 +28,12 @@ "LTX2ImageToVideoBlocks", "LTX2InContextBlocks", ] - _import_structure["modular_blocks_ltx25"] = ["LTX25AutoBlocks"] - _import_structure["modular_pipeline"] = ["LTX2ModularPipeline", "LTX25ModularPipeline"] + _import_structure["modular_blocks_ltx25"] = ["LTX25AutoBlocks", "LTX25TwoStageBlocks"] + _import_structure["modular_pipeline"] = [ + "LTX2ModularPipeline", + "LTX25ModularPipeline", + "LTX25TwoStageModularPipeline", + ] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -45,8 +49,8 @@ LTX2ImageToVideoBlocks, LTX2InContextBlocks, ) - from .modular_blocks_ltx25 import LTX25AutoBlocks - from .modular_pipeline import LTX2ModularPipeline, LTX25ModularPipeline + from .modular_blocks_ltx25 import LTX25AutoBlocks, LTX25TwoStageBlocks + from .modular_pipeline import LTX2ModularPipeline, LTX25ModularPipeline, LTX25TwoStageModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/ltx2/before_denoise.py b/src/diffusers/modular_pipelines/ltx2/before_denoise.py index 81ffc28188ea..b72d5607920b 100644 --- a/src/diffusers/modular_pipelines/ltx2/before_denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/before_denoise.py @@ -18,18 +18,20 @@ import numpy as np import torch -from ...models import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, LTX2VideoTransformer3DModel +from ...models import LTX2VideoTransformer3DModel # NOTE (modular.md gotcha #1): `LTX2ReferenceCondition` is a plain dataclass under `diffusers.pipelines.ltx2.*`, and # modular blocks must not import from `diffusers.pipelines.*`. It belongs in the same neutral-module relocation as # the other shared LTX-2 data/utilities enumerated in `encoders.py`. Imported from the pipelines path here only so # the draft is runnable. +from ...pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel from ...pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging from ...utils.torch_utils import randn_tensor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .decoders import _denormalize_latents logger = logging.get_logger(__name__) @@ -165,14 +167,6 @@ def _normalize_latents( return latents -def _normalize_audio_latents( - latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor -) -> torch.Tensor: - latents_mean = latents_mean.to(latents.device, latents.dtype) - latents_std = latents_std.to(latents.device, latents.dtype) - return (latents - latents_mean) / latents_std - - def _create_noised_state( latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None ) -> torch.Tensor: @@ -180,6 +174,92 @@ def _create_noised_state( return noise_scale * noise + (1 - noise_scale) * latents +def _downsample_mask_to_latent( + mask: torch.Tensor, latent_num_frames: int, latent_height: int, latent_width: int +) -> torch.Tensor: + """ + Downsample a pixel-space attention mask of shape `(B, 1, F, H, W)` (values in `[0, 1]`) to a flattened per-token + latent-space mask of shape `(B, latent_num_frames * latent_height * latent_width)`. Spatial downsampling is area + interpolation per frame; temporal downsampling is causal (the first frame is kept as-is). + """ + if mask.ndim != 5 or mask.shape[1] != 1: + raise ValueError(f"Expected `conditioning_attention_mask` of shape (B, 1, F, H, W), got {tuple(mask.shape)}.") + b, _, f_pix, _, _ = mask.shape + + mask_2d = mask.reshape(b * f_pix, 1, mask.shape[-2], mask.shape[-1]) + spatial_down = torch.nn.functional.interpolate(mask_2d, size=(latent_height, latent_width), mode="area") + spatial_down = spatial_down.reshape(b, 1, f_pix, latent_height, latent_width) + + first_frame = spatial_down[:, :, :1, :, :] + if f_pix > 1 and latent_num_frames > 1: + t = (f_pix - 1) // (latent_num_frames - 1) + if (f_pix - 1) % (latent_num_frames - 1) != 0: + raise ValueError( + f"Pixel frames ({f_pix}) not compatible with latent frames ({latent_num_frames}): " + f"(f_pix - 1) must be divisible by (latent_num_frames - 1)." + ) + rest = spatial_down[:, :, 1:, :, :] + rest = rest.reshape(b, 1, latent_num_frames - 1, t, latent_height, latent_width).mean(dim=3) + latent_mask = torch.cat([first_frame, rest], dim=2) + else: + latent_mask = first_frame + + return latent_mask.reshape(b, latent_num_frames * latent_height * latent_width) + + +def _build_video_self_attention_mask( + latents: torch.Tensor, + num_base_tokens: int, + num_ref_tokens: int, + reference_latents: list[torch.Tensor], + reference_token_counts: list[int], + conditioning_attention_strength: float, + conditioning_attention_mask: torch.Tensor | None, +) -> torch.Tensor: + """ + Builds the multiplicative video self-attention mask `[B, S, S]` over the `[base | keyframe | reference]` token + sequence of in-context generation, mirroring `build_attention_mask` in the reference implementation. Each reference + is its own attention group: + + - base <-> base, base <-> keyframe, keyframe <-> keyframe: 1.0 (full attention) + - base <-> reference group: that group's per-token strengths (`conditioning_attention_mask` downsampled to the + reference's latent grid, or ones, times `conditioning_attention_strength`), broadcast symmetrically + - reference group <-> itself: 1.0; reference group <-> any other appended group: 0.0 + + The cross blocks span only the *base* tokens: keyframe tokens are appended conditioning like the references, so the + two are masked off from each other. + """ + device = latents.device + batch_size, total_tokens, _ = latents.shape + num_prefix_tokens = total_tokens - num_ref_tokens + + cross = [] + for ref_latent, num_tokens in zip(reference_latents, reference_token_counts): + if conditioning_attention_mask is not None: + _, _, ref_latent_frames, ref_latent_height, ref_latent_width = ref_latent.shape + ref_cross = _downsample_mask_to_latent( + conditioning_attention_mask, ref_latent_frames, ref_latent_height, ref_latent_width + ).to(device=device, dtype=torch.float32) + else: + ref_cross = torch.ones((1, num_tokens), device=device, dtype=torch.float32) + cross.append(ref_cross * conditioning_attention_strength) + cross = torch.cat(cross, dim=1) + + # Start from zeros so the keyframe<->reference and reference<->reference blocks stay masked without explicit + # assignment. Each guidance pass is its own single-batch forward, so this is built at the generation batch size. + attn_mask = torch.zeros((batch_size, total_tokens, total_tokens), device=device, dtype=torch.float32) + attn_mask[:, :num_prefix_tokens, :num_prefix_tokens] = 1.0 + + offset = num_prefix_tokens + for group_cross in torch.split(cross, reference_token_counts, dim=1): + n = group_cross.shape[1] + attn_mask[:, :num_base_tokens, offset : offset + n] = group_cross.unsqueeze(1) + attn_mask[:, offset : offset + n, :num_base_tokens] = group_cross.unsqueeze(2) + attn_mask[:, offset : offset + n, offset : offset + n] = 1.0 + offset += n + return attn_mask + + def _prepare_keyframe_coords( keyframe_latent_num_frames: int, keyframe_latent_height: int, @@ -234,10 +314,11 @@ class LTX2TextInputStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Input processing step that expands the connector text conditioning (cond and uncond) by " - "`num_videos_per_prompt`, so it matches the `batch_size * num_videos_per_prompt` batch of the video and " - "audio latents. Runs at the head of the denoise stage, which keeps the text-conditioning stage's outputs " - "reusable across denoise runs with different `num_videos_per_prompt`." + "Input processing step that reports the prompt count (`batch_size`) and embedding `dtype`, and expands " + "the connector text conditioning (cond and uncond) by `num_videos_per_prompt`, so it matches the " + "`batch_size * num_videos_per_prompt` batch of the video and audio latents. Runs at the head of the " + "denoise stage, which keeps the text-conditioning stage's outputs reusable across denoise runs with " + "different `num_videos_per_prompt`." ) @property @@ -265,20 +346,17 @@ def inputs(self) -> list[InputParam]: InputParam( "negative_connector_prompt_embeds", type_hint=torch.Tensor, - required=True, - description="Video-branch text conditioning (uncond).", + description="Video-branch text conditioning (uncond), `None` without classifier-free guidance.", ), InputParam( "negative_connector_audio_prompt_embeds", type_hint=torch.Tensor, - required=True, - description="Audio-branch text conditioning (uncond).", + description="Audio-branch text conditioning (uncond), `None` without classifier-free guidance.", ), InputParam( "negative_connector_attention_mask", type_hint=torch.Tensor, - required=True, - description="Binary text attention mask (uncond).", + description="Binary text attention mask (uncond), `None` without classifier-free guidance.", ), ] @@ -315,17 +393,35 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=torch.Tensor, description="Binary text attention mask (uncond), expanded per prompt.", ), + OutputParam( + "batch_size", + type_hint=int, + description="The number of prompts being denoised (before per-prompt expansion).", + ), + OutputParam("dtype", type_hint=torch.dtype, description="The dtype of the text conditioning."), ] @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) + block_state.batch_size = block_state.connector_prompt_embeds.shape[0] + block_state.dtype = block_state.connector_prompt_embeds.dtype + # `repeat_interleave` keeps each prompt's copies contiguous, matching how the latents are laid out # (`batch_size * num_videos_per_prompt`, prompt-major) and how `image_latents` are expanded downstream. num_videos = block_state.num_videos_per_prompt - for name in self.intermediate_output_names: - setattr(block_state, name, getattr(block_state, name).repeat_interleave(num_videos, dim=0)) + for name in ( + "connector_prompt_embeds", + "connector_audio_prompt_embeds", + "connector_attention_mask", + "negative_connector_prompt_embeds", + "negative_connector_audio_prompt_embeds", + "negative_connector_attention_mask", + ): + value = getattr(block_state, name) + if value is not None: # the negative-prompt tensors are `None` without classifier-free guidance + setattr(block_state, name, value.repeat_interleave(num_videos, dim=0)) self.set_block_state(state, block_state) return components, state @@ -334,6 +430,26 @@ def __call__(self, components, state: PipelineState) -> PipelineState: class LTX2SetTimestepsStep(ModularPipelineBlocks): model_name = "ltx2" + def __init__( + self, sigmas_name: str = "sigmas", timesteps_name: str = "timesteps", sigmas_default: list[float] | None = None + ): + """ + Args: + sigmas_name (`str`, defaults to `"sigmas"`): + Name of the input that holds this pass's sigma schedule. Lets a first-pass and a second-pass copy of + the block sit in one pipeline that takes both `sigmas` and `stage_2_sigmas`. + timesteps_name (`str`, defaults to `"timesteps"`): + Name of the input that holds this pass's custom timesteps, for the same reason. + sigmas_default (`list[float]`, *optional*): + Default sigma schedule of the pass. Set where a blockset assembles the block for a checkpoint that runs + a fixed schedule (the LTX-2.5 distilled recipe); the block then exposes no `num_inference_steps`. + `None` leaves the schedule to `num_inference_steps`. + """ + self._sigmas_name = sigmas_name + self._timesteps_name = timesteps_name + self._sigmas_default = sigmas_default + super().__init__() + @property def description(self) -> str: return ( @@ -347,22 +463,22 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: - return [ - InputParam.template("num_inference_steps", default=30), - InputParam.template("timesteps"), - InputParam.template("sigmas"), + inputs = [ + InputParam.template("timesteps", name=self._timesteps_name), + InputParam.template("sigmas", name=self._sigmas_name, default=self._sigmas_default), InputParam.template("height", default=512), InputParam.template("width", default=704), InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), ] + # A block assembled with a fixed schedule has no step count to choose. + if self._sigmas_default is None: + inputs.insert(0, InputParam.template("num_inference_steps", default=30)) + return inputs @property def intermediate_outputs(self) -> list[OutputParam]: @@ -380,8 +496,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - num_inference_steps = block_state.num_inference_steps - sigmas = block_state.sigmas + num_inference_steps = getattr(block_state, "num_inference_steps", None) + timesteps = getattr(block_state, self._timesteps_name) + sigmas = getattr(block_state, self._sigmas_name) if sigmas is None: sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) @@ -406,11 +523,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ) block_state.audio_scheduler = copy.deepcopy(components.scheduler) - retrieve_timesteps( - block_state.audio_scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu - ) + retrieve_timesteps(block_state.audio_scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu) block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( - components.scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu + components.scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu ) # Set begin index to skip the nonzero().item() call in scheduler init, which triggers a GPU sync. @@ -427,16 +542,15 @@ class LTX2PrepareLatentsStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Prepares the packed video noise latents for text-to-video generation. `noise_scale` is declared with a " - "`None` default and resolved to 0.0 here rather than declared as 0.0: a blockset keeps the first " - "non-`None` default across its blocks, so a literal 0.0 would shadow the condition workflow's " - "`None -> sigmas[0] or 1.0` resolution wherever the two share a blockset (`LTX2AutoBlocks`). The " - "resolved value is written back to state for `LTX2PrepareAudioLatentsStep`." + "Samples the packed video noise latents for a first pass of text-to-video generation. Refining " + "existing latents is `LTX2Stage2PrepareLatentsStep`." ) @property def expected_components(self) -> list[ComponentSpec]: - return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)] + return [ + ComponentSpec("transformer", LTX2VideoTransformer3DModel), + ] @property def inputs(self) -> list[InputParam]: @@ -446,40 +560,114 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), - InputParam.template("latents"), InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam.template("generator"), + InputParam( + "batch_size", + type_hint=int, + required=True, + description="The number of prompts being denoised, used to expand conditioning per prompt.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [OutputParam("latents", type_hint=torch.Tensor, description="Packed noisy video latents.")] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + batch_size = block_state.batch_size * block_state.num_videos_per_prompt + num_channels_latents = components.transformer.config.in_channels + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + + shape = (batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width) + latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=torch.float32) + block_state.latents = _pack_latents( + latents, components.transformer_spatial_patch_size, components.transformer_temporal_patch_size + ) + + self.set_block_state(state, block_state) + return components, state + + +class LTX2Stage2PrepareLatentsStep(ModularPipelineBlocks): + model_name = "ltx2" + + def __init__(self, sigmas_name: str = "sigmas", sigmas_default: list[float] | None = None): + """ + Args: + sigmas_name (`str`, defaults to `"sigmas"`): + Name of the input that holds this pass's sigma schedule. Lets a first-pass and a second-pass copy of + the block sit in one pipeline that takes both `sigmas` and `stage_2_sigmas`. + sigmas_default (`list[float]`, *optional*): + Default sigma schedule of the pass, set where a blockset assembles the block for a checkpoint that runs + a fixed schedule (the LTX-2.5 distilled recipe). Read for the `noise_scale` default. + """ + self._sigmas_name = sigmas_name + self._sigmas_default = sigmas_default + super().__init__() + + @property + def description(self) -> str: + return ( + "Prepares the packed video latents for a second pass that refines existing latents: packs the normalized " + "`[B, C, F, H, W]` latents a first pass or `LTX2LatentUpsampleStep` leaves in state, then re-noises them " + "to `noise_scale` -- by default the first sigma of the pass, as in the reference two-stage recipe. The " + "resolved `noise_scale` is written back to state for `LTX2Stage2PrepareAudioLatentsStep`, and `height` / `width` / `num_frames` " + "are read off the latents for the blocks that follow." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", LTX2VideoTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latents", + type_hint=torch.Tensor, + required=True, + description="Video latents to refine, of shape [B, C, F, H, W] (normalized, not packed).", + ), InputParam( "noise_scale", type_hint=float, default=None, description=( - "Interpolation factor between random noise and any provided latents. `None` (default) resolves " - "to 0.0, which keeps the provided latents." + "Noise level the latents are re-noised to before the pass. `None` (default) resolves to " + "`sigmas[0]` when custom `sigmas` are supplied, else 1.0." ), ), + InputParam.template("sigmas", name=self._sigmas_name, default=self._sigmas_default), InputParam.template("generator"), - InputParam( - "batch_size", - type_hint=int, - required=True, - description="The number of prompts being denoised, used to expand conditioning per prompt.", - ), ] @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("latents", type_hint=torch.Tensor, description="Packed noisy video latents."), + OutputParam("latents", type_hint=torch.Tensor, description="Packed re-noised video latents."), + OutputParam("height", type_hint=int, description="Height of the pass in pixels, read off the latents."), + OutputParam("width", type_hint=int, description="Width of the pass in pixels, read off the latents."), + OutputParam( + "num_frames", + type_hint=int, + description="Frame count of the pass, read off the latents (grid-aligned).", + ), OutputParam( "noise_scale", type_hint=float, - description="The resolved interpolation factor, forwarded to the audio latents step.", + description="The resolved noise level, forwarded to the audio latents step.", ), ] @@ -488,34 +676,24 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - batch_size = block_state.batch_size * block_state.num_videos_per_prompt - num_channels_latents = components.transformer.config.in_channels - spatial_patch = components.transformer_spatial_patch_size - temporal_patch = components.transformer_temporal_patch_size + # The supplied latents fix the geometry of the pass; the blocks after this one read it from state. + _, _, latent_num_frames, latent_height, latent_width = block_state.latents.shape + block_state.height = latent_height * components.vae_spatial_compression_ratio + block_state.width = latent_width * components.vae_spatial_compression_ratio + block_state.num_frames = (latent_num_frames - 1) * components.vae_temporal_compression_ratio + 1 - if block_state.noise_scale is None: - block_state.noise_scale = 0.0 - - if block_state.latents is not None: - latents = block_state.latents - if latents.ndim == 5: - latents = _normalize_latents( - latents, - components.latents_mean, - components.latents_std, - components.vae_scaling_factor, - ) - latents = _pack_latents(latents, spatial_patch, temporal_patch) - latents = _create_noised_state(latents, block_state.noise_scale, block_state.generator) - block_state.latents = latents.to(device=device, dtype=torch.float32) - else: - latent_height = block_state.height // components.vae_spatial_compression_ratio - latent_width = block_state.width // components.vae_spatial_compression_ratio - latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + noise_scale = block_state.noise_scale + if noise_scale is None: + sigmas = getattr(block_state, self._sigmas_name) + noise_scale = sigmas[0] if sigmas is not None else 1.0 - shape = (batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width) - latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=torch.float32) - block_state.latents = _pack_latents(latents, spatial_patch, temporal_patch) + latents = _pack_latents( + block_state.latents, components.transformer_spatial_patch_size, components.transformer_temporal_patch_size + ) + # Re-noise in the latents' own dtype and cast afterwards, the order `LTX2Pipeline.prepare_latents` uses. + latents = _create_noised_state(latents.to(device), noise_scale, block_state.generator) + block_state.latents = latents.to(dtype=torch.float32) + block_state.noise_scale = noise_scale self.set_block_state(state, block_state) return components, state @@ -551,11 +729,8 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), InputParam( @@ -619,15 +794,7 @@ class LTX2PrepareAudioLatentsStep(ModularPipelineBlocks): @property def description(self) -> str: - return ( - "Prepares the packed audio noise latents and derives the audio latent frame count. `noise_scale` is " - "declared with a `None` default for the reason given on `LTX2PrepareLatentsStep`, which runs first and " - "writes the resolved value back to state." - ) - - @property - def expected_components(self) -> list[ComponentSpec]: - return [ComponentSpec("audio_vae", AutoencoderKLLTX2Audio)] + return "create the initial audio noise latents (packed) and derives the audio latent frame count.stage1 only" @property def inputs(self) -> list[InputParam]: @@ -635,30 +802,12 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam( "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." ), - InputParam( - "audio_latents", - type_hint=torch.Tensor, - default=None, - description="Optional pre-encoded audio latents; random noise is used when not provided.", - ), - InputParam( - "noise_scale", - type_hint=float, - default=None, - description=( - "Interpolation factor between random noise and any provided `audio_latents`. Resolved upstream " - "by `LTX2PrepareLatentsStep` (0.0, which keeps the provided latents)." - ), - ), InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), InputParam.template("generator"), InputParam( @@ -687,8 +836,8 @@ def __call__(self, components, state: PipelineState) -> PipelineState: device = components._execution_device batch_size = block_state.batch_size * block_state.num_videos_per_prompt - num_channels_latents = components.audio_vae.config.latent_channels - num_mel_bins = components.audio_vae.config.mel_bins + num_channels_latents = components.audio_latent_channels + latent_mel_bins = components.audio_latent_mel_bins duration_s = block_state.num_frames / block_state.frame_rate audio_latents_per_second = ( @@ -698,22 +847,88 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ) audio_num_frames = round(duration_s * audio_latents_per_second) - if block_state.audio_latents is not None: - audio_latents = block_state.audio_latents - if audio_latents.ndim == 4: - audio_num_frames = audio_latents.shape[2] - audio_latents = _pack_audio_latents(audio_latents) - audio_latents = _normalize_audio_latents( - audio_latents, components.audio_latents_mean, components.audio_latents_std - ) - audio_latents = _create_noised_state(audio_latents, block_state.noise_scale, block_state.generator) - block_state.audio_latents = audio_latents.to(device=device, dtype=torch.float32) - else: - latent_mel_bins = num_mel_bins // components.audio_vae_mel_compression_ratio - shape = (batch_size, num_channels_latents, audio_num_frames, latent_mel_bins) - audio_latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=torch.float32) - block_state.audio_latents = _pack_audio_latents(audio_latents) + shape = (batch_size, num_channels_latents, audio_num_frames, latent_mel_bins) + audio_latents = randn_tensor(shape, generator=block_state.generator, device=device, dtype=torch.float32) + block_state.audio_latents = _pack_audio_latents(audio_latents) + block_state.audio_num_frames = audio_num_frames + + self.set_block_state(state, block_state) + return components, state + + +class LTX2Stage2PrepareAudioLatentsStep(ModularPipelineBlocks): + model_name = "ltx2" + + def __init__(self, sigmas_name: str = "sigmas", sigmas_default: list[float] | None = None): + """ + Args: + sigmas_name (`str`, defaults to `"sigmas"`): + Name of the input that holds this pass's sigma schedule. Lets a first-pass and a second-pass copy of + the block sit in one pipeline that takes both `sigmas` and `stage_2_sigmas`. + sigmas_default (`list[float]`, *optional*): + Default sigma schedule of the pass, set where a blockset assembles the block for a checkpoint that runs + a fixed schedule (the LTX-2.5 distilled recipe). Read for the `noise_scale` default. + """ + self._sigmas_name = sigmas_name + self._sigmas_default = sigmas_default + super().__init__() + + @property + def description(self) -> str: + return ( + "Prepares the audio latents for stage2 that refines existing audio latents: packs the " + "normalized `[B, C, L, M]` latents from stage1, derives the audio latent frame count " + "from their shape, and re-noises them to `noise_scale`." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "audio_latents", + type_hint=torch.Tensor, + required=True, + description="Audio latents to refine, of shape [B, C, L, M] (normalized, not packed).", + ), + InputParam( + "noise_scale", + type_hint=float, + default=None, + description=( + "Noise level the audio latents are re-noised to before the pass. `None` (default) resolves to " + "`sigmas[0]` when custom `sigmas` are supplied, else 1.0." + ), + ), + InputParam.template("sigmas", name=self._sigmas_name, default=self._sigmas_default), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("audio_latents", type_hint=torch.Tensor, description="Packed re-noised audio latents."), + OutputParam( + "audio_num_frames", + type_hint=int, + kwargs_type="denoiser_input_fields", + description="Number of audio latent frames.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + noise_scale = block_state.noise_scale + if noise_scale is None: + sigmas = getattr(block_state, self._sigmas_name) + noise_scale = sigmas[0] if sigmas is not None else 1.0 + audio_num_frames = block_state.audio_latents.shape[2] + audio_latents = _pack_audio_latents(block_state.audio_latents) + audio_latents = _create_noised_state(audio_latents.to(device), noise_scale, block_state.generator) + block_state.audio_latents = audio_latents.to(dtype=torch.float32) block_state.audio_num_frames = audio_num_frames self.set_block_state(state, block_state) @@ -742,11 +957,8 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam( "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." @@ -799,9 +1011,117 @@ def __call__(self, components, state: PipelineState) -> PipelineState: return components, state +def _apply_frame_conditions(components, block_state, latents: torch.Tensor, noise_scale: float): + """ + Shared tail of the condition prepare-latents blocks: packs the base `[B, C, F, H, W]` latents, overwrites the + first-frame positions with latent-index-0 conditions, appends the other conditions as keyframe tokens with their + own RoPE coordinates, and samples the noise last, once the conditioning mask is known. + + Returns the packed noisy `latents`, `conditioning_mask`, `clean_latents`, `appended_coords` and `base_token_count`. + """ + device = components._execution_device + batch_size = block_state.batch_size * block_state.num_videos_per_prompt + spatial_patch = components.transformer_spatial_patch_size + temporal_patch = components.transformer_temporal_patch_size + frame_scale_factor = components.vae_temporal_compression_ratio + _, _, latent_num_frames, latent_height, latent_width = latents.shape + + conditioning_mask = latents.new_zeros((batch_size, 1, latent_num_frames, latent_height, latent_width)) + latents = _pack_latents(latents, spatial_patch, temporal_patch) + conditioning_mask = _pack_latents(conditioning_mask, spatial_patch, temporal_patch) # [B, S, 1] + + base_token_count = latents.shape[1] + condition_latents_packed = [ + _pack_latents(cond, spatial_patch, temporal_patch) for cond in block_state.condition_latents + ] + + # First-frame conditions (latent index 0): overwrite the tokens at the first-frame positions. Condition + # tensors carry batch 1 and broadcast across the generation batch. + clean_latents = torch.zeros_like(latents) + for cond, strength, latent_idx in zip( + condition_latents_packed, block_state.condition_strengths, block_state.condition_indices + ): + if latent_idx != 0: + continue + num_cond_tokens = cond.size(1) + latents[:, :num_cond_tokens] = cond + conditioning_mask[:, :num_cond_tokens] = strength + clean_latents[:, :num_cond_tokens] = cond + + # Non-first-frame ("keyframe") conditions (latent index > 0): append as extra tokens with an all-`strength` + # conditioning mask and their own coords. At denoising step i they see an effective noise level of + # (1 - strength) * sigma_i. + scale_factors = ( + frame_scale_factor, + components.vae_spatial_compression_ratio, + components.vae_spatial_compression_ratio, + ) + keyframe_tokens, keyframe_masks, keyframe_coords = [], [], [] + for cond_5d, cond_packed, strength, latent_idx, num_pixel_frames in zip( + block_state.condition_latents, + condition_latents_packed, + block_state.condition_strengths, + block_state.condition_indices, + block_state.condition_pixel_frames, + ): + if latent_idx == 0: + continue + + _, _, kf_latent_frames, kf_latent_height, kf_latent_width = cond_5d.shape + coords = _prepare_keyframe_coords( + keyframe_latent_num_frames=kf_latent_frames, + keyframe_latent_height=kf_latent_height, + keyframe_latent_width=kf_latent_width, + pixel_frame_idx=(latent_idx - 1) * frame_scale_factor + 1, + num_pixel_frames=num_pixel_frames, + fps=block_state.frame_rate, + patch_size=spatial_patch, + patch_size_t=temporal_patch, + scale_factors=scale_factors, + device=device, + ) + + keyframe_tokens.append(cond_packed.expand(batch_size, -1, -1)) + keyframe_masks.append( + torch.full( + (batch_size, cond_packed.shape[1], 1), + float(strength), + device=device, + dtype=conditioning_mask.dtype, + ) + ) + keyframe_coords.append(coords.expand(batch_size, -1, -1, -1)) + + if keyframe_tokens: + keyframe_tokens = torch.cat(keyframe_tokens, dim=1) + latents = torch.cat([latents, keyframe_tokens], dim=1) + clean_latents = torch.cat([clean_latents, keyframe_tokens], dim=1) + conditioning_mask = torch.cat([conditioning_mask, torch.cat(keyframe_masks, dim=1)], dim=1) + appended_coords = torch.cat(keyframe_coords, dim=2) + else: + appended_coords = torch.zeros((batch_size, 3, 0, 2), device=device, dtype=torch.float32) + + # Mask semantics: 0 -> fully noised, 1 -> kept clean, in between -> noise level (1 - mask) * noise_scale. + noise = randn_tensor(latents.shape, generator=block_state.generator, device=latents.device, dtype=latents.dtype) + scaled_mask = (1.0 - conditioning_mask) * noise_scale + latents = noise * scaled_mask + latents * (1 - scaled_mask) + + return latents, conditioning_mask, clean_latents, appended_coords, base_token_count + + class LTX2ConditionPrepareLatentsStep(ModularPipelineBlocks): model_name = "ltx2" + def __init__(self, sigmas_default: list[float] | None = None): + """ + Args: + sigmas_default (`list[float]`, *optional*): + Default sigma schedule of the pass, set where a blockset assembles the block for a checkpoint that runs + a fixed schedule (the LTX-2.5 distilled recipe). Read for the `noise_scale` default. + """ + self._sigmas_default = sigmas_default + super().__init__() + @property def description(self) -> str: return ( @@ -817,7 +1137,6 @@ def description(self) -> str: def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", LTX2VideoTransformer3DModel), - ComponentSpec("vae", AutoencoderKLLTX2Video), ] @property @@ -827,7 +1146,7 @@ def inputs(self) -> list[InputParam]: "condition_latents", type_hint=list, required=True, - description="Per-condition normalized VAE latents of shape [1, C, F, H, W].", + description="Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed).", ), InputParam( "condition_strengths", @@ -847,17 +1166,13 @@ def inputs(self) -> list[InputParam]: required=True, description="Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords.", ), - InputParam.template("latents"), InputParam.template("height", default=512), InputParam.template("width", default=704), InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam( "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." @@ -871,7 +1186,7 @@ def inputs(self) -> list[InputParam]: "when custom `sigmas` are supplied, else 1.0." ), ), - InputParam.template("sigmas"), + InputParam.template("sigmas", default=self._sigmas_default), InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), InputParam( "batch_size", @@ -929,13 +1244,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: device = components._execution_device batch_size = block_state.batch_size * block_state.num_videos_per_prompt - spatial_patch = components.transformer_spatial_patch_size - temporal_patch = components.transformer_temporal_patch_size - frame_scale_factor = components.vae_temporal_compression_ratio - latent_height = block_state.height // components.vae_spatial_compression_ratio latent_width = block_state.width // components.vae_spatial_compression_ratio - latent_num_frames = (block_state.num_frames - 1) // frame_scale_factor + 1 + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 # Noise level the un-conditioned tokens start at: the first (largest) sigma when custom sigmas are supplied, # else 1.0. Matches `LTX2ConditionPipeline.__call__`. @@ -949,116 +1260,194 @@ def __call__(self, components, state: PipelineState) -> PipelineState: f" list will be used for all (pseudo-)random operations." ) - if block_state.latents is not None: - latents = _normalize_latents( - block_state.latents, - components.latents_mean, - components.latents_std, - components.vae_scaling_factor, - ) - else: - # Zeros rather than a Gaussian sample: the noise is mixed in at the end, once the mask is known. - shape = ( - batch_size, - components.transformer.config.in_channels, - latent_num_frames, - latent_height, - latent_width, - ) - latents = torch.zeros(shape, device=device, dtype=torch.float32) + # Zeros rather than a Gaussian sample: the noise is mixed in at the end, once the mask is known. + shape = (batch_size, components.transformer.config.in_channels, latent_num_frames, latent_height, latent_width) + latents = torch.zeros(shape, device=device, dtype=torch.float32) + + ( + block_state.latents, + block_state.conditioning_mask, + block_state.clean_latents, + block_state.appended_coords, + block_state.base_token_count, + ) = _apply_frame_conditions(components, block_state, latents, noise_scale) + block_state.noise_scale = noise_scale - conditioning_mask = latents.new_zeros((batch_size, 1, latent_num_frames, latent_height, latent_width)) - latents = _pack_latents(latents, spatial_patch, temporal_patch) - conditioning_mask = _pack_latents(conditioning_mask, spatial_patch, temporal_patch) # [B, S, 1] + self.set_block_state(state, block_state) + return components, state - if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape[:2]: - raise ValueError( - f"Provided `latents` tensor packs to shape {latents.shape}, but the expected packed shape is " - f"{conditioning_mask.shape[:2] + (components.transformer.config.in_channels,)}." - ) - base_token_count = latents.shape[1] - condition_latents_packed = [ - _pack_latents(cond, spatial_patch, temporal_patch) for cond in block_state.condition_latents - ] +class LTX2ConditionStage2PrepareLatentsStep(ModularPipelineBlocks): + model_name = "ltx2" - # First-frame conditions (latent index 0): overwrite the tokens at the first-frame positions. Condition - # tensors carry batch 1 and broadcast across the generation batch. - clean_latents = torch.zeros_like(latents) - for cond, strength, latent_idx in zip( - condition_latents_packed, block_state.condition_strengths, block_state.condition_indices - ): - if latent_idx != 0: - continue - num_cond_tokens = cond.size(1) - latents[:, :num_cond_tokens] = cond - conditioning_mask[:, :num_cond_tokens] = strength - clean_latents[:, :num_cond_tokens] = cond + def __init__(self, sigmas_name: str = "sigmas", sigmas_default: list[float] | None = None): + """ + Args: + sigmas_name (`str`, defaults to `"sigmas"`): + Name of the input that holds this pass's sigma schedule. Lets a first-pass and a second-pass copy of + the block sit in one pipeline that takes both `sigmas` and `stage_2_sigmas`. + sigmas_default (`list[float]`, *optional*): + Default sigma schedule of the pass, set where a blockset assembles the block for a checkpoint that runs + a fixed schedule (the LTX-2.5 distilled recipe). Read for the `noise_scale` default. + """ + self._sigmas_name = sigmas_name + self._sigmas_default = sigmas_default + super().__init__() - # Non-first-frame ("keyframe") conditions (latent index > 0): append as extra tokens with an all-`strength` - # conditioning mask and their own coords. At denoising step i they see an effective noise level of - # (1 - strength) * sigma_i. - scale_factors = ( - frame_scale_factor, - components.vae_spatial_compression_ratio, - components.vae_spatial_compression_ratio, - ) - keyframe_tokens, keyframe_masks, keyframe_coords = [], [], [] - for cond_5d, cond_packed, strength, latent_idx, num_pixel_frames in zip( - block_state.condition_latents, - condition_latents_packed, - block_state.condition_strengths, - block_state.condition_indices, - block_state.condition_pixel_frames, - ): - if latent_idx == 0: - continue + @property + def description(self) -> str: + return ( + "Prepares the packed video latents for a second pass that refines existing latents under frame " + "conditions: the supplied normalized `[B, C, F, H, W]` latents (packed here) take the " + "place of the zeros `LTX2ConditionPrepareLatentsStep` starts from, then the same first-frame overwrite, " + "keyframe token append and mask-driven noising apply, with `noise_scale` -- by default the first sigma " + "of the pass -- as the level the un-conditioned tokens are re-noised to." + ) - _, _, kf_latent_frames, kf_latent_height, kf_latent_width = cond_5d.shape - coords = _prepare_keyframe_coords( - keyframe_latent_num_frames=kf_latent_frames, - keyframe_latent_height=kf_latent_height, - keyframe_latent_width=kf_latent_width, - pixel_frame_idx=(latent_idx - 1) * frame_scale_factor + 1, - num_pixel_frames=num_pixel_frames, - fps=block_state.frame_rate, - patch_size=spatial_patch, - patch_size_t=temporal_patch, - scale_factors=scale_factors, - device=device, - ) + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", LTX2VideoTransformer3DModel), + ] - keyframe_tokens.append(cond_packed.expand(batch_size, -1, -1)) - keyframe_masks.append( - torch.full( - (batch_size, cond_packed.shape[1], 1), - float(strength), - device=device, - dtype=conditioning_mask.dtype, - ) - ) - keyframe_coords.append(coords.expand(batch_size, -1, -1, -1)) + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "condition_latents", + type_hint=list, + required=True, + description="Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed).", + ), + InputParam( + "condition_strengths", + type_hint=list, + required=True, + description="Per-condition conditioning strengths.", + ), + InputParam( + "condition_indices", + type_hint=list, + required=True, + description="Per-condition latent frame index at which the condition is applied.", + ), + InputParam( + "condition_pixel_frames", + type_hint=list, + required=True, + description="Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords.", + ), + InputParam( + "latents", + type_hint=torch.Tensor, + required=True, + description=( + "Video latents to refine, of shape [B, C, F, H, W] (normalized, not packed) " + "of the generated video only (no appended condition tokens)." + ), + ), + InputParam( + "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." + ), + InputParam( + "noise_scale", + type_hint=float, + default=None, + description=( + "Noise level the un-conditioned tokens are re-noised to. `None` (default) resolves to " + "`sigmas[0]` when custom `sigmas` are supplied, else 1.0." + ), + ), + InputParam.template("sigmas", name=self._sigmas_name, default=self._sigmas_default), + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam( + "batch_size", + type_hint=int, + required=True, + description="The number of prompts being denoised, used to expand conditioning per prompt.", + ), + InputParam.template("generator"), + ] - if keyframe_tokens: - keyframe_tokens = torch.cat(keyframe_tokens, dim=1) - latents = torch.cat([latents, keyframe_tokens], dim=1) - clean_latents = torch.cat([clean_latents, keyframe_tokens], dim=1) - conditioning_mask = torch.cat([conditioning_mask, torch.cat(keyframe_masks, dim=1)], dim=1) - appended_coords = torch.cat(keyframe_coords, dim=2) - else: - appended_coords = torch.zeros((batch_size, 3, 0, 2), device=device, dtype=torch.float32) + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="Packed noisy video latents, with any keyframe condition tokens appended.", + ), + OutputParam( + "conditioning_mask", + type_hint=torch.Tensor, + description=( + "Packed per-token conditioning strengths of shape [B, S, 1] in [0, 1]: 1 at fully-conditioned " + "positions, 0 at free positions." + ), + ), + OutputParam( + "clean_latents", + type_hint=torch.Tensor, + description="Clean condition latents at conditioned positions, zeros elsewhere; same shape as `latents`.", + ), + OutputParam( + "appended_coords", + type_hint=torch.Tensor, + description=( + "RoPE coordinates of shape [B, 3, num_keyframe_tokens, 2] for the appended keyframe tokens, " + "zero-width when there are none." + ), + ), + OutputParam( + "base_token_count", + type_hint=int, + description="Number of generated-video tokens, i.e. the sequence length before appended tokens.", + ), + OutputParam("height", type_hint=int, description="Height of the pass in pixels, read off the latents."), + OutputParam("width", type_hint=int, description="Width of the pass in pixels, read off the latents."), + OutputParam( + "num_frames", + type_hint=int, + description="Frame count of the pass, read off the latents (grid-aligned).", + ), + OutputParam( + "noise_scale", + type_hint=float, + description="The resolved initial noise level, forwarded to the audio latents step.", + ), + ] - # Mask semantics: 0 -> fully noised, 1 -> kept clean, in between -> noise level (1 - mask) * noise_scale. - noise = randn_tensor( - latents.shape, generator=block_state.generator, device=latents.device, dtype=latents.dtype - ) - scaled_mask = (1.0 - conditioning_mask) * noise_scale - block_state.latents = noise * scaled_mask + latents * (1 - scaled_mask) + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device - block_state.conditioning_mask = conditioning_mask - block_state.clean_latents = clean_latents - block_state.appended_coords = appended_coords - block_state.base_token_count = base_token_count + # The supplied latents fix the geometry of the pass; the blocks after this one read it from state. + _, _, latent_num_frames, latent_height, latent_width = block_state.latents.shape + block_state.height = latent_height * components.vae_spatial_compression_ratio + block_state.width = latent_width * components.vae_spatial_compression_ratio + block_state.num_frames = (latent_num_frames - 1) * components.vae_temporal_compression_ratio + 1 + + noise_scale = block_state.noise_scale + if noise_scale is None: + sigmas = getattr(block_state, self._sigmas_name) + noise_scale = sigmas[0] if sigmas is not None else 1.0 + + if isinstance(block_state.generator, list): + logger.warning( + f"{self.__class__.__name__} does not support using a list of generators. The first generator in the" + f" list will be used for all (pseudo-)random operations." + ) + + latents = block_state.latents.to(device=device, dtype=torch.float32) + + ( + block_state.latents, + block_state.conditioning_mask, + block_state.clean_latents, + block_state.appended_coords, + block_state.base_token_count, + ) = _apply_frame_conditions(components, block_state, latents, noise_scale) block_state.noise_scale = noise_scale self.set_block_state(state, block_state) @@ -1068,6 +1457,16 @@ def __call__(self, components, state: PipelineState) -> PipelineState: class LTX2InContextPrepareLatentsStep(ModularPipelineBlocks): model_name = "ltx2" + def __init__(self, sigmas_default: list[float] | None = None): + """ + Args: + sigmas_default (`list[float]`, *optional*): + Default sigma schedule of the pass, set where a blockset assembles the block for a checkpoint that runs + a fixed schedule (the LTX-2.5 distilled recipe). Read for the `noise_scale` default. + """ + self._sigmas_default = sigmas_default + super().__init__() + @property def description(self) -> str: return ( @@ -1076,14 +1475,15 @@ def description(self) -> str: "encoded reference tokens after the keyframes with a per-token `conditioning_mask` of their own " "strength, giving a single `[base | keyframe | reference]` sequence. Mirrors " "`LTX2InContextPipeline.prepare_latents`, which likewise re-implements the condition version rather " - "than extending it -- the two are kept side by side so each reads top to bottom." + "than extending it -- the two are kept side by side so each reads top to bottom. First pass only: the " + "second pass of an in-context run needs no reference tokens and uses " + "`LTX2ConditionStage2PrepareLatentsStep` or `LTX2Stage2PrepareLatentsStep`." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("transformer", LTX2VideoTransformer3DModel), - ComponentSpec("vae", AutoencoderKLLTX2Video), ] @property @@ -1093,7 +1493,7 @@ def inputs(self) -> list[InputParam]: "condition_latents", type_hint=list, required=True, - description="Per-condition normalized VAE latents of shape [1, C, F, H, W].", + description="Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed).", ), InputParam( "condition_strengths", @@ -1124,36 +1524,31 @@ def inputs(self) -> list[InputParam]: ), InputParam( "reference_latents", - type_hint=torch.Tensor, + type_hint=list, default=None, description=( - "Packed reference tokens of shape [1, total_reference_tokens, C], or `None` when no reference " - "conditions were supplied (`LTX2AutoReferenceEncoderStep` is skipped)." + "Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from `LTX2ReferenceEncoderStep`, " + "or `None` when no reference conditions were supplied (`LTX2AutoReferenceEncoderStep` is " + "skipped)." ), ), InputParam( - "reference_coords", - type_hint=torch.Tensor, - default=None, - description="RoPE coordinates for the reference tokens.", - ), - InputParam( - "reference_token_counts", - type_hint=list, - default=None, - description="Per-reference token counts, in `reference_conditions` order.", + "reference_downscale_factor", + type_hint=int, + default=1, + description=( + "Ratio between the target and reference resolutions. The reference tokens' spatial coordinates " + "are scaled by it so they land in the target coordinate space, preserving the positional " + "relationship the IC-LoRA was trained on." + ), ), - InputParam.template("latents"), InputParam.template("height", default=512), InputParam.template("width", default=704), InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam( "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." @@ -1167,7 +1562,7 @@ def inputs(self) -> list[InputParam]: "when custom `sigmas` are supplied, else 1.0." ), ), - InputParam.template("sigmas"), + InputParam.template("sigmas", default=self._sigmas_default), InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), InputParam( "batch_size", @@ -1176,6 +1571,25 @@ def inputs(self) -> list[InputParam]: description="The number of prompts being denoised, used to expand conditioning per prompt.", ), InputParam.template("generator"), + InputParam( + "conditioning_attention_strength", + type_hint=float, + default=1.0, + description=( + "Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each " + "other. 1.0 (default) leaves attention unmasked." + ), + ), + InputParam( + "conditioning_attention_mask", + type_hint=torch.Tensor, + default=None, + description=( + "Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially " + "varying attention strength. Downsampled to each reference's latent grid and multiplied by " + "`conditioning_attention_strength`." + ), + ), ] @property @@ -1212,11 +1626,25 @@ def intermediate_outputs(self) -> list[OutputParam]: type_hint=int, description="Number of generated-video tokens, i.e. the sequence length before appended tokens.", ), + OutputParam( + "video_self_attention_mask", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description=( + "Multiplicative self-attention mask of shape [B, S, S] with values in [0, 1] over the " + "`[base | keyframe | reference]` tokens; `None` without reference tokens (full attention)." + ), + ), OutputParam( "num_ref_tokens", type_hint=int, description="Number of reference tokens, which sit at the very end of the sequence.", ), + OutputParam( + "reference_token_counts", + type_hint=list, + description="Per-reference token counts, in `reference_conditions` order, for the attention mask.", + ), OutputParam( "noise_scale", type_hint=float, @@ -1248,33 +1676,14 @@ def __call__(self, components, state: PipelineState) -> PipelineState: f" list will be used for all (pseudo-)random operations." ) - if block_state.latents is not None: - latents = _normalize_latents( - block_state.latents, - components.latents_mean, - components.latents_std, - components.vae_scaling_factor, - ) - else: - shape = ( - batch_size, - components.transformer.config.in_channels, - latent_num_frames, - latent_height, - latent_width, - ) - latents = torch.zeros(shape, device=device, dtype=torch.float32) + # Zeros rather than a Gaussian sample: the noise is mixed in at the end, once the mask is known. + shape = (batch_size, components.transformer.config.in_channels, latent_num_frames, latent_height, latent_width) + latents = torch.zeros(shape, device=device, dtype=torch.float32) conditioning_mask = latents.new_zeros((batch_size, 1, latent_num_frames, latent_height, latent_width)) latents = _pack_latents(latents, spatial_patch, temporal_patch) conditioning_mask = _pack_latents(conditioning_mask, spatial_patch, temporal_patch) # [B, S, 1] - if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape[:2]: - raise ValueError( - f"Provided `latents` tensor packs to shape {latents.shape}, but the expected packed shape is " - f"{conditioning_mask.shape[:2] + (components.transformer.config.in_channels,)}." - ) - base_token_count = latents.shape[1] condition_latents_packed = [ _pack_latents(cond, spatial_patch, temporal_patch) for cond in block_state.condition_latents @@ -1349,27 +1758,44 @@ def __call__(self, components, state: PipelineState) -> PipelineState: # reference implementation. Absent for IC-LoRAs that take no reference video (camera control, style, ...), # which `LTX2InContextPipeline` supports too -- `LTX2AutoReferenceEncoderStep` is then skipped. num_ref_tokens = 0 + reference_token_counts = [] if block_state.reference_latents is not None: reference_conditions = block_state.reference_conditions if isinstance(reference_conditions, LTX2ReferenceCondition): reference_conditions = [reference_conditions] - reference_tokens = block_state.reference_latents.expand(batch_size, -1, -1) - num_ref_tokens = reference_tokens.shape[1] - # Per-reference token counts come from the encoder rather than an equal split of the total, so - # references of differing lengths (a reference video shorter than `num_frames`) get the right strengths. - reference_masks = [ - torch.full( - (batch_size, num_tokens, 1), - float(ref_cond.strength), + reference_tokens, reference_masks, reference_coords = [], [], [] + for ref_latent, ref_cond in zip(block_state.reference_latents, reference_conditions): + _, _, ref_latent_frames, ref_latent_height, ref_latent_width = ref_latent.shape + tokens = _pack_latents(ref_latent, spatial_patch, temporal_patch) + # Coordinates on the reference's own latent grid, scaled spatially so the tokens map into the + # target's coordinate space. + coords = components.transformer.rope.prepare_video_coords( + batch_size=1, + num_frames=ref_latent_frames, + height=ref_latent_height, + width=ref_latent_width, device=device, - dtype=conditioning_mask.dtype, + fps=block_state.frame_rate, + ) + if block_state.reference_downscale_factor != 1: + coords[:, 1:, :, :] = coords[:, 1:, :, :] * block_state.reference_downscale_factor + reference_tokens.append(tokens.expand(batch_size, -1, -1)) + reference_masks.append( + torch.full( + (batch_size, tokens.shape[1], 1), + float(ref_cond.strength), + device=device, + dtype=conditioning_mask.dtype, + ) ) - for ref_cond, num_tokens in zip(reference_conditions, block_state.reference_token_counts) - ] + reference_coords.append(coords.expand(batch_size, -1, -1, -1)) + reference_token_counts.append(tokens.shape[1]) + reference_tokens = torch.cat(reference_tokens, dim=1) + num_ref_tokens = reference_tokens.shape[1] latents = torch.cat([latents, reference_tokens], dim=1) clean_latents = torch.cat([clean_latents, reference_tokens], dim=1) conditioning_mask = torch.cat([conditioning_mask, torch.cat(reference_masks, dim=1)], dim=1) - appended_coords.append(block_state.reference_coords.expand(batch_size, -1, -1, -1)) + appended_coords.append(torch.cat(reference_coords, dim=2)) # Mask semantics: 0 -> fully noised, 1 -> kept clean, in between -> noise level (1 - mask) * noise_scale. noise = randn_tensor( @@ -1383,103 +1809,22 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state.appended_coords = torch.cat(appended_coords, dim=2) block_state.base_token_count = base_token_count block_state.num_ref_tokens = num_ref_tokens - block_state.noise_scale = noise_scale - - self.set_block_state(state, block_state) - return components, state - - -class LTX2BuildVideoSelfAttentionMaskStep(ModularPipelineBlocks): - model_name = "ltx2" - - @property - def description(self) -> str: - return ( - "Builds the video self-attention mask over the `[base | keyframe | reference]` token sequence for " - "in-context generation, mirroring `build_attention_mask` in the reference implementation. Each " - "`LTX2ReferenceCondition` is its own attention group:\n" - " - base <-> base, base <-> keyframe, keyframe <-> keyframe: 1.0 (full attention)\n" - " - base <-> reference group: that group's slice of `reference_cross_mask`, broadcast symmetrically " - "across the base-token axis\n" - " - reference group <-> itself: 1.0 (a group fully attends to itself)\n" - " - reference group <-> any other appended group (keyframes, other references): 0.0\n" - "Note the cross blocks span only the *base* tokens: keyframe tokens are appended conditioning like the " - "references are, and the two are masked off from each other. Emitted tagged `denoiser_input_fields`, so " - "it reaches the transformer's `video_self_attention_mask` argument without any change to " - "`LTX2LoopDenoiser`. Only run when the references carry a per-token strength (see " - "`LTX2AutoBuildVideoSelfAttentionMaskStep`); attention is otherwise left unmasked." + block_state.reference_token_counts = reference_token_counts + # Without reference tokens there is nothing to mask: leave the attention unmasked rather than pass all ones. + block_state.video_self_attention_mask = ( + _build_video_self_attention_mask( + block_state.latents, + base_token_count, + num_ref_tokens, + block_state.reference_latents, + reference_token_counts, + block_state.conditioning_attention_strength, + block_state.conditioning_attention_mask, + ) + if num_ref_tokens > 0 + else None ) - - @property - def inputs(self) -> list[InputParam]: - return [ - InputParam.template("latents", required=True), - InputParam( - "base_token_count", - type_hint=int, - required=True, - description="Number of generated-video tokens, i.e. the sequence length before appended tokens.", - ), - InputParam( - "num_ref_tokens", - type_hint=int, - required=True, - description="Number of reference tokens, which sit at the very end of the sequence.", - ), - InputParam( - "reference_cross_mask", - type_hint=torch.Tensor, - required=True, - description="Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens].", - ), - InputParam( - "reference_token_counts", - type_hint=list, - required=True, - description="Per-reference token counts, used to split `reference_cross_mask` into attention groups.", - ), - ] - - @property - def intermediate_outputs(self) -> list[OutputParam]: - return [ - OutputParam( - "video_self_attention_mask", - type_hint=torch.Tensor, - kwargs_type="denoiser_input_fields", - description="Multiplicative self-attention mask of shape [B, S, S] with values in [0, 1].", - ), - ] - - @torch.no_grad() - def __call__(self, components, state: PipelineState) -> PipelineState: - block_state = self.get_block_state(state) - device = components._execution_device - - batch_size, total_tokens, _ = block_state.latents.shape - # Cross-attention partners are the base tokens only. The prefix (base + keyframe tokens) attends fully to - # itself, but keyframe tokens are appended conditioning just like the references, so the two groups are - # masked off from each other. - num_base_tokens = block_state.base_token_count - num_prefix_tokens = total_tokens - block_state.num_ref_tokens - cross = block_state.reference_cross_mask.to(device=device, dtype=torch.float32) - - # Start from zeros so the keyframe<->reference and reference<->reference blocks stay masked without explicit - # assignment. Each guidance pass is its own single-batch forward, so this is built at the generation batch - # size -- the standard pipeline's expand to 2B for the CFG batch has no counterpart here. - attn_mask = torch.zeros((batch_size, total_tokens, total_tokens), device=device, dtype=torch.float32) - attn_mask[:, :num_prefix_tokens, :num_prefix_tokens] = 1.0 - - # One attention group per reference condition, in the order the encoder emitted their tokens. - offset = num_prefix_tokens - for group_cross in torch.split(cross, block_state.reference_token_counts, dim=1): - n = group_cross.shape[1] - attn_mask[:, :num_base_tokens, offset : offset + n] = group_cross.unsqueeze(1) - attn_mask[:, offset : offset + n, :num_base_tokens] = group_cross.unsqueeze(2) - attn_mask[:, offset : offset + n, offset : offset + n] = 1.0 - offset += n - - block_state.video_self_attention_mask = attn_mask + block_state.noise_scale = noise_scale self.set_block_state(state, block_state) return components, state @@ -1488,6 +1833,26 @@ def __call__(self, components, state: PipelineState) -> PipelineState: class LTX2ConditionSetTimestepsStep(ModularPipelineBlocks): model_name = "ltx2" + def __init__( + self, sigmas_name: str = "sigmas", timesteps_name: str = "timesteps", sigmas_default: list[float] | None = None + ): + """ + Args: + sigmas_name (`str`, defaults to `"sigmas"`): + Name of the input that holds this pass's sigma schedule. Lets a first-pass and a second-pass copy of + the block sit in one pipeline that takes both `sigmas` and `stage_2_sigmas`. + timesteps_name (`str`, defaults to `"timesteps"`): + Name of the input that holds this pass's custom timesteps, for the same reason. + sigmas_default (`list[float]`, *optional*): + Default sigma schedule of the pass. Set where a blockset assembles the block for a checkpoint that runs + a fixed schedule (the LTX-2.5 distilled recipe); the block then exposes no `num_inference_steps`. + `None` leaves the schedule to `num_inference_steps`. + """ + self._sigmas_name = sigmas_name + self._timesteps_name = timesteps_name + self._sigmas_default = sigmas_default + super().__init__() + @property def description(self) -> str: return ( @@ -1504,12 +1869,15 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: - return [ - InputParam.template("num_inference_steps", default=30), - InputParam.template("timesteps"), - InputParam.template("sigmas"), + inputs = [ + InputParam.template("timesteps", name=self._timesteps_name), + InputParam.template("sigmas", name=self._sigmas_name, default=self._sigmas_default), InputParam.template("latents", required=True), ] + # A block assembled with a fixed schedule has no step count to choose. + if self._sigmas_default is None: + inputs.insert(0, InputParam.template("num_inference_steps", default=30)) + return inputs @property def intermediate_outputs(self) -> list[OutputParam]: @@ -1527,8 +1895,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - num_inference_steps = block_state.num_inference_steps - sigmas = block_state.sigmas + num_inference_steps = getattr(block_state, "num_inference_steps", None) + timesteps = getattr(block_state, self._timesteps_name) + sigmas = getattr(block_state, self._sigmas_name) if sigmas is None: sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) @@ -1541,11 +1910,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ) block_state.audio_scheduler = copy.deepcopy(components.scheduler) - retrieve_timesteps( - block_state.audio_scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu - ) + retrieve_timesteps(block_state.audio_scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu) block_state.timesteps, block_state.num_inference_steps = retrieve_timesteps( - components.scheduler, num_inference_steps, device, block_state.timesteps, sigmas=sigmas, mu=mu + components.scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu ) # Set begin index to skip the nonzero().item() call in scheduler init, which triggers a GPU sync. @@ -1562,53 +1929,28 @@ class LTX2ConditionPrepareAudioLatentsStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Prepares the packed audio noise latents and derives the audio latent frame count, for condition-based " - "generation. Two deliberate differences from `LTX2PrepareAudioLatentsStep`:\n" - " 1. the noise is sampled directly in the packed shape [B, L, C * M], matching " - "`LTX2ConditionPipeline.prepare_audio_latents`. The text-to-video/image-to-video pipelines sample " - "unpacked [B, C, L, M] and pack afterwards; both draw the same number of values from the generator but " - "lay them out differently, so sampling the wrong way silently desynchronizes the audio noise (and, " - "through the joint attention, the video too).\n" - " 2. `noise_scale` is declared with a `None` default: a blockset keeps the first non-`None` default " - "across its blocks, so the text-to-video default of 0.0 would otherwise shadow the condition workflow's " - "`None -> sigmas[0] or 1.0` resolution. `LTX2ConditionPrepareLatentsStep` runs first and writes the " - "resolved value back to state." + "Samples the packed audio noise latents for a first pass of condition-based generation and derives the " + "audio latent frame count. One deliberate difference from `LTX2PrepareAudioLatentsStep`: the noise is " + "sampled directly in the packed shape [B, L, C * M], matching `LTX2ConditionPipeline." + "prepare_audio_latents`. The text-to-video/image-to-video pipelines sample unpacked [B, C, L, M] and " + "pack afterwards; both draw the same number of values from the generator but lay them out differently, " + "so sampling the wrong way silently desynchronizes the audio noise (and, through the joint attention, " + "the video too). Refining existing audio latents is `LTX2Stage2PrepareAudioLatentsStep`, which the " + "condition workflow shares with text-to-video." ) - @property - def expected_components(self) -> list[ComponentSpec]: - return [ComponentSpec("audio_vae", AutoencoderKLLTX2Audio)] - @property def inputs(self) -> list[InputParam]: return [ InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam( "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." ), - InputParam( - "audio_latents", - type_hint=torch.Tensor, - default=None, - description="Optional pre-encoded audio latents; random noise is used when not provided.", - ), - InputParam( - "noise_scale", - type_hint=float, - default=None, - description=( - "Initial noise level applied to any provided `audio_latents`. Resolved upstream by " - "`LTX2ConditionPrepareLatentsStep` (`sigmas[0]` when custom sigmas are supplied, else 1.0)." - ), - ), InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), InputParam.template("generator"), InputParam( @@ -1637,8 +1979,8 @@ def __call__(self, components, state: PipelineState) -> PipelineState: device = components._execution_device batch_size = block_state.batch_size * block_state.num_videos_per_prompt - num_channels_latents = components.audio_vae.config.latent_channels - latent_mel_bins = components.audio_vae.config.mel_bins // components.audio_vae_mel_compression_ratio + num_channels_latents = components.audio_latent_channels + latent_mel_bins = components.audio_latent_mel_bins duration_s = block_state.num_frames / block_state.frame_rate audio_latents_per_second = ( @@ -1648,24 +1990,12 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ) audio_num_frames = round(duration_s * audio_latents_per_second) - if block_state.audio_latents is not None: - audio_latents = block_state.audio_latents - if audio_latents.ndim == 4: - audio_num_frames = audio_latents.shape[2] - audio_latents = _pack_audio_latents(audio_latents) - audio_latents = _normalize_audio_latents( - audio_latents, components.audio_latents_mean, components.audio_latents_std - ) - audio_latents = _create_noised_state(audio_latents, block_state.noise_scale, block_state.generator) - block_state.audio_latents = audio_latents.to(device=device, dtype=torch.float32) - else: - # Sample directly in packed shape, following `LTX2ConditionPipeline.prepare_audio_latents` -- see the - # block description for why the unpacked-then-pack order used by text-to-video is not interchangeable. - packed_shape = (batch_size, audio_num_frames, num_channels_latents * latent_mel_bins) - block_state.audio_latents = randn_tensor( - packed_shape, generator=block_state.generator, device=device, dtype=torch.float32 - ) - + # Sample directly in packed shape, following `LTX2ConditionPipeline.prepare_audio_latents` -- see the block + # description for why the unpacked-then-pack order used by text-to-video is not interchangeable. + packed_shape = (batch_size, audio_num_frames, num_channels_latents * latent_mel_bins) + block_state.audio_latents = randn_tensor( + packed_shape, generator=block_state.generator, device=device, dtype=torch.float32 + ) block_state.audio_num_frames = audio_num_frames self.set_block_state(state, block_state) @@ -1696,11 +2026,8 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam( "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." @@ -1758,3 +2085,70 @@ def __call__(self, components, state: PipelineState) -> PipelineState: self.set_block_state(state, block_state) return components, state + + +class LTX2LatentUpsampleStep(ModularPipelineBlocks): + model_name = "ltx2" + + @property + def description(self) -> str: + return ( + "Spatially upsamples video latents by 2x with the `latent_upsampler`: the bridge between the two passes of " + "the two-stage recipe. Takes the normalized `[B, C, F, H, W]` latents a denoise pass leaves in state, " + "denormalizes them for the upsampler (which works on raw VAE latents) and re-normalizes the result, " + "handing back the same form at twice the resolution with `height` / `width` doubled to match. Matches " + '`LTX2LatentUpsamplePipeline` with `latents_normalized=True`, `output_type="latent"` and no AdaIN or ' + "tone mapping." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("latent_upsampler", LTX2LatentUpsamplerModel), + ComponentSpec("transformer", LTX2VideoTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latents", + type_hint=torch.Tensor, + required=True, + description="Video latents to upsample, of shape [B, C, F, H, W] (normalized, not packed).", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="Upsampled video latents of shape [B, C, F, 2H, 2W] (normalized, not packed).", + ), + OutputParam("height", type_hint=int, description="Height of the upsampled latents, in pixels."), + OutputParam("width", type_hint=int, description="Width of the upsampled latents, in pixels."), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + upsampler = components.latent_upsampler + + # The upsampler works on raw VAE latents: denormalize, upsample, re-normalize. + latents = _denormalize_latents( + block_state.latents, components.latents_mean, components.latents_std, components.vae_scaling_factor + ) + latents = upsampler(latents.to(device=upsampler.device, dtype=upsampler.dtype)) + latents = _normalize_latents( + latents, components.latents_mean, components.latents_std, components.vae_scaling_factor + ) + + block_state.latents = latents + # The second pass and any re-encoding ahead of it read the upsampled resolution from state. + block_state.height = latents.shape[-2] * components.vae_spatial_compression_ratio + block_state.width = latents.shape[-1] * components.vae_spatial_compression_ratio + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/ltx2/decoders.py b/src/diffusers/modular_pipelines/ltx2/decoders.py index fc957a3f9925..d9b289f35470 100644 --- a/src/diffusers/modular_pipelines/ltx2/decoders.py +++ b/src/diffusers/modular_pipelines/ltx2/decoders.py @@ -12,8 +12,6 @@ # See the License for the specific language governing permissions and # limitations under the License. -from typing import Any - import torch from ...configuration_utils import FrozenDict @@ -62,8 +60,11 @@ def _unpack_latents( def _denormalize_audio_latents( latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor ) -> torch.Tensor: - latents_mean = latents_mean.to(latents.device, latents.dtype) - latents_std = latents_std.to(latents.device, latents.dtype) + # Denormalizes audio latents of shape [B, C, L, M]. The statistics are stored per (channel, mel bin), flattened in + # the order the packed `[B, L, C * M]` layout uses, so they broadcast as [1, C, 1, M] here. + num_channels, num_mel_bins = latents.shape[1], latents.shape[3] + latents_mean = latents_mean.view(1, num_channels, 1, num_mel_bins).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, num_channels, 1, num_mel_bins).to(latents.device, latents.dtype) return (latents * latents_std) + latents_mean @@ -85,14 +86,98 @@ def _unpack_audio_latents( return latents +class LTX2UnpackLatentsStep(ModularPipelineBlocks): + model_name = "ltx2" + + @property + def description(self) -> str: + return ( + "Unpacks the denoised video and audio latents from the transformer's token layout back into the " + "`[B, C, F, H, W]` / `[B, C, L, M]` form the VAE encoders emit (still normalized). Closes every core " + "denoise group, so the decode and upsample blocks that follow take the same form the encoders produce " + "and need no geometry inputs." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "latents", + type_hint=torch.Tensor, + required=True, + description="Denoised video latents, packed and normalized, of shape [B, S, C].", + ), + InputParam( + "audio_latents", + type_hint=torch.Tensor, + required=True, + description="Denoised audio latents, packed and normalized, of shape [B, L, C * M].", + ), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam( + "num_frames", + type_hint=int, + required=True, + description="The number of frames in the generated video.", + ), + InputParam( + "audio_num_frames", + type_hint=int, + required=True, + description="Number of audio latent frames, used to unpack the audio latents.", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="Video latents of shape [B, C, F, H, W] (normalized, not packed).", + ), + OutputParam( + "audio_latents", + type_hint=torch.Tensor, + description="Audio latents of shape [B, C, L, M] (normalized, not packed).", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + latents = _unpack_latents( + block_state.latents, + latent_num_frames, + latent_height, + latent_width, + components.transformer_spatial_patch_size, + components.transformer_temporal_patch_size, + ) + block_state.latents = latents + + latent_mel_bins = components.audio_latent_mel_bins + block_state.audio_latents = _unpack_audio_latents( + block_state.audio_latents, block_state.audio_num_frames, num_mel_bins=latent_mel_bins + ) + + self.set_block_state(state, block_state) + return components, state + + class LTX2TrimConditionTokensStep(ModularPipelineBlocks): model_name = "ltx2" @property def description(self) -> str: return ( - "Drops the appended keyframe-condition tokens from the denoised latents, leaving only the " - "generated-video tokens for the decoders." + "Drops the appended keyframe-condition tokens from the denoised packed latents, leaving only the " + "generated-video tokens. Runs ahead of `LTX2UnpackLatentsStep`, which needs the plain video token grid." ) @property @@ -131,10 +216,10 @@ class LTX2DiffusionVaeDecoderStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Step that unpacks and decodes the denoised video latents with the LTX-2 diffusion decoder (or returns " - "latents). Swap this in for `LTX2VaeDecoderStep` on checkpoints that ship the diffusion decoder, which " - "from LTX-2.5 on is the native default. The decoder denoises rather than deterministically decoding, so " - "it draws its own noise from `generator` and takes no decode timestep." + "Decodes the video latents with the LTX-2 diffusion decoder. Swap this in for `LTX2VaeDecoderStep` on " + "checkpoints that ship the diffusion decoder, which from LTX-2.5 on is the native default. The decoder " + "denoises rather than deterministically decoding, so it draws its own noise from `generator` and takes " + "no decode timestep." ) @property @@ -150,17 +235,16 @@ def expected_components(self) -> list[ComponentSpec]: ] @property - def inputs(self) -> list[tuple[str, Any]]: + def inputs(self) -> list[InputParam]: return [ - InputParam.template("latents", required=True), - InputParam.template("output_type", default="pil"), - InputParam.template("height", default=512), - InputParam.template("width", default=704), InputParam( - "num_frames", type_hint=int, default=121, description="The number of frames in the generated video." + "latents", + type_hint=torch.Tensor, + required=True, + description="Video latents of shape [B, C, F, H, W] (normalized, not packed).", ), + InputParam.template("output_type", default="pil"), InputParam.template("generator"), - InputParam.template("dtype", required=True), ] @property @@ -172,29 +256,10 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) decoder = components.diffusion_decoder - latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 - latent_height = block_state.height // components.vae_spatial_compression_ratio - latent_width = block_state.width // components.vae_spatial_compression_ratio - - latents = _unpack_latents( - block_state.latents, - latent_num_frames, - latent_height, - latent_width, - components.transformer_spatial_patch_size, - components.transformer_temporal_patch_size, - ) - - if block_state.output_type == "latent": - block_state.videos = _denormalize_latents( - latents, components.latents_mean, components.latents_std, components.vae_scaling_factor - ) - self.set_block_state(state, block_state) - return components, state - - latents = latents.to(block_state.dtype) + # Denormalize in the loop's float32 before casting to the decoder dtype -- the same order as running the + # pipeline with `output_type="latent"` and decoding with `LTX2VideoDiffusionDecodePipeline`. latents = _denormalize_latents( - latents, components.latents_mean, components.latents_std, components.vae_scaling_factor + block_state.latents, components.latents_mean, components.latents_std, components.vae_scaling_factor ) latents = latents.to(decoder.dtype) # It samples the noise it denoises, so pass the generator to keep decoding reproducible. @@ -210,7 +275,7 @@ class LTX2VaeDecoderStep(ModularPipelineBlocks): @property def description(self) -> str: - return "Step that unpacks and decodes the denoised video latents into videos (or returns latents)." + return "Decodes the video latents with the video VAE into the output video." @property def expected_components(self) -> list[ComponentSpec]: @@ -225,21 +290,15 @@ def expected_components(self) -> list[ComponentSpec]: ] @property - def inputs(self) -> list[tuple[str, Any]]: + def inputs(self) -> list[InputParam]: return [ - InputParam.template("latents", required=True), - InputParam.template("output_type", default="pil"), - InputParam.template("height", default=512), - InputParam.template("width", default=704), InputParam( - "num_frames", - type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + "latents", + type_hint=torch.Tensor, + required=True, + description="Video latents of shape [B, C, F, H, W] (normalized, not packed).", ), + InputParam.template("output_type", default="pil"), InputParam( "decode_timestep", default=0.0, description="The timestep at which the VAE decodes the final latents." ), @@ -262,34 +321,9 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) vae = components.vae - latents = block_state.latents - height = block_state.height - width = block_state.width - num_frames = block_state.num_frames - - latent_num_frames = (num_frames - 1) // components.vae_temporal_compression_ratio + 1 - latent_height = height // components.vae_spatial_compression_ratio - latent_width = width // components.vae_spatial_compression_ratio - - latents = _unpack_latents( - latents, - latent_num_frames, - latent_height, - latent_width, - components.transformer_spatial_patch_size, - components.transformer_temporal_patch_size, - ) - - if block_state.output_type == "latent": - block_state.videos = _denormalize_latents( - latents, components.latents_mean, components.latents_std, components.vae_scaling_factor - ) - self.set_block_state(state, block_state) - return components, state - # LTX-2 applies the optional decode-time noise on the *normalized* latents, then denormalizes # (the reverse of LTX-1's decoder order). - latents = latents.to(block_state.dtype) + latents = block_state.latents.to(block_state.dtype) if not vae.config.timestep_conditioning: timestep = None else: @@ -328,10 +362,7 @@ class LTX2AudioDecoderStep(ModularPipelineBlocks): @property def description(self) -> str: - return ( - "Step that unpacks and decodes the denoised audio latents into a waveform via the audio VAE and vocoder " - "(or returns the unpacked audio latents when `output_type='latent'`)." - ) + return "Decodes the audio latents with the audio VAE into a mel spectrogram and vocodes it into a waveform." @property def expected_components(self) -> list[ComponentSpec]: @@ -343,16 +374,14 @@ def expected_components(self) -> list[ComponentSpec]: ] @property - def inputs(self) -> list[tuple[str, Any]]: + def inputs(self) -> list[InputParam]: return [ - InputParam("audio_latents", type_hint=torch.Tensor, required=True, description="Denoised audio latents."), InputParam( - "audio_num_frames", - type_hint=int, + "audio_latents", + type_hint=torch.Tensor, required=True, - description="Number of audio latent frames, used to unpack the audio latent sequence.", + description="Audio latents of shape [B, C, L, M] (normalized, not packed).", ), - InputParam.template("output_type", default="pil"), ] @property @@ -366,22 +395,12 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) audio_vae = components.audio_vae - num_mel_bins = audio_vae.config.mel_bins - latent_mel_bins = num_mel_bins // components.audio_vae_mel_compression_ratio - audio_latents = _denormalize_audio_latents( block_state.audio_latents, components.audio_latents_mean, components.audio_latents_std ) - audio_latents = _unpack_audio_latents( - audio_latents, block_state.audio_num_frames, num_mel_bins=latent_mel_bins - ) - - if block_state.output_type == "latent": - block_state.audio = audio_latents - else: - audio_latents = audio_latents.to(audio_vae.dtype) - generated_mel_spectrograms = audio_vae.decode(audio_latents, return_dict=False)[0] - block_state.audio = components.vocoder(generated_mel_spectrograms) + audio_latents = audio_latents.to(audio_vae.dtype) + generated_mel_spectrograms = audio_vae.decode(audio_latents, return_dict=False)[0] + block_state.audio = components.vocoder(generated_mel_spectrograms) self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/ltx2/denoise.py b/src/diffusers/modular_pipelines/ltx2/denoise.py index b1c4657d4d04..f8b3d791c274 100644 --- a/src/diffusers/modular_pipelines/ltx2/denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/denoise.py @@ -14,7 +14,6 @@ import inspect -from typing import Any import torch @@ -184,7 +183,7 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) # Default per-pass conditioning map for `LTX2LoopDenoiser`: transformer argument -> block-state attribute names # indexed [cond, uncond, stg, modality]. STG and modality-isolation reuse the conditional (positive) tensors and # differ only in their per-pass model flags, which the denoiser sets after preparation. -_DEFAULT_GUIDER_INPUT_FIELDS = { +_GUIDER_INPUT_FIELDS = { "encoder_hidden_states": ( "connector_prompt_embeds", "negative_connector_prompt_embeds", @@ -215,36 +214,12 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) class LTX2LoopDenoiser(ModularPipelineBlocks): model_name = "ltx2" - def __init__(self, guider_input_fields: dict[str, Any] = _DEFAULT_GUIDER_INPUT_FIELDS): - """Initialize the joint video+audio denoiser block for LTX-2.X. - - Args: - guider_input_fields: Maps each transformer argument (e.g. "encoder_hidden_states") to the block-state - attribute names the guiders read for each guidance pass. Each value is a 4-tuple of names indexed - [cond, uncond, stg, modality] -- for example {"encoder_hidden_states": ("connector_prompt_embeds", - "negative_connector_prompt_embeds", "connector_prompt_embeds", "connector_prompt_embeds")} reads the - positive embeds for the conditional/STG/modality passes and the negative embeds for the unconditional - pass. A guider builds only the passes it declares active, so a swapped-in guider that uses fewer passes - (e.g. `ClassifierFreeGuidance` -> cond/uncond) reads only the first two slots. - - Note: - `audio_num_frames`, `video_coords`, and `audio_coords` reach this block via the `denoiser_input_fields` tag - (their producer blocks declare `kwargs_type="denoiser_input_fields"`), not as named inputs. In a full - pipeline they arrive automatically; when running this block standalone (without those upstream blocks) they - must be passed through `denoiser_input_fields={...}` -- passing them as plain named kwargs is silently - ignored (modular.md's `kwargs_type` standalone gotcha). - """ - if not isinstance(guider_input_fields, dict): - raise ValueError(f"guider_input_fields must be a dictionary but is {type(guider_input_fields)}") - self._guider_input_fields = guider_input_fields - super().__init__() - @property def description(self) -> str: return ( "Joint video+audio denoiser. Runs the transformer once per guidance pass (each a single batch), with " "each pass's conditioning assembled by the guiders via `prepare_inputs_from_block_state` (driven by " - "`guider_input_fields`) and unioned across the video `guider` and audio `audio_guider`; the per-pass " + "`_GUIDER_INPUT_FIELDS`) and unioned across the video `guider` and audio `audio_guider`; the per-pass " "model flags (STG blocks, modality isolation) are set by identifier afterwards. Converts each pass's " "velocity to x0 and delegates the per-modality CFG + STG + modality-isolation combine to the two guiders." ) @@ -303,37 +278,51 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam( "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." ), + InputParam.template("attention_kwargs"), + ] + # The text conditioning the guiders read off block_state, per `_GUIDER_INPUT_FIELDS`. The negative tensors + # exist only under classifier-free guidance. + inputs += [ InputParam( - "use_cross_timestep", - type_hint=bool, - default=True, - description="Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+).", + "connector_prompt_embeds", + type_hint=torch.Tensor, + required=True, + description="Video-branch text conditioning (cond), expanded per prompt.", + ), + InputParam( + "connector_audio_prompt_embeds", + type_hint=torch.Tensor, + required=True, + description="Audio-branch text conditioning (cond), expanded per prompt.", + ), + InputParam( + "connector_attention_mask", + type_hint=torch.Tensor, + required=True, + description="Binary text attention mask (cond), expanded per prompt.", + ), + InputParam( + "negative_connector_prompt_embeds", + type_hint=torch.Tensor, + description="Video-branch text conditioning (uncond); read only under classifier-free guidance.", + ), + InputParam( + "negative_connector_audio_prompt_embeds", + type_hint=torch.Tensor, + description="Audio-branch text conditioning (uncond); read only under classifier-free guidance.", + ), + InputParam( + "negative_connector_attention_mask", + type_hint=torch.Tensor, + description="Binary text attention mask (uncond); read only under classifier-free guidance.", ), - InputParam.template("attention_kwargs"), ] - # The per-pass conditioning tensors the guiders read off block_state, declared from the field map so a - # custom `guider_input_fields` stays self-describing. - guider_input_names = [] - for value in self._guider_input_fields.values(): - guider_input_names.extend(value if isinstance(value, tuple) else (value,)) - for name in dict.fromkeys(guider_input_names): - inputs.append( - InputParam( - name, - type_hint=torch.Tensor, - required=True, - description="Per-pass text conditioning read by the guiders via `guider_input_fields`.", - ) - ) return inputs @torch.no_grad() @@ -353,7 +342,7 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) height=latent_height, width=latent_width, fps=block_state.frame_rate, - use_cross_timestep=block_state.use_cross_timestep, + use_cross_timestep=components.use_cross_timestep, attention_kwargs=block_state.attention_kwargs, perturbation_mask=None, ) @@ -362,13 +351,31 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) components.audio_guider.set_state(step=i, num_inference_steps=block_state.num_inference_steps, timestep=t) # Each guider maps block-state conditioning into one identifier-tagged batch per active pass via - # `_guider_input_fields` (transformer arg -> per-pass block-state attribute names, indexed + # `_GUIDER_INPUT_FIELDS` (transformer arg -> per-pass block-state attribute names, indexed # [cond, uncond, stg, modality]). A pass runs if *either* modality wants it, so union both guiders' batches # by identifier (same identifier => identical conditioning, built from the same map). identifier_key = LTX2Guidance._identifier_key + if any( + "pred_uncond" in guider.active_predictions() for guider in (components.guider, components.audio_guider) + ): + missing = [ + name + for name in ( + "negative_connector_prompt_embeds", + "negative_connector_audio_prompt_embeds", + "negative_connector_attention_mask", + ) + if getattr(block_state, name, None) is None + ] + if missing: + raise ValueError( + f"The guider runs classifier-free guidance but the unconditional conditioning {missing} is " + "missing. The text encoder produces it when the pipeline's guider has classifier-free guidance " + "enabled; when running the blocks separately, pass a `negative_prompt` to the text encoder." + ) batches_by_id = {} for guider in (components.guider, components.audio_guider): - for batch in guider.prepare_inputs_from_block_state(block_state, self._guider_input_fields): + for batch in guider.prepare_inputs_from_block_state(block_state, _GUIDER_INPUT_FIELDS): batches_by_id.setdefault(getattr(batch, identifier_key), batch) guider_state = list(batches_by_id.values()) @@ -401,7 +408,7 @@ def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor) # against `LTX2Pipeline` on fp32 and treat bf16 as close-but-not-bitwise. for batch in guider_state: components.guider.prepare_models(components.transformer) - cond_kwargs = {name: getattr(batch, name) for name in self._guider_input_fields} + cond_kwargs = {name: getattr(batch, name) for name in _GUIDER_INPUT_FIELDS} cond_kwargs["spatio_temporal_guidance_blocks"] = batch.spatio_temporal_guidance_blocks cond_kwargs["isolate_modalities"] = batch.isolate_modalities with components.transformer.cache_context(getattr(batch, identifier_key)): @@ -505,11 +512,8 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), ] diff --git a/src/diffusers/modular_pipelines/ltx2/encoders.py b/src/diffusers/modular_pipelines/ltx2/encoders.py index b261597c0f68..9bcf81762ef0 100644 --- a/src/diffusers/modular_pipelines/ltx2/encoders.py +++ b/src/diffusers/modular_pipelines/ltx2/encoders.py @@ -21,7 +21,7 @@ from transformers import PreTrainedModel, PreTrainedTokenizerBase, ProcessorMixin from ...configuration_utils import FrozenDict -from ...models import AutoencoderKLLTX2Video, LTX2VideoTransformer3DModel +from ...models import AutoencoderKLLTX2Video # NOTE (modular.md gotcha #1): `LTX2TextConnectors`, `LTX2DurationHead`, `LTX2VideoCondition`, # `LTX2ReferenceCondition`, the prompt-enhancement config/helpers, and the system prompts live under @@ -44,7 +44,6 @@ LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT, LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT, apply_image_conditioning_crf, - resolve_default_image_crf, ) from ...utils import logging from ...video_processor import VideoProcessor @@ -176,14 +175,6 @@ def expected_components(self) -> list[ComponentSpec]: def inputs(self) -> list[InputParam]: return [ InputParam.template("prompt", required=True), - InputParam( - "enable_prompt_enhancement", - type_hint=bool, - default=False, - description=( - "Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines." - ), - ), InputParam( "system_prompt", type_hint=str, @@ -224,9 +215,6 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - if not block_state.enable_prompt_enhancement: - self.set_block_state(state, block_state) # leave prompt unchanged - return components, state if getattr(components, "prompt_enhancer", None) is None: raise ValueError( "`enable_prompt_enhancement=True` but no `prompt_enhancer` component is loaded. Load a " @@ -273,14 +261,6 @@ def inputs(self) -> list[InputParam]: return [ InputParam.template("prompt", required=True), InputParam.template("image", required=True), - InputParam( - "enable_prompt_enhancement", - type_hint=bool, - default=False, - description=( - "Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines." - ), - ), InputParam( "system_prompt", type_hint=str, @@ -321,9 +301,6 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - if not block_state.enable_prompt_enhancement: - self.set_block_state(state, block_state) # leave prompt unchanged - return components, state if getattr(components, "prompt_enhancer", None) is None: raise ValueError( "`enable_prompt_enhancement=True` but no `prompt_enhancer` component is loaded. Load a " @@ -379,14 +356,6 @@ def inputs(self) -> list[InputParam]: "of the generated video." ), ), - InputParam( - "enable_prompt_enhancement", - type_hint=bool, - default=False, - description=( - "Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines." - ), - ), InputParam( "system_prompt", type_hint=str, @@ -430,9 +399,6 @@ def intermediate_outputs(self) -> list[OutputParam]: def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - if not block_state.enable_prompt_enhancement: - self.set_block_state(state, block_state) # leave prompt unchanged - return components, state if getattr(components, "prompt_enhancer", None) is None: raise ValueError( "`enable_prompt_enhancement=True` but no `prompt_enhancer` component is loaded. Load a " @@ -484,15 +450,16 @@ class LTX2TextEncoderStep(ModularPipelineBlocks): @property def description(self) -> str: return ( - "Text encoder step. Encodes `prompt` and `negative_prompt` into packed per-layer Gemma hidden states " - "that the connectors adapt for the video and audio branches, and reports the prompt count (`batch_size`) " - "and embedding `dtype`." + "Text encoder step. Encodes `prompt` -- and, when needed, `negative_prompt` -- into packed per-layer " + "Gemma hidden states that the connectors adapt for the video and audio branches. Whether the negative " + "prompt is encoded follows the guiders: with a guider registered and running classifier-free guidance " + "it is always encoded (defaulting to the empty prompt); with a guider registered but not running CFG " + "it is skipped, with a warning if one was passed; with no guider registered (the block on its own) it is " + "encoded only when a `negative_prompt` is passed." ) @property def expected_components(self) -> list[ComponentSpec]: - # No `guider`: LTX-2 applies CFG (+ STG + modality-isolation) manually in the denoise loop, so the encoder - # always produces both conditional and unconditional embeddings and the denoiser decides what to use. return [ ComponentSpec("text_encoder", PreTrainedModel), ComponentSpec("tokenizer", PreTrainedTokenizerBase), @@ -522,19 +489,13 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "negative_prompt_embeds", type_hint=torch.Tensor, - description="Packed per-layer Gemma hidden states for the negative prompt.", + description="Packed per-layer Gemma hidden states for the negative prompt, `None` when not encoded.", ), OutputParam( "negative_prompt_attention_mask", type_hint=torch.Tensor, - description="Binary attention mask for `negative_prompt_embeds`.", + description="Binary attention mask for `negative_prompt_embeds`, `None` when not encoded.", ), - OutputParam( - "batch_size", - type_hint=int, - description="The number of prompts being denoised (before per-prompt expansion).", - ), - OutputParam("dtype", type_hint=torch.dtype, description="The dtype of the prompt embeddings."), ] @staticmethod @@ -556,14 +517,27 @@ def __call__(self, components, state: PipelineState) -> PipelineState: components, prompt, max_sequence_length, device, dtype ) - negative_prompt = block_state.negative_prompt or "" - negative_prompt = len(prompt) * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt - block_state.negative_prompt_embeds, block_state.negative_prompt_attention_mask = _get_gemma_prompt_embeds( - components, negative_prompt, max_sequence_length, device, dtype + negative_prompt = block_state.negative_prompt + has_guider = ( + getattr(components, "guider", None) is not None or getattr(components, "audio_guider", None) is not None ) - - block_state.batch_size = block_state.prompt_embeds.shape[0] - block_state.dtype = block_state.prompt_embeds.dtype + block_state.negative_prompt_embeds = block_state.negative_prompt_attention_mask = None + if has_guider and not components.requires_unconditional_embeds: + # If guider is registered and not running CFG, nothing would consume + # unconditional embeddings. + if negative_prompt is not None: + logger.warning( + "`negative_prompt` was passed but the guider is not running classifier-free guidance, so it is " + "ignored." + ) + elif has_guider or negative_prompt is not None: + # Classifier-free guidance (the negative prompt defaults to the empty prompt), or the block on its own + # (without the guider) and with an explicit negative prompt. + negative_prompt = negative_prompt or "" + negative_prompt = len(prompt) * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + block_state.negative_prompt_embeds, block_state.negative_prompt_attention_mask = _get_gemma_prompt_embeds( + components, negative_prompt, max_sequence_length, device, dtype + ) self.set_block_state(state, block_state) return components, state @@ -576,24 +550,20 @@ class LTX2TextConnectorStep(ModularPipelineBlocks): def description(self) -> str: return ( "Connector step. Adapts the Gemma hidden states into the separate video- and audio-branch text " - "conditioning consumed by the transformer, for both the conditional and unconditional prompts." + "conditioning consumed by the transformer." ) @property def expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("connectors", LTX2TextConnectors), - # Declared only to read `padding_side` (used by the LTX-2.0 connector branch); LTX-2.5 defaults to "left". - ComponentSpec("tokenizer", PreTrainedTokenizerBase), - ] + return [ComponentSpec("connectors", LTX2TextConnectors)] @property def inputs(self) -> list[InputParam]: return [ InputParam("prompt_embeds", type_hint=torch.Tensor, required=True), InputParam("prompt_attention_mask", type_hint=torch.Tensor, required=True), - InputParam("negative_prompt_embeds", type_hint=torch.Tensor, required=True), - InputParam("negative_prompt_attention_mask", type_hint=torch.Tensor, required=True), + InputParam("negative_prompt_embeds", type_hint=torch.Tensor), + InputParam("negative_prompt_attention_mask", type_hint=torch.Tensor), ] @property @@ -616,46 +586,45 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "negative_connector_prompt_embeds", type_hint=torch.Tensor, - description="Video-branch text conditioning (uncond).", + description="Video-branch text conditioning (uncond), `None` when no negative prompt was encoded.", ), OutputParam( "negative_connector_audio_prompt_embeds", type_hint=torch.Tensor, - description="Audio-branch text conditioning (uncond).", + description="Audio-branch text conditioning (uncond), `None` when no negative prompt was encoded.", ), OutputParam( "negative_connector_attention_mask", type_hint=torch.Tensor, - description="Binary text attention mask (uncond).", + description="Binary text attention mask (uncond), `None` when no negative prompt was encoded.", ), ] @torch.no_grad() def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) - padding_side = components.tokenizer.padding_side - - # Run the connector once on the CFG-concatenated `[uncond, cond]` batch, matching the standard pipeline - # (`LTX2Pipeline` concatenates before the single `self.connectors(...)` call). The connector is applied per - # batch element, so cond/uncond are mathematically independent either way, but a single batched call keeps the - # results bitwise-identical to the standard pipeline: the connector's GEMM/attention kernels round the same for - # a given row only at batch >= 2, so running the branches separately would diverge by ~1e-6 at batch size 1. - num_negative = block_state.negative_prompt_embeds.shape[0] - prompt_embeds = torch.cat([block_state.negative_prompt_embeds, block_state.prompt_embeds], dim=0) - prompt_attention_mask = torch.cat( - [block_state.negative_prompt_attention_mask, block_state.prompt_attention_mask], dim=0 - ) - connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = components.connectors( - prompt_embeds, prompt_attention_mask, padding_side=padding_side - ) + padding_side = "left" - # Split back into uncond (first `num_negative`) and cond (rest). - block_state.negative_connector_prompt_embeds = connector_prompt_embeds[:num_negative] - block_state.negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:num_negative] - block_state.negative_connector_attention_mask = connector_attention_mask[:num_negative] - block_state.connector_prompt_embeds = connector_prompt_embeds[num_negative:] - block_state.connector_audio_prompt_embeds = connector_audio_prompt_embeds[num_negative:] - block_state.connector_attention_mask = connector_attention_mask[num_negative:] + ( + block_state.connector_prompt_embeds, + block_state.connector_audio_prompt_embeds, + block_state.connector_attention_mask, + ) = components.connectors( + block_state.prompt_embeds, block_state.prompt_attention_mask, padding_side=padding_side + ) + block_state.negative_connector_prompt_embeds = None + block_state.negative_connector_audio_prompt_embeds = None + block_state.negative_connector_attention_mask = None + if block_state.negative_prompt_embeds is not None: + ( + block_state.negative_connector_prompt_embeds, + block_state.negative_connector_audio_prompt_embeds, + block_state.negative_connector_attention_mask, + ) = components.connectors( + block_state.negative_prompt_embeds, + block_state.negative_prompt_attention_mask, + padding_side=padding_side, + ) self.set_block_state(state, block_state) return components, state @@ -679,6 +648,15 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ + InputParam( + "num_frames", + type_hint=int, + default=None, + description=( + "The number of frames in the generated video. Omit to have this step predict it with the " + "`duration_head`; the denoise blocks then take the predicted count." + ), + ), InputParam( "min_seconds", type_hint=float, @@ -706,12 +684,6 @@ def inputs(self) -> list[InputParam]: required=True, description="Audio-branch text conditioning from the connector (positive prompt).", ), - InputParam( - "batch_size", - type_hint=int, - required=True, - description="The number of prompts being denoised, used to expand conditioning per prompt.", - ), ] @property @@ -732,9 +704,10 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ) # The head predicts one duration; prompts with different natural lengths cannot share a single frame count. - if block_state.batch_size > 1: + batch_size = block_state.connector_prompt_embeds.shape[0] + if batch_size > 1: raise ValueError( - f"`num_frames` was omitted so the duration head would auto-predict, but {block_state.batch_size} " + f"`num_frames` was omitted so the duration head would auto-predict, but {batch_size} " "prompts were supplied. The duration head predicts one duration -- run one prompt at a time, or pass " "`num_frames` as an integer." ) @@ -779,62 +752,12 @@ def retrieve_latents( def _normalize_latents( latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 ) -> torch.Tensor: - # Normalize video latents across the channel dimension [B, C, F, H, W]. latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) latents = (latents - latents_mean) * scaling_factor / latents_std return latents -def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: - batch_size, num_channels, num_frames, height, width = latents.shape - latents = latents.reshape( - batch_size, - -1, - num_frames // patch_size_t, - patch_size_t, - height // patch_size, - patch_size, - width // patch_size, - patch_size, - ) - latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) - return latents - - -def _downsample_mask_to_latent( - mask: torch.Tensor, latent_num_frames: int, latent_height: int, latent_width: int -) -> torch.Tensor: - """ - Downsample a pixel-space attention mask of shape `(B, 1, F, H, W)` (values in `[0, 1]`) to a flattened per-token - latent-space mask of shape `(B, latent_num_frames * latent_height * latent_width)`. Spatial downsampling is area - interpolation per frame; temporal downsampling is causal (the first frame is kept as-is). - """ - if mask.ndim != 5 or mask.shape[1] != 1: - raise ValueError(f"Expected `conditioning_attention_mask` of shape (B, 1, F, H, W), got {tuple(mask.shape)}.") - b, _, f_pix, _, _ = mask.shape - - mask_2d = mask.reshape(b * f_pix, 1, mask.shape[-2], mask.shape[-1]) - spatial_down = torch.nn.functional.interpolate(mask_2d, size=(latent_height, latent_width), mode="area") - spatial_down = spatial_down.reshape(b, 1, f_pix, latent_height, latent_width) - - first_frame = spatial_down[:, :, :1, :, :] - if f_pix > 1 and latent_num_frames > 1: - t = (f_pix - 1) // (latent_num_frames - 1) - if (f_pix - 1) % (latent_num_frames - 1) != 0: - raise ValueError( - f"Pixel frames ({f_pix}) not compatible with latent frames ({latent_num_frames}): " - f"(f_pix - 1) must be divisible by (latent_num_frames - 1)." - ) - rest = spatial_down[:, :, 1:, :, :] - rest = rest.reshape(b, 1, latent_num_frames - 1, t, latent_height, latent_width).mean(dim=3) - latent_mask = torch.cat([first_frame, rest], dim=2) - else: - latent_mask = first_frame - - return latent_mask.reshape(b, latent_num_frames * latent_height * latent_width) - - class LTX2VaeEncoderStep(ModularPipelineBlocks): model_name = "ltx2" @@ -846,8 +769,6 @@ def description(self) -> str: def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKLLTX2Video), - # Only used to resolve the default `image_crf` from the text-encoder generation. - ComponentSpec("text_encoder", PreTrainedModel), ComponentSpec( "video_processor", VideoProcessor, @@ -868,8 +789,8 @@ def inputs(self) -> list[InputParam]: default=None, description=( "H.264 CRF used to re-compress the conditioning `image` before VAE encode, matching the " - "compression the model was trained against. `None` (default) resolves from the text-encoder " - "generation (33 through LTX-2.3, 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a " + "compression the model was trained against. `None` (default) uses the pipeline's " + "`default_image_crf` (33 through LTX-2.3, 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a " "`PIL.Image.Image` when re-compression runs." ), ), @@ -882,7 +803,10 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "image_latents", type_hint=torch.Tensor, - description="Normalized image latents (a single latent frame) for image-to-video conditioning.", + description=( + "Image latents for image-to-video conditioning: a single latent frame of shape [B, C, 1, H, W] " + "(normalized, not packed)." + ), ), ] @@ -894,11 +818,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: image = block_state.image if not isinstance(image, torch.Tensor): # H.264 re-compress before resize/normalize (ltx-pipelines `load_image_and_preprocess`). - crf = ( - block_state.image_crf - if block_state.image_crf is not None - else resolve_default_image_crf(components.text_encoder) - ) + crf = block_state.image_crf if block_state.image_crf is not None else components.default_image_crf if crf != 0: if not isinstance(image, PIL.Image.Image): raise ValueError( @@ -951,8 +871,6 @@ def expected_components(self) -> list[ComponentSpec]: # anti-alias prefilters on downscale. The reference uses a plain `F.interpolate`, reproduced in `__call__`. return [ ComponentSpec("vae", AutoencoderKLLTX2Video), - # Only used to resolve a condition's default `crf` from the text-encoder generation. - ComponentSpec("text_encoder", PreTrainedModel), ] @property @@ -972,11 +890,8 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), + required=True, + description="The number of frames in the generated video.", ), InputParam.template("generator"), ] @@ -987,7 +902,7 @@ def intermediate_outputs(self) -> list[OutputParam]: OutputParam( "condition_latents", type_hint=list, - description="Per-condition normalized VAE latents of shape [1, C, F, H, W].", + description="Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed).", ), OutputParam("condition_strengths", type_hint=list, description="Per-condition conditioning strengths."), OutputParam( @@ -1049,9 +964,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: # Single-frame image keyframes are H.264 re-compressed at the model CRF (ltx-pipelines # `ImageConditioner.resolve_crf` + `media_io.preprocess`). Multi-frame video conditions are not. if arr.shape[0] == 1: - crf = ( - condition.crf if condition.crf is not None else resolve_default_image_crf(components.text_encoder) - ) + crf = condition.crf if condition.crf is not None else components.default_image_crf if crf != 0 and arr.dtype != np.uint8: raise ValueError( f"Image conditioning CRF expects a uint8 RGB frame, got dtype={arr.dtype}. " @@ -1122,17 +1035,14 @@ class LTX2ReferenceEncoderStep(ModularPipelineBlocks): def description(self) -> str: return ( "Reference encoder step for in-context (IC-LoRA) generation. Preprocesses each reference video to the " - "(optionally downscaled) target resolution, VAE-encodes and packs it into tokens, and computes the " - "positional coordinates that map those tokens into the target coordinate space. When " - "`conditioning_attention_strength < 1.0` or a pixel-space `conditioning_attention_mask` is supplied it " - "also produces the per-token cross-attention strengths driving the video self-attention mask." + "(optionally downscaled) target resolution and VAE-encodes it into normalized latents, one entry per " + "reference." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("vae", AutoencoderKLLTX2Video), - ComponentSpec("transformer", LTX2VideoTransformer3DModel), ComponentSpec( "video_processor", VideoProcessor, @@ -1159,27 +1069,7 @@ def inputs(self) -> list[InputParam]: default=1, description=( "Ratio between the target and reference resolutions; 2 means the reference is preprocessed at " - "half the target resolution. Spatial coordinates are scaled by this factor so the reference " - "tokens land in the target coordinate space. Must match the factor the IC-LoRA was trained with." - ), - ), - InputParam( - "conditioning_attention_strength", - type_hint=float, - default=1.0, - description=( - "Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each " - "other. 1.0 (default) leaves attention unmasked." - ), - ), - InputParam( - "conditioning_attention_mask", - type_hint=torch.Tensor, - default=None, - description=( - "Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially " - "varying attention strength. Downsampled to the reference's latent grid and multiplied by " - "`conditioning_attention_strength`." + "half the target resolution. Must match the factor the IC-LoRA was trained with." ), ), InputParam.template("height", default=512), @@ -1187,14 +1077,8 @@ def inputs(self) -> list[InputParam]: InputParam( "num_frames", type_hint=int, - default=None, - description=( - "The number of frames in the generated video. Omit to auto-predict via the `duration_head` " - "(see `LTX2AutoDurationStep`)." - ), - ), - InputParam( - "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." + required=True, + description="The number of frames in the generated video.", ), InputParam.template("generator"), ] @@ -1204,25 +1088,10 @@ def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam( "reference_latents", - type_hint=torch.Tensor, - description="Packed reference tokens of shape [1, total_reference_tokens, C].", - ), - OutputParam( - "reference_coords", - type_hint=torch.Tensor, - description="RoPE coordinates for the reference tokens, of shape [1, 3, total_reference_tokens, 2].", - ), - OutputParam( - "reference_token_counts", type_hint=list, - description="Per-reference token counts, in `reference_conditions` order.", - ), - OutputParam( - "reference_cross_mask", - type_hint=torch.Tensor, description=( - "Per-reference-token noisy<->reference attention strengths of shape [1, " - "total_reference_tokens], or `None` when attention is left unmasked." + "Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed), in " + "`reference_conditions` order." ), ), ] @@ -1241,16 +1110,11 @@ def __call__(self, components, state: PipelineState) -> PipelineState: "take no reference video -- `LTX2AutoReferenceEncoderStep` then skips this step." ) - downscale_factor = block_state.reference_downscale_factor - ref_height = block_state.height // downscale_factor - ref_width = block_state.width // downscale_factor - strength = block_state.conditioning_attention_strength - attention_mask = block_state.conditioning_attention_mask - # An all-ones mask at full strength is the same as no mask, so only materialize one when it can bite. - mask_needed = strength < 1.0 or attention_mask is not None + ref_height = block_state.height // block_state.reference_downscale_factor + ref_width = block_state.width // block_state.reference_downscale_factor generator = block_state.generator[0] if isinstance(block_state.generator, list) else block_state.generator - all_latents, all_coords, all_cross_masks, token_counts = [], [], [], [] + reference_latents = [] for ref_cond in reference_conditions: if isinstance(ref_cond.frames, PIL.Image.Image): video_like = [ref_cond.frames] @@ -1268,45 +1132,13 @@ def __call__(self, components, state: PipelineState) -> PipelineState: ref_pixels = ref_pixels.to(dtype=components.vae.dtype, device=device) ref_latent = retrieve_latents(components.vae.encode(ref_pixels), generator=generator, sample_mode="argmax") - ref_latent = _normalize_latents(ref_latent, components.latents_mean, components.latents_std).to( - device=device, dtype=torch.float32 - ) - _, _, ref_latent_frames, ref_latent_height, ref_latent_width = ref_latent.shape - ref_latent_packed = _pack_latents( - ref_latent, components.transformer_spatial_patch_size, components.transformer_temporal_patch_size + reference_latents.append( + _normalize_latents(ref_latent, components.latents_mean, components.latents_std).to( + device=device, dtype=torch.float32 + ) ) - # Coordinates are computed on the reference's own latent grid, then scaled spatially so the tokens map - # into the target's coordinate space (preserving the positional relationship the IC-LoRA was trained on). - ref_coords = components.transformer.rope.prepare_video_coords( - batch_size=1, - num_frames=ref_latent_frames, - height=ref_latent_height, - width=ref_latent_width, - device=device, - fps=block_state.frame_rate, - ) - if downscale_factor != 1: - ref_coords[:, 1, :, :] = ref_coords[:, 1, :, :] * downscale_factor - ref_coords[:, 2, :, :] = ref_coords[:, 2, :, :] * downscale_factor - - if mask_needed: - if attention_mask is not None: - ref_cross = _downsample_mask_to_latent( - attention_mask, ref_latent_frames, ref_latent_height, ref_latent_width - ).to(device=device, dtype=torch.float32) - else: - ref_cross = torch.ones((1, ref_latent_packed.shape[1]), device=device, dtype=torch.float32) - all_cross_masks.append(ref_cross * strength) - - all_latents.append(ref_latent_packed) - all_coords.append(ref_coords) - token_counts.append(ref_latent_packed.shape[1]) - - block_state.reference_latents = torch.cat(all_latents, dim=1) - block_state.reference_coords = torch.cat(all_coords, dim=2) - block_state.reference_token_counts = token_counts - block_state.reference_cross_mask = torch.cat(all_cross_masks, dim=1) if mask_needed else None + block_state.reference_latents = reference_latents self.set_block_state(state, block_state) return components, state diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py index 86428328a5a6..b0c48219c9b4 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2.py @@ -16,9 +16,8 @@ from ...utils import logging from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import OutputParam +from ..modular_pipeline_utils import InputParam, OutputParam from .before_denoise import ( - LTX2BuildVideoSelfAttentionMaskStep, LTX2ConditionPrepareAudioLatentsStep, LTX2ConditionPrepareCoordsStep, LTX2ConditionPrepareLatentsStep, @@ -34,6 +33,7 @@ from .decoders import ( LTX2AudioDecoderStep, LTX2TrimConditionTokensStep, + LTX2UnpackLatentsStep, LTX2VaeDecoderStep, ) from .denoise import LTX2ConditionDenoiseStep, LTX2DenoiseStep, LTX2Image2VideoDenoiseStep @@ -73,8 +73,6 @@ class LTX2AutoPromptEnhancerStep(ConditionalPipelineBlocks): conditions (`list`, *optional*): `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the generated video. - enable_prompt_enhancement (`bool`, *optional*, defaults to False): - Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. system_prompt (`str`, *optional*): System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. @@ -89,6 +87,8 @@ class LTX2AutoPromptEnhancerStep(ConditionalPipelineBlocks): Torch generator for deterministic generation. image (`Image | list`, *optional*): Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. Outputs: prompt (`list`): @@ -100,6 +100,20 @@ class LTX2AutoPromptEnhancerStep(ConditionalPipelineBlocks): block_names = ["condition", "image2video", "text2video"] block_trigger_inputs = ["conditions", "image", "enable_prompt_enhancement"] + @property + def inputs(self): + # The trigger belongs to this wrapper, not to the enhancer steps: each of them always enhances. + inputs = super().inputs + inputs.append( + InputParam( + "enable_prompt_enhancement", + type_hint=bool, + default=False, + description="Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines.", + ) + ) + return inputs + def select_block(self, conditions=None, image=None, enable_prompt_enhancement=False) -> str | None: # `conditions` is checked before `image`: the condition and in-context workflows place their reference # frames in `conditions` and never take a raw `image`. @@ -146,13 +160,9 @@ class LTX2TextConditioningStep(SequentialPipelineBlocks): prompt_attention_mask (`Tensor`): Binary attention mask for `prompt_embeds`. negative_prompt_embeds (`Tensor`): - Packed per-layer Gemma hidden states for the negative prompt. + Packed per-layer Gemma hidden states for the negative prompt, `None` when not encoded. negative_prompt_attention_mask (`Tensor`): - Binary attention mask for `negative_prompt_embeds`. - batch_size (`int`): - The number of prompts being denoised (before per-prompt expansion). - dtype (`dtype`): - The dtype of the prompt embeddings. + Binary attention mask for `negative_prompt_embeds`, `None` when not encoded. connector_prompt_embeds (`Tensor`): Video-branch text conditioning (cond). connector_audio_prompt_embeds (`Tensor`): @@ -160,11 +170,11 @@ class LTX2TextConditioningStep(SequentialPipelineBlocks): connector_attention_mask (`Tensor`): Binary text attention mask (cond). negative_connector_prompt_embeds (`Tensor`): - Video-branch text conditioning (uncond). + Video-branch text conditioning (uncond), `None` when no negative prompt was encoded. negative_connector_audio_prompt_embeds (`Tensor`): - Audio-branch text conditioning (uncond). + Audio-branch text conditioning (uncond), `None` when no negative prompt was encoded. negative_connector_attention_mask (`Tensor`): - Binary text attention mask (uncond). + Binary text attention mask (uncond), `None` when no negative prompt was encoded. """ model_name = "ltx2" @@ -192,6 +202,9 @@ class LTX2AutoDurationStep(ConditionalPipelineBlocks): duration_head (`LTX2DurationHead`) Inputs: + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. min_seconds (`float`, *optional*, defaults to 1.0): Lower bound on the auto-predicted duration. max_seconds (`float`, *optional*, defaults to 20.0): @@ -202,8 +215,6 @@ class LTX2AutoDurationStep(ConditionalPipelineBlocks): Video-branch text conditioning from the connector (positive prompt). connector_audio_prompt_embeds (`Tensor`, *optional*): Audio-branch text conditioning from the connector (positive prompt). - batch_size (`int`, *optional*): - The number of prompts being denoised, used to expand conditioning per prompt. Outputs: num_frames (`int`): @@ -235,7 +246,7 @@ class LTX2AutoVaeEncoderStep(AutoPipelineBlocks): - Skipped otherwise. Components: - vae (`AutoencoderKLLTX2Video`) text_encoder (`PreTrainedModel`) video_processor (`VideoProcessor`) + vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) Inputs: image (`Image | list`, *optional*): @@ -246,15 +257,15 @@ class LTX2AutoVaeEncoderStep(AutoPipelineBlocks): The width in pixels of the generated image. image_crf (`int`, *optional*): H.264 CRF used to re-compress the conditioning `image` before VAE encode, matching the compression the - model was trained against. `None` (default) resolves from the text-encoder generation (33 through - LTX-2.3, 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when - re-compression runs. + model was trained against. `None` (default) uses the pipeline's `default_image_crf` (33 through LTX-2.3, + 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when re-compression runs. generator (`Generator`, *optional*): Torch generator for deterministic generation. Outputs: image_latents (`Tensor`): - Normalized image latents (a single latent frame) for image-to-video conditioning. + Image latents for image-to-video conditioning: a single latent frame of shape [B, C, 1, H, W] + (normalized, not packed). """ model_name = "ltx2" @@ -278,8 +289,8 @@ class LTX2CoreDenoiseStep(SequentialPipelineBlocks): latents and runs the joint denoising loop. Components: - scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) audio_vae - (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) Inputs: num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -290,12 +301,12 @@ class LTX2CoreDenoiseStep(SequentialPipelineBlocks): Audio-branch text conditioning (cond). connector_attention_mask (`Tensor`): Binary text attention mask (cond). - negative_connector_prompt_embeds (`Tensor`): - Video-branch text conditioning (uncond). - negative_connector_audio_prompt_embeds (`Tensor`): - Audio-branch text conditioning (uncond). - negative_connector_attention_mask (`Tensor`): - Binary text attention mask (uncond). + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond), `None` without classifier-free guidance. num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. timesteps (`Tensor`, *optional*): @@ -306,28 +317,14 @@ class LTX2CoreDenoiseStep(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. - noise_scale (`float`, *optional*): - Interpolation factor between random noise and any provided latents. `None` (default) resolves to 0.0, - which keeps the provided latents. + num_frames (`int`): + The number of frames in the generated video. generator (`Generator`, *optional*): Torch generator for deterministic generation. - batch_size (`int`): - The number of prompts being denoised, used to expand conditioning per prompt. frame_rate (`float`, *optional*, defaults to 24.0): Frames per second of the generated video. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. - dtype (`dtype`): - The dtype the model inputs are cast to. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. @@ -346,6 +343,7 @@ class LTX2CoreDenoiseStep(SequentialPipelineBlocks): LTX2PrepareAudioLatentsStep, LTX2PrepareCoordsStep, LTX2DenoiseStep, + LTX2UnpackLatentsStep, ] block_names = [ "input", @@ -354,6 +352,7 @@ class LTX2CoreDenoiseStep(SequentialPipelineBlocks): "prepare_audio_latents", "prepare_coords", "denoise", + "unpack", ] @property @@ -378,8 +377,8 @@ class LTX2Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): conditioning and runs the joint denoising loop. Components: - scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) audio_vae - (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) Inputs: num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -390,12 +389,12 @@ class LTX2Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): Audio-branch text conditioning (cond). connector_attention_mask (`Tensor`): Binary text attention mask (cond). - negative_connector_prompt_embeds (`Tensor`): - Video-branch text conditioning (uncond). - negative_connector_audio_prompt_embeds (`Tensor`): - Audio-branch text conditioning (uncond). - negative_connector_attention_mask (`Tensor`): - Binary text attention mask (uncond). + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond), `None` without classifier-free guidance. num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. timesteps (`Tensor`, *optional*): @@ -406,30 +405,16 @@ class LTX2Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. - noise_scale (`float`, *optional*): - Interpolation factor between random noise and any provided latents. `None` (default) resolves to 0.0, - which keeps the provided latents. + num_frames (`int`): + The number of frames in the generated video. generator (`Generator`, *optional*): Torch generator for deterministic generation. - batch_size (`int`): - The number of prompts being denoised, used to expand conditioning per prompt. image_latents (`Tensor`): VAE-encoded reference-image latents used for image-to-video conditioning. frame_rate (`float`, *optional*, defaults to 24.0): Frames per second of the generated video. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. - dtype (`dtype`): - The dtype the model inputs are cast to. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. @@ -449,6 +434,7 @@ class LTX2Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): LTX2PrepareAudioLatentsStep, LTX2PrepareCoordsStep, LTX2Image2VideoDenoiseStep, + LTX2UnpackLatentsStep, ] block_names = [ "input", @@ -458,6 +444,7 @@ class LTX2Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): "prepare_audio_latents", "prepare_coords", "denoise", + "unpack", ] @property @@ -482,9 +469,8 @@ class LTX2ConditionCoreDenoiseStep(SequentialPipelineBlocks): conditions to the video latents and runs the joint denoising loop. Components: - transformer (`LTX2VideoTransformer3DModel`) vae (`AutoencoderKLLTX2Video`) scheduler - (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider - (`LTX2Guidance`) + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) Inputs: num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -495,29 +481,26 @@ class LTX2ConditionCoreDenoiseStep(SequentialPipelineBlocks): Audio-branch text conditioning (cond). connector_attention_mask (`Tensor`): Binary text attention mask (cond). - negative_connector_prompt_embeds (`Tensor`): - Video-branch text conditioning (uncond). - negative_connector_audio_prompt_embeds (`Tensor`): - Audio-branch text conditioning (uncond). - negative_connector_attention_mask (`Tensor`): - Binary text attention mask (uncond). + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond), `None` without classifier-free guidance. condition_latents (`list`): - Per-condition normalized VAE latents of shape [1, C, F, H, W]. + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). condition_strengths (`list`): Per-condition conditioning strengths. condition_indices (`list`): Per-condition latent frame index at which the condition is applied. condition_pixel_frames (`list`): Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). + num_frames (`int`): + The number of frames in the generated video. frame_rate (`float`, *optional*, defaults to 24.0): Frames per second of the generated video. noise_scale (`float`, *optional*): @@ -525,22 +508,14 @@ class LTX2ConditionCoreDenoiseStep(SequentialPipelineBlocks): `sigmas` are supplied, else 1.0. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - batch_size (`int`): - The number of prompts being denoised, used to expand conditioning per prompt. generator (`Generator`, *optional*): Torch generator for deterministic generation. num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. timesteps (`Tensor`, *optional*): Timesteps for the denoising process. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. - dtype (`dtype`): - The dtype the model inputs are cast to. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. @@ -563,8 +538,19 @@ class LTX2ConditionCoreDenoiseStep(SequentialPipelineBlocks): LTX2ConditionPrepareAudioLatentsStep, LTX2ConditionPrepareCoordsStep, LTX2ConditionDenoiseStep, + LTX2TrimConditionTokensStep, + LTX2UnpackLatentsStep, + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_audio_latents", + "prepare_coords", + "denoise", + "trim_condition_tokens", + "unpack", ] - block_names = ["input", "prepare_latents", "set_timesteps", "prepare_audio_latents", "prepare_coords", "denoise"] @property def description(self): @@ -591,7 +577,7 @@ class LTX2AutoConditionEncoderStep(ConditionalPipelineBlocks): - Skipped for text-to-video and image-to-video. Components: - vae (`AutoencoderKLLTX2Video`) text_encoder (`PreTrainedModel`) + vae (`AutoencoderKLLTX2Video`) Inputs: conditions (`list`, *optional*): @@ -602,14 +588,13 @@ class LTX2AutoConditionEncoderStep(ConditionalPipelineBlocks): width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). + The number of frames in the generated video. generator (`Generator`, *optional*): Torch generator for deterministic generation. Outputs: condition_latents (`list`): - Per-condition normalized VAE latents of shape [1, C, F, H, W]. + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). condition_strengths (`list`): Per-condition conditioning strengths. condition_indices (`list`): @@ -650,7 +635,7 @@ class LTX2AutoReferenceEncoderStep(ConditionalPipelineBlocks): - Skipped otherwise, for IC-LoRAs that take no reference video. Components: - vae (`AutoencoderKLLTX2Video`) transformer (`LTX2VideoTransformer3DModel`) video_processor (`VideoProcessor`) + vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) Inputs: reference_conditions (`list`, *optional*): @@ -658,37 +643,20 @@ class LTX2AutoReferenceEncoderStep(ConditionalPipelineBlocks): adapter attends to. reference_downscale_factor (`int`, *optional*, defaults to 1): Ratio between the target and reference resolutions; 2 means the reference is preprocessed at half the - target resolution. Spatial coordinates are scaled by this factor so the reference tokens land in the - target coordinate space. Must match the factor the IC-LoRA was trained with. - conditioning_attention_strength (`float`, *optional*, defaults to 1.0): - Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 - (default) leaves attention unmasked. - conditioning_attention_mask (`Tensor`, *optional*): - Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying - attention strength. Downsampled to the reference's latent grid and multiplied by - `conditioning_attention_strength`. + target resolution. Must match the factor the IC-LoRA was trained with. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). - frame_rate (`float`, *optional*, defaults to 24.0): - Frames per second of the generated video. + The number of frames in the generated video. generator (`Generator`, *optional*): Torch generator for deterministic generation. Outputs: - reference_latents (`Tensor`): - Packed reference tokens of shape [1, total_reference_tokens, C]. - reference_coords (`Tensor`): - RoPE coordinates for the reference tokens, of shape [1, 3, total_reference_tokens, 2]. - reference_token_counts (`list`): - Per-reference token counts, in `reference_conditions` order. - reference_cross_mask (`Tensor`): - Per-reference-token noisy<->reference attention strengths of shape [1, total_reference_tokens], or `None` - when attention is left unmasked. + reference_latents (`list`): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed), in `reference_conditions` + order. """ model_name = "ltx2" @@ -711,69 +679,6 @@ def description(self): ) -# auto_docstring -class LTX2AutoBuildVideoSelfAttentionMaskStep(ConditionalPipelineBlocks): - """ - Conditional video self-attention mask step, run only when the reference tokens carry a per-token attention - strength. - - `LTX2BuildVideoSelfAttentionMaskStep` when `conditioning_attention_strength < 1.0` or a - `conditioning_attention_mask` is supplied. - - Skipped otherwise, leaving attention unmasked. - - Inputs: - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. - base_token_count (`int`, *optional*): - Number of generated-video tokens, i.e. the sequence length before appended tokens. - num_ref_tokens (`int`, *optional*): - Number of reference tokens, which sit at the very end of the sequence. - reference_cross_mask (`Tensor`, *optional*): - Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. - reference_token_counts (`list`, *optional*): - Per-reference token counts, used to split `reference_cross_mask` into attention groups. - - Outputs: - video_self_attention_mask (`Tensor`): - Multiplicative self-attention mask of shape [B, S, S] with values in [0, 1]. - """ - - model_name = "ltx2" - block_classes = [LTX2BuildVideoSelfAttentionMaskStep] - block_names = ["attention_mask"] - block_trigger_inputs = [ - "reference_conditions", - "conditioning_attention_strength", - "conditioning_attention_mask", - ] - - def select_block( - self, reference_conditions=None, conditioning_attention_strength=None, conditioning_attention_mask=None - ) -> str | None: - # The mask only ever governs noisy<->reference attention, so it is meaningless without reference tokens. - # Beyond that, an all-ones mask at full strength is the same as no mask, so only build one when it can - # actually bite -- the same `mask_needed` test `LTX2ReferenceEncoderStep` uses to decide whether to emit - # `reference_cross_mask` at all. `conditioning_attention_strength` defaults to `None` rather than its - # input default of 1.0 because `get_execution_blocks` passes every unset trigger input as an explicit - # `None`; `None` and 1.0 both mean "leave attention unmasked". - if not reference_conditions: - return None - if conditioning_attention_mask is not None: - return "attention_mask" - if conditioning_attention_strength is not None and conditioning_attention_strength < 1.0: - return "attention_mask" - return None - - @property - def description(self): - return ( - "Conditional video self-attention mask step, run only when the reference tokens carry a per-token " - "attention strength.\n" - " - `LTX2BuildVideoSelfAttentionMaskStep` when `conditioning_attention_strength < 1.0` or a " - "`conditioning_attention_mask` is supplied.\n" - " - Skipped otherwise, leaving attention unmasked." - ) - - # auto_docstring class LTX2InContextCoreDenoiseStep(SequentialPipelineBlocks): """ @@ -783,9 +688,8 @@ class LTX2InContextCoreDenoiseStep(SequentialPipelineBlocks): the reference implementation's uniform treatment of both. Components: - transformer (`LTX2VideoTransformer3DModel`) vae (`AutoencoderKLLTX2Video`) scheduler - (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider - (`LTX2Guidance`) + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) Inputs: num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -796,14 +700,14 @@ class LTX2InContextCoreDenoiseStep(SequentialPipelineBlocks): Audio-branch text conditioning (cond). connector_attention_mask (`Tensor`): Binary text attention mask (cond). - negative_connector_prompt_embeds (`Tensor`): - Video-branch text conditioning (uncond). - negative_connector_audio_prompt_embeds (`Tensor`): - Audio-branch text conditioning (uncond). - negative_connector_attention_mask (`Tensor`): - Binary text attention mask (uncond). + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond), `None` without classifier-free guidance. condition_latents (`list`): - Per-condition normalized VAE latents of shape [1, C, F, H, W]. + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). condition_strengths (`list`): Per-condition conditioning strengths. condition_indices (`list`): @@ -813,22 +717,20 @@ class LTX2InContextCoreDenoiseStep(SequentialPipelineBlocks): reference_conditions (`list`, *optional*): `LTX2ReferenceCondition` (or list of them); only their `strength` is read here. Omit for IC-LoRAs that carry their behavior in the adapter weights and take no reference video. - reference_latents (`Tensor`, *optional*): - Packed reference tokens of shape [1, total_reference_tokens, C], or `None` when no reference conditions - were supplied (`LTX2AutoReferenceEncoderStep` is skipped). - reference_coords (`Tensor`, *optional*): - RoPE coordinates for the reference tokens. - reference_token_counts (`list`, *optional*): - Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). + reference_downscale_factor (`int`, *optional*, defaults to 1): + Ratio between the target and reference resolutions. The reference tokens' spatial coordinates are scaled + by it so they land in the target coordinate space, preserving the positional relationship the IC-LoRA was + trained on. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). + num_frames (`int`): + The number of frames in the generated video. frame_rate (`float`, *optional*, defaults to 24.0): Frames per second of the generated video. noise_scale (`float`, *optional*): @@ -836,24 +738,21 @@ class LTX2InContextCoreDenoiseStep(SequentialPipelineBlocks): `sigmas` are supplied, else 1.0. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - batch_size (`int`): - The number of prompts being denoised, used to expand conditioning per prompt. generator (`Generator`, *optional*): Torch generator for deterministic generation. - reference_cross_mask (`Tensor`, *optional*): - Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. timesteps (`Tensor`, *optional*): Timesteps for the denoising process. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. - dtype (`dtype`): - The dtype the model inputs are cast to. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. @@ -870,20 +769,22 @@ class LTX2InContextCoreDenoiseStep(SequentialPipelineBlocks): block_classes = [ LTX2TextInputStep, LTX2InContextPrepareLatentsStep, - LTX2AutoBuildVideoSelfAttentionMaskStep, LTX2ConditionSetTimestepsStep, LTX2ConditionPrepareAudioLatentsStep, LTX2ConditionPrepareCoordsStep, LTX2ConditionDenoiseStep, + LTX2TrimConditionTokensStep, + LTX2UnpackLatentsStep, ] block_names = [ "input", "prepare_latents", - "attention_mask", "set_timesteps", "prepare_audio_latents", "prepare_coords", "denoise", + "trim_condition_tokens", + "unpack", ] @property @@ -914,9 +815,8 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): - `LTX2CoreDenoiseStep` otherwise (text-to-video). Components: - transformer (`LTX2VideoTransformer3DModel`) vae (`AutoencoderKLLTX2Video`) scheduler - (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider - (`LTX2Guidance`) + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) Inputs: num_videos_per_prompt (`int`, *optional*, defaults to 1): @@ -927,14 +827,14 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): Audio-branch text conditioning (cond). connector_attention_mask (`Tensor`): Binary text attention mask (cond). - negative_connector_prompt_embeds (`Tensor`): - Video-branch text conditioning (uncond). - negative_connector_audio_prompt_embeds (`Tensor`): - Audio-branch text conditioning (uncond). - negative_connector_attention_mask (`Tensor`): - Binary text attention mask (uncond). + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond), `None` without classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond), `None` without classifier-free guidance. condition_latents (`list`, *optional*): - Per-condition normalized VAE latents of shape [1, C, F, H, W]. + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). condition_strengths (`list`, *optional*): Per-condition conditioning strengths. condition_indices (`list`, *optional*): @@ -944,22 +844,20 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): reference_conditions (`list`, *optional*): `LTX2ReferenceCondition` (or list of them); only their `strength` is read here. Omit for IC-LoRAs that carry their behavior in the adapter weights and take no reference video. - reference_latents (`Tensor`, *optional*): - Packed reference tokens of shape [1, total_reference_tokens, C], or `None` when no reference conditions - were supplied (`LTX2AutoReferenceEncoderStep` is skipped). - reference_coords (`Tensor`, *optional*): - RoPE coordinates for the reference tokens. - reference_token_counts (`list`, *optional*): - Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`): - Pre-generated noisy latents for image generation. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). + reference_downscale_factor (`int`, *optional*, defaults to 1): + Ratio between the target and reference resolutions. The reference tokens' spatial coordinates are scaled + by it so they land in the target coordinate space, preserving the positional relationship the IC-LoRA was + trained on. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). + num_frames (`int`): + The number of frames in the generated video. frame_rate (`float`, *optional*, defaults to 24.0): Frames per second of the generated video. noise_scale (`float`, *optional*): @@ -967,24 +865,21 @@ class LTX2AutoCoreDenoiseStep(ConditionalPipelineBlocks): `sigmas` are supplied, else 1.0. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - batch_size (`int`): - The number of prompts being denoised, used to expand conditioning per prompt. generator (`Generator`, *optional*): Torch generator for deterministic generation. - reference_cross_mask (`Tensor`, *optional*): - Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. num_inference_steps (`int`): The number of denoising steps. timesteps (`Tensor`): Timesteps for the denoising process. - audio_latents (`Tensor`): - Optional pre-encoded audio latents; random noise is used when not provided. - dtype (`dtype`): - The dtype the model inputs are cast to. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. image_latents (`Tensor`, *optional*): @@ -1035,7 +930,7 @@ def description(self): # auto_docstring class LTX2DecoderStep(SequentialPipelineBlocks): """ - Decode stage: VAE-decodes the video latents and vocodes the audio latents (or returns latents). + Decode stage: VAE-decodes the video latents and vocodes the audio latents. Components: vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) @@ -1043,16 +938,9 @@ class LTX2DecoderStep(SequentialPipelineBlocks): Inputs: latents (`Tensor`): - Pre-generated noisy latents for image generation. + Video latents of shape [B, C, F, H, W] (normalized, not packed). output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 704): - The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). decode_timestep (`None`, *optional*, defaults to 0.0): The timestep at which the VAE decodes the final latents. decode_noise_scale (`None`, *optional*): @@ -1065,9 +953,7 @@ class LTX2DecoderStep(SequentialPipelineBlocks): dtype (`dtype`): The dtype of the model inputs, can be generated in input step. audio_latents (`Tensor`): - Denoised audio latents. - audio_num_frames (`int`): - Number of audio latent frames, used to unpack the audio latent sequence. + Audio latents of shape [B, C, L, M] (normalized, not packed). Outputs: videos (`list`): @@ -1082,146 +968,7 @@ class LTX2DecoderStep(SequentialPipelineBlocks): @property def description(self): - return "Decode stage: VAE-decodes the video latents and vocodes the audio latents (or returns latents)." - - @property - def outputs(self): - return [ - OutputParam.template("videos"), - OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform."), - ] - - -# auto_docstring -class LTX2ConditionDecoderStep(SequentialPipelineBlocks): - """ - Decode stage for condition workflows: drops the appended keyframe-condition tokens, then VAE-decodes the video - latents and vocodes the audio latents (or returns latents). - - Components: - vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) - vocoder (`LTX2Vocoder`) - - Inputs: - latents (`Tensor`): - Pre-generated noisy latents for image generation. - base_token_count (`int`): - Number of generated-video tokens, i.e. the sequence length before appended tokens. - output_type (`str`, *optional*, defaults to pil): - Output format: 'pil', 'np', 'pt'. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 704): - The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). - decode_timestep (`None`, *optional*, defaults to 0.0): - The timestep at which the VAE decodes the final latents. - decode_noise_scale (`None`, *optional*): - Noise interpolation factor applied to the latents at the decode timestep. - generator (`Generator`, *optional*): - Torch generator for deterministic generation. - batch_size (`int`, *optional*, defaults to 1): - Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can - be generated in input step. - dtype (`dtype`): - The dtype of the model inputs, can be generated in input step. - audio_latents (`Tensor`): - Denoised audio latents. - audio_num_frames (`int`): - Number of audio latent frames, used to unpack the audio latent sequence. - - Outputs: - videos (`list`): - The generated videos. - audio (`Tensor`): - The generated audio waveform. - """ - - model_name = "ltx2" - block_classes = [LTX2TrimConditionTokensStep, LTX2VaeDecoderStep, LTX2AudioDecoderStep] - block_names = ["trim_condition_tokens", "video_decode", "audio_decode"] - - @property - def description(self): - return ( - "Decode stage for condition workflows: drops the appended keyframe-condition tokens, then VAE-decodes " - "the video latents and vocodes the audio latents (or returns latents)." - ) - - @property - def outputs(self): - return [ - OutputParam.template("videos"), - OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform."), - ] - - -# auto_docstring -class LTX2AutoDecoderStep(AutoPipelineBlocks): - """ - Auto decode block that selects the decoder based on inputs. - - `LTX2ConditionDecoderStep` when `base_token_count` is present, i.e. the denoised sequence carries appended - keyframe / reference tokens (condition, in-context). - - `LTX2DecoderStep` otherwise (text-to-video, image-to-video). - - Components: - vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) - vocoder (`LTX2Vocoder`) - - Inputs: - latents (`Tensor`): - Pre-generated noisy latents for image generation. - base_token_count (`int`, *optional*): - Number of generated-video tokens, i.e. the sequence length before appended tokens. - output_type (`str`, *optional*, defaults to pil): - Output format: 'pil', 'np', 'pt'. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 704): - The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). - decode_timestep (`None`, *optional*, defaults to 0.0): - The timestep at which the VAE decodes the final latents. - decode_noise_scale (`None`, *optional*): - Noise interpolation factor applied to the latents at the decode timestep. - generator (`Generator`, *optional*): - Torch generator for deterministic generation. - batch_size (`int`, *optional*, defaults to 1): - Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can - be generated in input step. - dtype (`dtype`): - The dtype of the model inputs, can be generated in input step. - audio_latents (`Tensor`): - Denoised audio latents. - audio_num_frames (`int`): - Number of audio latent frames, used to unpack the audio latent sequence. - - Outputs: - videos (`list`): - The generated videos. - audio (`Tensor`): - The generated audio waveform. - """ - - model_name = "ltx2" - # `base_token_count` is emitted only by the condition / in-context prepare-latents steps, so it is exactly the - # signal for "the sequence carries appended tokens that have to come off before decoding". - block_classes = [LTX2ConditionDecoderStep, LTX2DecoderStep] - block_names = ["condition", "default"] - block_trigger_inputs = ["base_token_count", None] - - @property - def description(self): - return ( - "Auto decode block that selects the decoder based on inputs.\n" - " - `LTX2ConditionDecoderStep` when `base_token_count` is present, i.e. the denoised sequence carries " - "appended keyframe / reference tokens (condition, in-context).\n" - " - `LTX2DecoderStep` otherwise (text-to-video, image-to-video)." - ) + return "Decode stage: VAE-decodes the video latents and vocodes the audio latents." @property def outputs(self): @@ -1239,9 +986,9 @@ class LTX2Blocks(SequentialPipelineBlocks): Components: prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) scheduler - (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) audio_vae - (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) vae - (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) vocoder (`LTX2Vocoder`) + (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) guider (`LTX2Guidance`) + audio_guider (`LTX2Guidance`) vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) audio_vae + (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) Inputs: prompt (`str`, *optional*): @@ -1249,8 +996,6 @@ class LTX2Blocks(SequentialPipelineBlocks): conditions (`list`, *optional*): `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the generated video. - enable_prompt_enhancement (`bool`, *optional*, defaults to False): - Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. system_prompt (`str`, *optional*): System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. @@ -1265,10 +1010,15 @@ class LTX2Blocks(SequentialPipelineBlocks): Torch generator for deterministic generation. image (`Image | list`, *optional*): Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. negative_prompt (`str`, *optional*): The prompt or prompts not to guide the image generation. max_sequence_length (`int`, *optional*, defaults to 1024): Maximum sequence length for prompt encoding. + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. min_seconds (`float`, *optional*, defaults to 1.0): Lower bound on the auto-predicted duration. max_seconds (`float`, *optional*, defaults to 20.0): @@ -1287,20 +1037,8 @@ class LTX2Blocks(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. - noise_scale (`float`, *optional*): - Interpolation factor between random noise and any provided latents. `None` (default) resolves to 0.0, - which keeps the provided latents. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. output_type (`str`, *optional*, defaults to pil): @@ -1348,8 +1086,8 @@ class LTX2ImageToVideoBlocks(SequentialPipelineBlocks): prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) scheduler (`FlowMatchEulerDiscreteScheduler`) - transformer (`LTX2VideoTransformer3DModel`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) - audio_guider (`LTX2Guidance`) vocoder (`LTX2Vocoder`) + transformer (`LTX2VideoTransformer3DModel`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) audio_vae + (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) Inputs: prompt (`str`, *optional*): @@ -1357,8 +1095,6 @@ class LTX2ImageToVideoBlocks(SequentialPipelineBlocks): conditions (`list`, *optional*): `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the generated video. - enable_prompt_enhancement (`bool`, *optional*, defaults to False): - Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. system_prompt (`str`, *optional*): System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. @@ -1373,10 +1109,15 @@ class LTX2ImageToVideoBlocks(SequentialPipelineBlocks): Torch generator for deterministic generation. image (`Image | list`, *optional*): Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. negative_prompt (`str`, *optional*): The prompt or prompts not to guide the image generation. max_sequence_length (`int`, *optional*, defaults to 1024): Maximum sequence length for prompt encoding. + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. min_seconds (`float`, *optional*, defaults to 1.0): Lower bound on the auto-predicted duration. max_seconds (`float`, *optional*, defaults to 20.0): @@ -1389,9 +1130,8 @@ class LTX2ImageToVideoBlocks(SequentialPipelineBlocks): The width in pixels of the generated image. image_crf (`int`, *optional*): H.264 CRF used to re-compress the conditioning `image` before VAE encode, matching the compression the - model was trained against. `None` (default) resolves from the text-encoder generation (33 through - LTX-2.3, 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when - re-compression runs. + model was trained against. `None` (default) uses the pipeline's `default_image_crf` (33 through LTX-2.3, + 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when re-compression runs. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. num_inference_steps (`int`, *optional*, defaults to 30): @@ -1400,22 +1140,10 @@ class LTX2ImageToVideoBlocks(SequentialPipelineBlocks): Timesteps for the denoising process. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. - noise_scale (`float`, *optional*): - Interpolation factor between random noise and any provided latents. `None` (default) resolves to 0.0, - which keeps the provided latents. image_latents (`Tensor`): VAE-encoded reference-image latents used for image-to-video conditioning. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. output_type (`str`, *optional*, defaults to pil): @@ -1465,8 +1193,8 @@ class LTX2ConditionBlocks(SequentialPipelineBlocks): prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) vae (`AutoencoderKLLTX2Video`) transformer (`LTX2VideoTransformer3DModel`) scheduler - (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider - (`LTX2Guidance`) video_processor (`VideoProcessor`) vocoder (`LTX2Vocoder`) + (`FlowMatchEulerDiscreteScheduler`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) video_processor + (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) Inputs: prompt (`str`, *optional*): @@ -1474,8 +1202,6 @@ class LTX2ConditionBlocks(SequentialPipelineBlocks): conditions (`list`, *optional*): `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the generated video. - enable_prompt_enhancement (`bool`, *optional*, defaults to False): - Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. system_prompt (`str`, *optional*): System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. @@ -1490,10 +1216,15 @@ class LTX2ConditionBlocks(SequentialPipelineBlocks): Torch generator for deterministic generation. image (`Image | list`, *optional*): Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. negative_prompt (`str`, *optional*): The prompt or prompts not to guide the image generation. max_sequence_length (`int`, *optional*, defaults to 1024): Maximum sequence length for prompt encoding. + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. min_seconds (`float`, *optional*, defaults to 1.0): Lower bound on the auto-predicted duration. max_seconds (`float`, *optional*, defaults to 20.0): @@ -1504,13 +1235,8 @@ class LTX2ConditionBlocks(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. noise_scale (`float`, *optional*): Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom `sigmas` are supplied, else 1.0. @@ -1520,12 +1246,8 @@ class LTX2ConditionBlocks(SequentialPipelineBlocks): The number of denoising steps. timesteps (`Tensor`, *optional*): Timesteps for the denoising process. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. output_type (`str`, *optional*, defaults to pil): @@ -1549,7 +1271,7 @@ class LTX2ConditionBlocks(SequentialPipelineBlocks): LTX2AutoDurationStep, LTX2ConditionEncoderStep, LTX2ConditionCoreDenoiseStep, - LTX2ConditionDecoderStep, + LTX2DecoderStep, ] block_names = ["prompt_enhancer", "text_encoder", "duration", "condition_encoder", "denoise", "decode"] @@ -1579,10 +1301,10 @@ class LTX2InContextBlocks(SequentialPipelineBlocks): Components: prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer - (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) vae (`AutoencoderKLLTX2Video`) transformer - (`LTX2VideoTransformer3DModel`) video_processor (`VideoProcessor`) scheduler - (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) audio_guider - (`LTX2Guidance`) vocoder (`LTX2Vocoder`) + (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) vae (`AutoencoderKLLTX2Video`) video_processor + (`VideoProcessor`) transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) audio_vae (`AutoencoderKLLTX2Audio`) vocoder + (`LTX2Vocoder`) Inputs: prompt (`str`, *optional*): @@ -1590,8 +1312,6 @@ class LTX2InContextBlocks(SequentialPipelineBlocks): conditions (`list`, *optional*): `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the generated video. - enable_prompt_enhancement (`bool`, *optional*, defaults to False): - Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. system_prompt (`str`, *optional*): System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. @@ -1606,6 +1326,8 @@ class LTX2InContextBlocks(SequentialPipelineBlocks): Torch generator for deterministic generation. image (`Image | list`, *optional*): Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. negative_prompt (`str`, *optional*): The prompt or prompts not to guide the image generation. max_sequence_length (`int`, *optional*, defaults to 1024): @@ -1614,53 +1336,40 @@ class LTX2InContextBlocks(SequentialPipelineBlocks): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). + num_frames (`int`): + The number of frames in the generated video. reference_conditions (`list`, *optional*): `LTX2ReferenceCondition` (or list of them) whose videos are encoded into extra latent tokens the IC-LoRA adapter attends to. reference_downscale_factor (`int`, *optional*, defaults to 1): Ratio between the target and reference resolutions; 2 means the reference is preprocessed at half the - target resolution. Spatial coordinates are scaled by this factor so the reference tokens land in the - target coordinate space. Must match the factor the IC-LoRA was trained with. - conditioning_attention_strength (`float`, *optional*, defaults to 1.0): - Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 - (default) leaves attention unmasked. - conditioning_attention_mask (`Tensor`, *optional*): - Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying - attention strength. Downsampled to the reference's latent grid and multiplied by - `conditioning_attention_strength`. - frame_rate (`float`, *optional*, defaults to 24.0): - Frames per second of the generated video. + target resolution. Must match the factor the IC-LoRA was trained with. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. - reference_latents (`Tensor`, *optional*): - Packed reference tokens of shape [1, total_reference_tokens, C], or `None` when no reference conditions - were supplied (`LTX2AutoReferenceEncoderStep` is skipped). - reference_coords (`Tensor`, *optional*): - RoPE coordinates for the reference tokens. - reference_token_counts (`list`, *optional*): - Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`, *optional*): - Pre-generated noisy latents for image generation. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. noise_scale (`float`, *optional*): Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom `sigmas` are supplied, else 1.0. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - reference_cross_mask (`Tensor`, *optional*): - Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. num_inference_steps (`int`, *optional*, defaults to 30): The number of denoising steps. timesteps (`Tensor`, *optional*): Timesteps for the denoising process. - audio_latents (`Tensor`, *optional*): - Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. output_type (`str`, *optional*, defaults to pil): @@ -1688,7 +1397,7 @@ class LTX2InContextBlocks(SequentialPipelineBlocks): LTX2ConditionEncoderStep, LTX2AutoReferenceEncoderStep, LTX2InContextCoreDenoiseStep, - LTX2ConditionDecoderStep, + LTX2DecoderStep, ] block_names = [ "prompt_enhancer", @@ -1733,8 +1442,8 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) transformer (`LTX2VideoTransformer3DModel`) - scheduler (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) - audio_guider (`LTX2Guidance`) vocoder (`LTX2Vocoder`) + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) audio_vae + (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) Inputs: prompt (`str`, *optional*): @@ -1742,8 +1451,6 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): conditions (`list`, *optional*): `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the generated video. - enable_prompt_enhancement (`bool`, *optional*, defaults to False): - Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. system_prompt (`str`, *optional*): System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. @@ -1758,10 +1465,15 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): Torch generator for deterministic generation. image (`Image | list`, *optional*): Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. negative_prompt (`str`, *optional*): The prompt or prompts not to guide the image generation. max_sequence_length (`int`, *optional*, defaults to 1024): Maximum sequence length for prompt encoding. + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. min_seconds (`float`, *optional*, defaults to 1.0): Lower bound on the auto-predicted duration. max_seconds (`float`, *optional*, defaults to 20.0): @@ -1774,62 +1486,46 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): The width in pixels of the generated image. image_crf (`int`, *optional*): H.264 CRF used to re-compress the conditioning `image` before VAE encode, matching the compression the - model was trained against. `None` (default) resolves from the text-encoder generation (33 through - LTX-2.3, 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when - re-compression runs. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). + model was trained against. `None` (default) uses the pipeline's `default_image_crf` (33 through LTX-2.3, + 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when re-compression runs. reference_conditions (`list`, *optional*): `LTX2ReferenceCondition` (or list of them) whose videos are encoded into extra latent tokens the IC-LoRA adapter attends to. reference_downscale_factor (`int`, *optional*, defaults to 1): Ratio between the target and reference resolutions; 2 means the reference is preprocessed at half the - target resolution. Spatial coordinates are scaled by this factor so the reference tokens land in the - target coordinate space. Must match the factor the IC-LoRA was trained with. - conditioning_attention_strength (`float`, *optional*, defaults to 1.0): - Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 - (default) leaves attention unmasked. - conditioning_attention_mask (`Tensor`, *optional*): - Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying - attention strength. Downsampled to the reference's latent grid and multiplied by - `conditioning_attention_strength`. + target resolution. Must match the factor the IC-LoRA was trained with. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. condition_latents (`list`, *optional*): - Per-condition normalized VAE latents of shape [1, C, F, H, W]. + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). condition_strengths (`list`, *optional*): Per-condition conditioning strengths. condition_indices (`list`, *optional*): Per-condition latent frame index at which the condition is applied. condition_pixel_frames (`list`, *optional*): Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. - reference_latents (`Tensor`, *optional*): - Packed reference tokens of shape [1, total_reference_tokens, C], or `None` when no reference conditions - were supplied (`LTX2AutoReferenceEncoderStep` is skipped). - reference_coords (`Tensor`, *optional*): - RoPE coordinates for the reference tokens. - reference_token_counts (`list`, *optional*): - Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`): - Pre-generated noisy latents for image generation. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). noise_scale (`float`, *optional*): Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom `sigmas` are supplied, else 1.0. sigmas (`list`, *optional*): Custom sigmas for the denoising process. - reference_cross_mask (`Tensor`, *optional*): - Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. num_inference_steps (`int`): The number of denoising steps. timesteps (`Tensor`): Timesteps for the denoising process. - audio_latents (`Tensor`): - Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. image_latents (`Tensor`, *optional*): @@ -1857,7 +1553,7 @@ class LTX2AutoBlocks(SequentialPipelineBlocks): LTX2AutoConditionEncoderStep, LTX2AutoReferenceEncoderStep, LTX2AutoCoreDenoiseStep, - LTX2AutoDecoderStep, + LTX2DecoderStep, ] block_names = [ "prompt_enhancer", diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py index 7c77aba94a74..53d4b3137123 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py @@ -14,29 +14,1354 @@ import torch -from ..modular_pipeline import AutoPipelineBlocks, SequentialPipelineBlocks -from ..modular_pipeline_utils import OutputParam +from ...pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES +from ...utils import logging +from ..modular_pipeline import AutoPipelineBlocks, ConditionalPipelineBlocks, SequentialPipelineBlocks +from ..modular_pipeline_utils import InputParam, InsertableDict, OutputParam +from .before_denoise import ( + LTX2ConditionPrepareAudioLatentsStep, + LTX2ConditionPrepareCoordsStep, + LTX2ConditionPrepareLatentsStep, + LTX2ConditionSetTimestepsStep, + LTX2ConditionStage2PrepareLatentsStep, + LTX2Image2VideoPrepareLatentsStep, + LTX2InContextPrepareLatentsStep, + LTX2LatentUpsampleStep, + LTX2PrepareAudioLatentsStep, + LTX2PrepareCoordsStep, + LTX2PrepareLatentsStep, + LTX2SetTimestepsStep, + LTX2Stage2PrepareAudioLatentsStep, + LTX2Stage2PrepareLatentsStep, + LTX2TextInputStep, +) from .decoders import ( LTX2AudioDecoderStep, LTX2DiffusionVaeDecoderStep, LTX2TrimConditionTokensStep, + LTX2UnpackLatentsStep, +) +from .denoise import ( + LTX2ConditionDenoiseStep, + LTX2DenoiseStep, + LTX2Image2VideoDenoiseStep, +) +from .encoders import ( + LTX2ConditionEncoderStep, + LTX2ConditionPromptEnhancerStep, + LTX2DurationStep, + LTX2ImageToVideoPromptEnhancerStep, + LTX2PromptEnhancerStep, + LTX2ReferenceEncoderStep, + LTX2TextConnectorStep, + LTX2TextEncoderStep, + LTX2VaeEncoderStep, +) + + +logger = logging.get_logger(__name__) + + +# auto_docstring +class LTX25AutoPromptEnhancerStep(ConditionalPipelineBlocks): + """ + Conditional prompt-enhancer step, run only when `enable_prompt_enhancement` is truthy. + - `LTX2ConditionPromptEnhancerStep` when `conditions` are provided (condition-to-video, in-context); grounds the + rewrite in the first `PIL.Image.Image` frame found in `conditions`, falling back to a text-only rewrite when + there is none. + - `LTX2ImageToVideoPromptEnhancerStep` when an `image` is provided (image-to-video). + - `LTX2PromptEnhancerStep` otherwise (text-to-video). + - Skipped when `enable_prompt_enhancement` is falsy / not provided. + + Components: + prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + conditions (`list`, *optional*): + `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the + generated video. + system_prompt (`str`, *optional*): + System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` + condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. + prompt_max_new_tokens (`int`, *optional*): + Maximum number of new tokens to generate during prompt enhancement. Defaults to 600, the LTX-2.5 Gemma-4 + enhancer's budget. + prompt_enhancement_kwargs (`dict`, *optional*): + Keyword arguments for the enhancer's `.generate` call. Defaults to greedy decoding. + prompt_enhancement_seed (`int`, *optional*, defaults to 10): + Random seed for prompt enhancement (inert under LTX-2.5's greedy decoding). + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. + + Outputs: + prompt (`list`): + The prompt(s) after prompt-enhancer rewriting. + """ + + model_name = "ltx2.5" + block_classes = [LTX2ConditionPromptEnhancerStep, LTX2ImageToVideoPromptEnhancerStep, LTX2PromptEnhancerStep] + block_names = ["condition", "image2video", "text2video"] + block_trigger_inputs = ["conditions", "image", "enable_prompt_enhancement"] + + @property + def inputs(self): + # The trigger belongs to this wrapper, not to the enhancer steps: each of them always enhances. + inputs = super().inputs + inputs.append( + InputParam( + "enable_prompt_enhancement", + type_hint=bool, + default=False, + description="Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines.", + ) + ) + return inputs + + def select_block(self, conditions=None, image=None, enable_prompt_enhancement=False) -> str | None: + # `conditions` is checked before `image`: the condition and in-context workflows place their reference + # frames in `conditions` and never take a raw `image`. + if not enable_prompt_enhancement: + return None + if conditions is not None: + return "condition" + return "image2video" if image is not None else "text2video" + + @property + def description(self): + return ( + "Conditional prompt-enhancer step, run only when `enable_prompt_enhancement` is truthy.\n" + " - `LTX2ConditionPromptEnhancerStep` when `conditions` are provided (condition-to-video, in-context); " + "grounds the rewrite in the first `PIL.Image.Image` frame found in `conditions`, falling back to a " + "text-only rewrite when there is none.\n" + " - `LTX2ImageToVideoPromptEnhancerStep` when an `image` is provided (image-to-video).\n" + " - `LTX2PromptEnhancerStep` otherwise (text-to-video).\n" + " - Skipped when `enable_prompt_enhancement` is falsy / not provided." + ) + + +# auto_docstring +class LTX25TextConditioningStep(SequentialPipelineBlocks): + """ + Text-conditioning stage for LTX-2.X: encodes the prompt(s), then runs the text connectors to produce the + video/audio-branch connector embeddings the denoiser consumes. Outputs stay at one row per prompt -- the denoise + stage expands them by `num_videos_per_prompt` -- so they can be reused across denoise runs. + + Components: + text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) + + Inputs: + prompt (`str`): + The prompt or prompts to guide image generation. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + + Outputs: + prompt_embeds (`Tensor`): + Packed per-layer Gemma hidden states for the prompt. + prompt_attention_mask (`Tensor`): + Binary attention mask for `prompt_embeds`. + negative_prompt_embeds (`Tensor`): + Packed per-layer Gemma hidden states for the negative prompt, `None` when not encoded. + negative_prompt_attention_mask (`Tensor`): + Binary attention mask for `negative_prompt_embeds`, `None` when not encoded. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond). + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond). + connector_attention_mask (`Tensor`): + Binary text attention mask (cond). + negative_connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (uncond), `None` when no negative prompt was encoded. + negative_connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (uncond), `None` when no negative prompt was encoded. + negative_connector_attention_mask (`Tensor`): + Binary text attention mask (uncond), `None` when no negative prompt was encoded. + """ + + model_name = "ltx2.5" + block_classes = [LTX2TextEncoderStep, LTX2TextConnectorStep] + block_names = ["text_encoder", "connectors"] + + @property + def description(self): + return ( + "Text-conditioning stage for LTX-2.X: encodes the prompt(s), then runs the text connectors to produce " + "the video/audio-branch connector embeddings the denoiser consumes. Outputs stay at one row per prompt " + "-- the denoise stage expands them by `num_videos_per_prompt` -- so they can be reused across denoise " + "runs." + ) + + +# auto_docstring +class LTX25AutoDurationStep(ConditionalPipelineBlocks): + """ + Conditional duration-prediction step, run only when `num_frames` is omitted. + - `LTX2DurationStep` predicts `num_frames` from the connector text conditioning via the `duration_head`. + - Skipped when `num_frames` is supplied as an integer. + + Components: + duration_head (`LTX2DurationHead`) + + Inputs: + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. + min_seconds (`float`, *optional*, defaults to 1.0): + Lower bound on the auto-predicted duration. + max_seconds (`float`, *optional*, defaults to 20.0): + Upper bound on the auto-predicted duration. Must be strictly greater than `min_seconds`. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning from the connector (positive prompt). + connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning from the connector (positive prompt). + + Outputs: + num_frames (`int`): + The predicted number of frames to generate. + """ + + model_name = "ltx2.5" + block_classes = [LTX2DurationStep] + block_names = ["duration"] + block_trigger_inputs = ["num_frames"] + + def select_block(self, num_frames=None) -> str | None: + return "duration" if num_frames is None else None + + @property + def description(self): + return ( + "Conditional duration-prediction step, run only when `num_frames` is omitted.\n" + " - `LTX2DurationStep` predicts `num_frames` from the connector text conditioning via the `duration_head`.\n" + " - Skipped when `num_frames` is supplied as an integer." + ) + + +# auto_docstring +class LTX25AutoVaeEncoderStep(AutoPipelineBlocks): + """ + VAE encoder step that encodes the reference `image` into latents for image-to-video. + - `LTX2VaeEncoderStep` runs when `image` is provided. + - Skipped otherwise. + + Components: + vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) + + Inputs: + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + image_crf (`int`, *optional*): + H.264 CRF used to re-compress the conditioning `image` before VAE encode, matching the compression the + model was trained against. `None` (default) uses the pipeline's `default_image_crf` (33 through LTX-2.3, + 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when re-compression runs. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + image_latents (`Tensor`): + Image latents for image-to-video conditioning: a single latent frame of shape [B, C, 1, H, W] + (normalized, not packed). + """ + + model_name = "ltx2.5" + block_classes = [LTX2VaeEncoderStep] + block_names = ["vae_encoder"] + block_trigger_inputs = ["image"] + + @property + def description(self): + return ( + "VAE encoder step that encodes the reference `image` into latents for image-to-video.\n" + " - `LTX2VaeEncoderStep` runs when `image` is provided.\n" + " - Skipped otherwise." + ) + + +# auto_docstring +class LTX25AutoConditionEncoderStep(ConditionalPipelineBlocks): + """ + Conditional condition-encoder step, run only for the condition and in-context workflows. + - `LTX2ConditionEncoderStep` VAE-encodes the frame `conditions` into per-condition latents. + - Also runs when only `reference_conditions` are supplied, emitting empty per-condition lists for + `LTX2InContextPrepareLatentsStep`. + - Skipped for text-to-video and image-to-video. + + Components: + vae (`AutoencoderKLLTX2Video`) + + Inputs: + conditions (`list`, *optional*): + `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the + generated video. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`, *optional*): + The number of frames in the generated video. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + condition_latents (`list`): + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). + condition_strengths (`list`): + Per-condition conditioning strengths. + condition_indices (`list`): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`): + Per-condition trimmed pixel frame count, used to clamp the temporal extent of single-frame keyframe + coordinates. + """ + + model_name = "ltx2.5" + block_classes = [LTX2ConditionEncoderStep] + block_names = ["condition_encoder"] + block_trigger_inputs = ["conditions", "reference_conditions"] + + def select_block(self, conditions=None, reference_conditions=None) -> str | None: + # Also runs for a reference-only in-context request: `LTX2InContextPrepareLatentsStep` requires the + # per-condition lists, and this step emits them empty when there are no frame `conditions`. + if conditions is not None or reference_conditions: + return "condition_encoder" + return None + + @property + def description(self): + return ( + "Conditional condition-encoder step, run only for the condition and in-context workflows.\n" + " - `LTX2ConditionEncoderStep` VAE-encodes the frame `conditions` into per-condition latents.\n" + " - Also runs when only `reference_conditions` are supplied, emitting empty per-condition lists for " + "`LTX2InContextPrepareLatentsStep`.\n" + " - Skipped for text-to-video and image-to-video." + ) + + +# auto_docstring +class LTX25AutoReferenceEncoderStep(ConditionalPipelineBlocks): + """ + Conditional reference-encoder step, run only when `reference_conditions` are supplied. + - `LTX2ReferenceEncoderStep` encodes the reference videos into extra latent tokens. + - Skipped otherwise, for IC-LoRAs that take no reference video. + + Components: + vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) + + Inputs: + reference_conditions (`list`, *optional*): + `LTX2ReferenceCondition` (or list of them) whose videos are encoded into extra latent tokens the IC-LoRA + adapter attends to. + reference_downscale_factor (`int`, *optional*, defaults to 1): + Ratio between the target and reference resolutions; 2 means the reference is preprocessed at half the + target resolution. Must match the factor the IC-LoRA was trained with. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`, *optional*): + The number of frames in the generated video. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + reference_latents (`list`): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed), in `reference_conditions` + order. + """ + + model_name = "ltx2.5" + block_classes = [LTX2ReferenceEncoderStep] + block_names = ["reference_encoder"] + block_trigger_inputs = ["reference_conditions"] + + def select_block(self, reference_conditions=None) -> str | None: + # Falsy covers both `None` and an empty list. `LTX2InContextPipeline` likewise treats a missing reference + # as "no reference tokens" rather than an error: IC-LoRAs that carry their behavior in the adapter weights + # (camera control, style, ...) take no reference video at all. + return "reference_encoder" if reference_conditions else None + + @property + def description(self): + return ( + "Conditional reference-encoder step, run only when `reference_conditions` are supplied.\n" + " - `LTX2ReferenceEncoderStep` encodes the reference videos into extra latent tokens.\n" + " - Skipped otherwise, for IC-LoRAs that take no reference video." + ) + + +LTX25_T2V_BLOCKS = InsertableDict( + [ + ("set_timesteps", LTX2SetTimestepsStep(sigmas_default=DISTILLED_SIGMA_VALUES)), + ("prepare_latents", LTX2PrepareLatentsStep()), + ("prepare_audio_latents", LTX2PrepareAudioLatentsStep()), + ("prepare_coords", LTX2PrepareCoordsStep()), + ("denoise", LTX2DenoiseStep()), + ("unpack", LTX2UnpackLatentsStep()), + ] ) -from .modular_blocks_ltx2 import ( - LTX2AutoConditionEncoderStep, - LTX2AutoCoreDenoiseStep, - LTX2AutoDurationStep, - LTX2AutoPromptEnhancerStep, - LTX2AutoReferenceEncoderStep, - LTX2AutoVaeEncoderStep, - LTX2TextConditioningStep, + + +# auto_docstring +class LTX25CoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block (text-to-video) that expands the text conditioning by `num_videos_per_prompt`, prepares video/audio + latents and runs the joint denoising loop. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + sigmas (`list`, *optional*, defaults to [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`): + The number of frames in the generated video. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + block_classes = LTX25_T2V_BLOCKS.values() + block_names = LTX25_T2V_BLOCKS.keys() + + @property + def description(self): + return ( + "Denoise block (text-to-video) that expands the text conditioning by `num_videos_per_prompt`, prepares " + "video/audio latents and runs the joint denoising loop." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="Denoised audio latents."), + ] + + +LTX25_T2V_STAGE_2_BLOCKS = InsertableDict( + [ + ( + "prepare_latents", + LTX2Stage2PrepareLatentsStep(sigmas_name="stage_2_sigmas", sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES), + ), + ( + "set_timesteps", + LTX2SetTimestepsStep( + sigmas_name="stage_2_sigmas", + timesteps_name="stage_2_timesteps", + sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES, + ), + ), + ( + "prepare_audio_latents", + LTX2Stage2PrepareAudioLatentsStep( + sigmas_name="stage_2_sigmas", sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES + ), + ), + ("prepare_coords", LTX2PrepareCoordsStep()), + ("denoise", LTX2DenoiseStep()), + ("unpack", LTX2UnpackLatentsStep()), + ] ) +# auto_docstring +class LTX25Stage2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block (text-to-video, second pass) that expands the text conditioning by `num_videos_per_prompt`, re-noises + the supplied video/audio latents to `noise_scale` and runs the joint denoising loop over them on `stage_2_sigmas` + (the LTX-2 stage-2 distilled schedule by default) -- the refinement pass of the two-stage recipe, or any run that + starts from existing latents. + + Components: + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + latents (`Tensor`): + Video latents to refine, of shape [B, C, F, H, W] (normalized, not packed). + noise_scale (`float`, *optional*): + Noise level the latents are re-noised to before the pass. `None` (default) resolves to `sigmas[0]` when + custom `sigmas` are supplied, else 1.0. + stage_2_sigmas (`list`, *optional*, defaults to [0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + stage_2_timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + audio_latents (`Tensor`): + Audio latents to refine, of shape [B, C, L, M] (normalized, not packed). + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + block_classes = LTX25_T2V_STAGE_2_BLOCKS.values() + block_names = LTX25_T2V_STAGE_2_BLOCKS.keys() + + @property + def description(self): + return ( + "Denoise block (text-to-video, second pass) that expands the text conditioning by " + "`num_videos_per_prompt`, re-noises the supplied video/audio latents to `noise_scale` and runs the joint " + "denoising loop over them on `stage_2_sigmas` (the LTX-2 stage-2 distilled schedule by default) -- the " + "refinement pass of the two-stage recipe, or any run that starts from existing latents." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="Denoised audio latents."), + ] + + +# auto_docstring +class LTX25Image2VideoCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block (image-to-video) that expands the text conditioning by `num_videos_per_prompt`, adds image + conditioning and runs the joint denoising loop. + + Components: + scheduler (`FlowMatchEulerDiscreteScheduler`) transformer (`LTX2VideoTransformer3DModel`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + sigmas (`list`, *optional*, defaults to [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`): + The number of frames in the generated video. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + image_latents (`Tensor`): + VAE-encoded reference-image latents used for image-to-video conditioning. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + block_classes = [ + LTX2SetTimestepsStep(sigmas_default=DISTILLED_SIGMA_VALUES), + LTX2PrepareLatentsStep, + LTX2Image2VideoPrepareLatentsStep, + LTX2PrepareAudioLatentsStep, + LTX2PrepareCoordsStep, + LTX2Image2VideoDenoiseStep, + LTX2UnpackLatentsStep, + ] + block_names = [ + "set_timesteps", + "prepare_latents", + "prepare_i2v_latents", + "prepare_audio_latents", + "prepare_coords", + "denoise", + "unpack", + ] + + @property + def description(self): + return ( + "Denoise block (image-to-video) that expands the text conditioning by `num_videos_per_prompt`, adds " + "image conditioning and runs the joint denoising loop." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="Denoised audio latents."), + ] + + +# auto_docstring +class LTX25Image2VideoStage2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block (image-to-video, second pass) that expands the text conditioning by `num_videos_per_prompt`, + re-noises the supplied video/audio latents to `noise_scale`, adds image conditioning and runs the joint denoising + loop over them on `stage_2_sigmas` (the LTX-2 stage-2 distilled schedule by default). + + Components: + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + latents (`Tensor`): + Video latents to refine, of shape [B, C, F, H, W] (normalized, not packed). + noise_scale (`float`, *optional*): + Noise level the latents are re-noised to before the pass. `None` (default) resolves to `sigmas[0]` when + custom `sigmas` are supplied, else 1.0. + stage_2_sigmas (`list`, *optional*, defaults to [0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + stage_2_timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + image_latents (`Tensor`): + VAE-encoded reference-image latents used for image-to-video conditioning. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + audio_latents (`Tensor`): + Audio latents to refine, of shape [B, C, L, M] (normalized, not packed). + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + block_classes = [ + LTX2Stage2PrepareLatentsStep(sigmas_name="stage_2_sigmas", sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES), + LTX2SetTimestepsStep( + sigmas_name="stage_2_sigmas", + timesteps_name="stage_2_timesteps", + sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES, + ), + LTX2Image2VideoPrepareLatentsStep, + LTX2Stage2PrepareAudioLatentsStep(sigmas_name="stage_2_sigmas", sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES), + LTX2PrepareCoordsStep, + LTX2Image2VideoDenoiseStep, + LTX2UnpackLatentsStep, + ] + block_names = [ + "prepare_latents", + "set_timesteps", + "prepare_i2v_latents", + "prepare_audio_latents", + "prepare_coords", + "denoise", + "unpack", + ] + + @property + def description(self): + return ( + "Denoise block (image-to-video, second pass) that expands the text conditioning by " + "`num_videos_per_prompt`, re-noises the supplied video/audio latents to `noise_scale`, adds image " + "conditioning and runs the joint denoising loop over them on `stage_2_sigmas` (the LTX-2 stage-2 " + "distilled schedule by default)." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="Denoised audio latents."), + ] + + +# auto_docstring +class LTX25ConditionCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block (condition-to-video) that expands the text conditioning by `num_videos_per_prompt`, applies the frame + conditions to the video latents and runs the joint denoising loop. + + Components: + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + condition_latents (`list`): + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). + condition_strengths (`list`): + Per-condition conditioning strengths. + condition_indices (`list`): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`): + The number of frames in the generated video. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + noise_scale (`float`, *optional*): + Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom + `sigmas` are supplied, else 1.0. + sigmas (`list`, *optional*, defaults to [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + # NOTE: prepare-latents runs *before* set-timesteps here, unlike the text-to-video / image-to-video steps. The + # resolution-aware shift `mu` is computed from the packed latent sequence length, which for condition workflows + # includes the appended keyframe tokens, so the latents have to exist first. This mirrors `LTX2ConditionPipeline` + # (its section 4 runs before section 5). + block_classes = [ + LTX2ConditionPrepareLatentsStep(sigmas_default=DISTILLED_SIGMA_VALUES), + LTX2ConditionSetTimestepsStep(sigmas_default=DISTILLED_SIGMA_VALUES), + LTX2ConditionPrepareAudioLatentsStep, + LTX2ConditionPrepareCoordsStep, + LTX2ConditionDenoiseStep, + LTX2TrimConditionTokensStep, + LTX2UnpackLatentsStep, + ] + block_names = [ + "prepare_latents", + "set_timesteps", + "prepare_audio_latents", + "prepare_coords", + "denoise", + "trim_condition_tokens", + "unpack", + ] + + @property + def description(self): + return ( + "Denoise block (condition-to-video) that expands the text conditioning by `num_videos_per_prompt`, " + "applies the frame conditions to the video latents and runs the joint denoising loop." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="Denoised audio latents."), + ] + + +# auto_docstring +class LTX25ConditionStage2CoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block (condition-to-video, second pass) that expands the text conditioning by `num_videos_per_prompt`, + applies the frame conditions on top of the supplied video latents, re-noises them and the supplied audio latents to + `noise_scale` and runs the joint denoising loop over them on `stage_2_sigmas` (the LTX-2 stage-2 distilled schedule + by default). + + Components: + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + condition_latents (`list`): + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). + condition_strengths (`list`): + Per-condition conditioning strengths. + condition_indices (`list`): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + latents (`Tensor`): + Video latents to refine, of shape [B, C, F, H, W] (normalized, not packed) of the generated video only + (no appended condition tokens). + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + noise_scale (`float`, *optional*): + Noise level the un-conditioned tokens are re-noised to. `None` (default) resolves to `sigmas[0]` when + custom `sigmas` are supplied, else 1.0. + stage_2_sigmas (`list`, *optional*, defaults to [0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + stage_2_timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + audio_latents (`Tensor`): + Audio latents to refine, of shape [B, C, L, M] (normalized, not packed). + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + # NOTE: prepare-latents runs *before* set-timesteps here, unlike the text-to-video / image-to-video steps. The + # resolution-aware shift `mu` is computed from the packed latent sequence length, which for condition workflows + # includes the appended keyframe tokens, so the latents have to exist first. This mirrors `LTX2ConditionPipeline` + # (its section 4 runs before section 5). + block_classes = [ + LTX2ConditionStage2PrepareLatentsStep( + sigmas_name="stage_2_sigmas", sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES + ), + LTX2ConditionSetTimestepsStep( + sigmas_name="stage_2_sigmas", + timesteps_name="stage_2_timesteps", + sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES, + ), + LTX2Stage2PrepareAudioLatentsStep(sigmas_name="stage_2_sigmas", sigmas_default=STAGE_2_DISTILLED_SIGMA_VALUES), + LTX2ConditionPrepareCoordsStep, + LTX2ConditionDenoiseStep, + LTX2TrimConditionTokensStep, + LTX2UnpackLatentsStep, + ] + block_names = [ + "prepare_latents", + "set_timesteps", + "prepare_audio_latents", + "prepare_coords", + "denoise", + "trim_condition_tokens", + "unpack", + ] + + @property + def description(self): + return ( + "Denoise block (condition-to-video, second pass) that expands the text conditioning by " + "`num_videos_per_prompt`, applies the frame conditions on top of the supplied video latents, re-noises " + "them and the supplied audio latents to `noise_scale` and runs the joint denoising loop over them on " + "`stage_2_sigmas` (the LTX-2 stage-2 distilled schedule by default)." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="Denoised audio latents."), + ] + + +# auto_docstring +class LTX25InContextCoreDenoiseStep(SequentialPipelineBlocks): + """ + Denoise block (in-context) that expands the text conditioning by `num_videos_per_prompt`, folds the frame + conditions and the IC-LoRA reference tokens into one latent sequence and runs the joint denoising loop. Reuses the + condition denoise step unchanged: reference tokens are pinned by the same x0 blend as frame conditions, matching + the reference implementation's uniform treatment of both. + + Components: + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + condition_latents (`list`): + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). + condition_strengths (`list`): + Per-condition conditioning strengths. + condition_indices (`list`): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + reference_conditions (`list`, *optional*): + `LTX2ReferenceCondition` (or list of them); only their `strength` is read here. Omit for IC-LoRAs that + carry their behavior in the adapter weights and take no reference video. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). + reference_downscale_factor (`int`, *optional*, defaults to 1): + Ratio between the target and reference resolutions. The reference tokens' spatial coordinates are scaled + by it so they land in the target coordinate space, preserving the positional relationship the IC-LoRA was + trained on. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`): + The number of frames in the generated video. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + noise_scale (`float`, *optional*): + Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom + `sigmas` are supplied, else 1.0. + sigmas (`list`, *optional*, defaults to [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + # Same ordering rationale as `LTX25ConditionCoreDenoiseStep`: prepare-latents precedes set-timesteps because + # `mu` is read off the packed sequence length, which here includes both keyframe and reference tokens. + block_classes = [ + LTX2InContextPrepareLatentsStep(sigmas_default=DISTILLED_SIGMA_VALUES), + LTX2ConditionSetTimestepsStep(sigmas_default=DISTILLED_SIGMA_VALUES), + LTX2ConditionPrepareAudioLatentsStep, + LTX2ConditionPrepareCoordsStep, + LTX2ConditionDenoiseStep, + LTX2TrimConditionTokensStep, + LTX2UnpackLatentsStep, + ] + block_names = [ + "prepare_latents", + "set_timesteps", + "prepare_audio_latents", + "prepare_coords", + "denoise", + "trim_condition_tokens", + "unpack", + ] + + @property + def description(self): + return ( + "Denoise block (in-context) that expands the text conditioning by `num_videos_per_prompt`, folds the " + "frame conditions and the IC-LoRA reference tokens into one latent sequence and runs the joint denoising " + "loop. Reuses the condition denoise step unchanged: " + "reference tokens are pinned by the same x0 blend as frame conditions, matching the reference " + "implementation's uniform treatment of both." + ) + + @property + def outputs(self): + return [ + OutputParam.template("latents"), + OutputParam("audio_latents", type_hint=torch.Tensor, description="Denoised audio latents."), + ] + + +# auto_docstring +class LTX25AutoCoreDenoiseStep(ConditionalPipelineBlocks): + """ + Auto denoise block that selects the workflow based on inputs. + - `LTX25InContextCoreDenoiseStep` when `reference_conditions` are provided (in-context / IC-LoRA). + - `LTX25ConditionCoreDenoiseStep` when `condition_latents` are provided (condition-to-video). + - `LTX25Image2VideoCoreDenoiseStep` when `image_latents` is provided. + - `LTX25CoreDenoiseStep` otherwise (text-to-video). + + Components: + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + condition_latents (`list`, *optional*): + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). + condition_strengths (`list`, *optional*): + Per-condition conditioning strengths. + condition_indices (`list`, *optional*): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`, *optional*): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + reference_conditions (`list`, *optional*): + `LTX2ReferenceCondition` (or list of them); only their `strength` is read here. Omit for IC-LoRAs that + carry their behavior in the adapter weights and take no reference video. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). + reference_downscale_factor (`int`, *optional*, defaults to 1): + Ratio between the target and reference resolutions. The reference tokens' spatial coordinates are scaled + by it so they land in the target coordinate space, preserving the positional relationship the IC-LoRA was + trained on. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`): + The number of frames in the generated video. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + noise_scale (`float`, *optional*): + Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom + `sigmas` are supplied, else 1.0. + sigmas (`list`, *optional*, defaults to [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. + timesteps (`Tensor`): + Timesteps for the denoising process. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + image_latents (`Tensor`, *optional*): + VAE-encoded reference-image latents used for image-to-video conditioning. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + block_classes = [ + LTX25InContextCoreDenoiseStep, + LTX25ConditionCoreDenoiseStep, + LTX25Image2VideoCoreDenoiseStep, + LTX25CoreDenoiseStep, + ] + block_names = ["in_context", "condition", "image2video", "text2video"] + block_trigger_inputs = ["reference_conditions", "condition_latents", "image_latents"] + default_block_name = "text2video" + + def select_block(self, reference_conditions=None, condition_latents=None, image_latents=None) -> str | None: + # An IC-LoRA that takes no reference video lands on the condition branch, which is the right answer rather + # than a fallback: `LTX2InContextPrepareLatentsStep` with no reference tokens does exactly what + # `LTX2ConditionPrepareLatentsStep` does, and the extra `num_ref_tokens` it emits is only read by the + # attention-mask construction, which is skipped in that case anyway. + if reference_conditions: + return "in_context" + if condition_latents is not None: + return "condition" + if image_latents is not None: + return "image2video" + return "text2video" + + @property + def description(self): + return ( + "Auto denoise block that selects the workflow based on inputs.\n" + " - `LTX25InContextCoreDenoiseStep` when `reference_conditions` are provided (in-context / IC-LoRA).\n" + " - `LTX25ConditionCoreDenoiseStep` when `condition_latents` are provided (condition-to-video).\n" + " - `LTX25Image2VideoCoreDenoiseStep` when `image_latents` is provided.\n" + " - `LTX25CoreDenoiseStep` otherwise (text-to-video).\n" + ) + + +# auto_docstring +class LTX25AutoStage2CoreDenoiseStep(ConditionalPipelineBlocks): + """ + Auto denoise block for the second pass of the two-stage recipe, selecting the workflow based on inputs. Each branch re-noises the video / audio latents in state on `stage_2_sigmas` instead of sampling fresh noise: + - `LTX25ConditionStage2CoreDenoiseStep` when `condition_latents` are provided (condition-to-video; also the + second pass of an in-context run, whose references shape the first pass only). + - `LTX25Image2VideoStage2CoreDenoiseStep` when `image_latents` is provided. + - `LTX25Stage2CoreDenoiseStep` otherwise (text-to-video). + + Components: + transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) guider + (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + + Inputs: + condition_latents (`list`, *optional*): + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). + condition_strengths (`list`, *optional*): + Per-condition conditioning strengths. + condition_indices (`list`, *optional*): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`, *optional*): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + latents (`Tensor`): + Video latents to refine, of shape [B, C, F, H, W] (normalized, not packed) of the generated video only + (no appended condition tokens). + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + noise_scale (`float`, *optional*): + Noise level the un-conditioned tokens are re-noised to. `None` (default) resolves to `sigmas[0]` when + custom `sigmas` are supplied, else 1.0. + stage_2_sigmas (`list`, *optional*, defaults to [0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + stage_2_timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + audio_latents (`Tensor`): + Audio latents to refine, of shape [B, C, L, M] (normalized, not packed). + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`, *optional*): + Video-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_audio_prompt_embeds (`Tensor`, *optional*): + Audio-branch text conditioning (uncond); read only under classifier-free guidance. + negative_connector_attention_mask (`Tensor`, *optional*): + Binary text attention mask (uncond); read only under classifier-free guidance. + image_latents (`Tensor`, *optional*): + VAE-encoded reference-image latents used for image-to-video conditioning. + + Outputs: + latents (`Tensor`): + Denoised latents. + audio_latents (`Tensor`): + Denoised audio latents. + """ + + model_name = "ltx2.5" + block_classes = [ + LTX25ConditionStage2CoreDenoiseStep, + LTX25Image2VideoStage2CoreDenoiseStep, + LTX25Stage2CoreDenoiseStep, + ] + block_names = ["condition", "image2video", "text2video"] + block_trigger_inputs = ["condition_latents", "image_latents"] + default_block_name = "text2video" + + def select_block(self, condition_latents=None, image_latents=None) -> str | None: + # The second pass of an in-context run is a plain condition pass: the references only shape the first pass, + # as in `LTX2InContextPipeline`. + if condition_latents is not None: + return "condition" + if image_latents is not None: + return "image2video" + return "text2video" + + @property + def description(self): + return ( + "Auto denoise block for the second pass of the two-stage recipe, selecting the workflow based on " + "inputs. Each branch re-noises the video / audio latents in state on `stage_2_sigmas` instead of " + "sampling fresh noise:\n" + " - `LTX25ConditionStage2CoreDenoiseStep` when `condition_latents` are provided (condition-to-video; " + "also the second pass of an in-context run, whose references shape the first pass only).\n" + " - `LTX25Image2VideoStage2CoreDenoiseStep` when `image_latents` is provided.\n" + " - `LTX25Stage2CoreDenoiseStep` otherwise (text-to-video)." + ) + + # auto_docstring class LTX25DecoderStep(SequentialPipelineBlocks): """ - Decode stage for LTX-2.5: denoises the video latents with the diffusion decoder and vocodes the audio latents (or - returns latents). + Decode stage for LTX-2.5: denoises the video latents with the diffusion decoder and vocodes the audio latents. Components: diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) video_processor (`VideoProcessor`) audio_vae @@ -44,23 +1369,13 @@ class LTX25DecoderStep(SequentialPipelineBlocks): Inputs: latents (`Tensor`): - Pre-generated noisy latents for image generation. + Video latents of shape [B, C, F, H, W] (normalized, not packed). output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 704): - The width in pixels of the generated image. - num_frames (`int`, *optional*, defaults to 121): - The number of frames in the generated video. generator (`Generator`, *optional*): Torch generator for deterministic generation. - dtype (`dtype`): - The dtype of the model inputs, can be generated in input step. audio_latents (`Tensor`): - Denoised audio latents. - audio_num_frames (`int`): - Number of audio latent frames, used to unpack the audio latent sequence. + Audio latents of shape [B, C, L, M] (normalized, not packed). Outputs: videos (`list`): @@ -77,7 +1392,7 @@ class LTX25DecoderStep(SequentialPipelineBlocks): def description(self): return ( "Decode stage for LTX-2.5: denoises the video latents with the diffusion decoder and vocodes the audio " - "latents (or returns latents)." + "latents." ) @property @@ -89,36 +1404,112 @@ def outputs(self): # auto_docstring -class LTX25ConditionDecoderStep(SequentialPipelineBlocks): +class LTX25AutoBlocks(SequentialPipelineBlocks): """ - Decode stage for LTX-2.5 condition workflows: drops the appended keyframe-condition tokens, then denoises the video - latents with the diffusion decoder and vocodes the audio latents (or returns latents). + Auto blocks for LTX-2.5 supporting text-to-video, image-to-video, condition-to-video and in-context (IC-LoRA) + generation (joint video + audio). Identical to `LTX2AutoBlocks` except that the video decoder is + `LTX2DiffusionVaeDecoderStep`, since the diffusion decoder is the native default from LTX-2.5 on. To decode with + the convolutional VAE instead, swap the decode block: `blocks.sub_blocks["decode"] = LTX2DecoderStep()`. + + Supported workflows: + - `text2video`: requires `prompt` + - `image2video`: requires `image`, `prompt` + - `condition`: requires `conditions`, `prompt` + - `in_context`: requires `reference_conditions`, `num_frames`, `prompt` Components: - diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) video_processor (`VideoProcessor`) audio_vae - (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) + prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer + (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) vae + (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) transformer (`LTX2VideoTransformer3DModel`) + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) audio_vae (`AutoencoderKLLTX2Audio`) vocoder + (`LTX2Vocoder`) Inputs: - latents (`Tensor`): - Pre-generated noisy latents for image generation. - base_token_count (`int`): - Number of generated-video tokens, i.e. the sequence length before appended tokens. - output_type (`str`, *optional*, defaults to pil): - Output format: 'pil', 'np', 'pt'. + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + conditions (`list`, *optional*): + `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the + generated video. + system_prompt (`str`, *optional*): + System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` + condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. + prompt_max_new_tokens (`int`, *optional*): + Maximum number of new tokens to generate during prompt enhancement. Defaults to 600, the LTX-2.5 Gemma-4 + enhancer's budget. + prompt_enhancement_kwargs (`dict`, *optional*): + Keyword arguments for the enhancer's `.generate` call. Defaults to greedy decoding. + prompt_enhancement_seed (`int`, *optional*, defaults to 10): + Random seed for prompt enhancement (inert under LTX-2.5's greedy decoding). + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. + min_seconds (`float`, *optional*, defaults to 1.0): + Lower bound on the auto-predicted duration. + max_seconds (`float`, *optional*, defaults to 20.0): + Upper bound on the auto-predicted duration. Must be strictly greater than `min_seconds`. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. height (`int`, *optional*, defaults to 512): The height in pixels of the generated image. width (`int`, *optional*, defaults to 704): The width in pixels of the generated image. - num_frames (`int`, *optional*, defaults to 121): - The number of frames in the generated video. - generator (`Generator`, *optional*): - Torch generator for deterministic generation. - dtype (`dtype`): - The dtype of the model inputs, can be generated in input step. - audio_latents (`Tensor`): - Denoised audio latents. - audio_num_frames (`int`): - Number of audio latent frames, used to unpack the audio latent sequence. + image_crf (`int`, *optional*): + H.264 CRF used to re-compress the conditioning `image` before VAE encode, matching the compression the + model was trained against. `None` (default) uses the pipeline's `default_image_crf` (33 through LTX-2.3, + 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when re-compression runs. + reference_conditions (`list`, *optional*): + `LTX2ReferenceCondition` (or list of them) whose videos are encoded into extra latent tokens the IC-LoRA + adapter attends to. + reference_downscale_factor (`int`, *optional*, defaults to 1): + Ratio between the target and reference resolutions; 2 means the reference is preprocessed at half the + target resolution. Must match the factor the IC-LoRA was trained with. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + condition_latents (`list`, *optional*): + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). + condition_strengths (`list`, *optional*): + Per-condition conditioning strengths. + condition_indices (`list`, *optional*): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`, *optional*): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). + noise_scale (`float`, *optional*): + Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom + `sigmas` are supplied, else 1.0. + sigmas (`list`, *optional*, defaults to [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. + timesteps (`Tensor`): + Timesteps for the denoising process. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + image_latents (`Tensor`, *optional*): + VAE-encoded reference-image latents used for image-to-video conditioning. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. Outputs: videos (`list`): @@ -128,15 +1519,46 @@ class LTX25ConditionDecoderStep(SequentialPipelineBlocks): """ model_name = "ltx2.5" - block_classes = [LTX2TrimConditionTokensStep, LTX2DiffusionVaeDecoderStep, LTX2AudioDecoderStep] - block_names = ["trim_condition_tokens", "video_decode", "audio_decode"] + block_classes = [ + LTX25AutoPromptEnhancerStep, + LTX25TextConditioningStep, + LTX25AutoDurationStep, + LTX25AutoVaeEncoderStep, + LTX25AutoConditionEncoderStep, + LTX25AutoReferenceEncoderStep, + LTX2TextInputStep, + LTX25AutoCoreDenoiseStep, + LTX25DecoderStep, + ] + block_names = [ + "prompt_enhancer", + "text_encoder", + "duration", + "vae_encoder", + "condition_encoder", + "reference_encoder", + "input", + "denoise", + "decode", + ] + # `num_frames` on `in_context` is a requirement, not a trigger: the in-context checkpoints ship without a + # `duration_head`, so the workflow drops `LTX25AutoDurationStep` and `LTX2ConditionEncoderStep` raises if + # `num_frames` is still `None`. Matches `LTX2InContextBlocks`, which omits the duration step outright. + _workflow_map = { + "text2video": {"prompt": True}, + "image2video": {"image": True, "prompt": True}, + "condition": {"conditions": True, "prompt": True}, + "in_context": {"reference_conditions": True, "num_frames": True, "prompt": True}, + } @property def description(self): return ( - "Decode stage for LTX-2.5 condition workflows: drops the appended keyframe-condition tokens, then " - "denoises the video latents with the diffusion decoder and vocodes the audio latents (or returns " - "latents)." + "Auto blocks for LTX-2.5 supporting text-to-video, image-to-video, condition-to-video and in-context " + "(IC-LoRA) generation (joint video + audio). Identical to `LTX2AutoBlocks` except that the video decoder " + "is `LTX2DiffusionVaeDecoderStep`, since the diffusion decoder is the native default from LTX-2.5 on. To " + 'decode with the convolutional VAE instead, swap the decode block: `blocks.sub_blocks["decode"] = ' + "LTX2DecoderStep()`." ) @property @@ -148,76 +1570,52 @@ def outputs(self): # auto_docstring -class LTX25AutoDecoderStep(AutoPipelineBlocks): +class LTX25UpsampleStep(SequentialPipelineBlocks): """ - Auto decode block for LTX-2.5 that selects the decoder based on inputs. - - `LTX25ConditionDecoderStep` when `base_token_count` is present, i.e. the denoised sequence carries appended - keyframe / reference tokens (condition, in-context). - - `LTX25DecoderStep` otherwise (text-to-video, image-to-video). + Upsample stage of the two-stage recipe: `LTX2LatentUpsampleStep` as an LTX-2.5 block, so that popped into its own + pipeline it resolves the LTX-2.5 latent statistics when no autoencoder is loaded. Components: - diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) video_processor (`VideoProcessor`) audio_vae - (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) + latent_upsampler (`LTX2LatentUpsamplerModel`) transformer (`LTX2VideoTransformer3DModel`) Inputs: latents (`Tensor`): - Pre-generated noisy latents for image generation. - base_token_count (`int`, *optional*): - Number of generated-video tokens, i.e. the sequence length before appended tokens. - output_type (`str`, *optional*, defaults to pil): - Output format: 'pil', 'np', 'pt'. - height (`int`, *optional*, defaults to 512): - The height in pixels of the generated image. - width (`int`, *optional*, defaults to 704): - The width in pixels of the generated image. - num_frames (`int`, *optional*, defaults to 121): - The number of frames in the generated video. - generator (`Generator`, *optional*): - Torch generator for deterministic generation. - dtype (`dtype`): - The dtype of the model inputs, can be generated in input step. - audio_latents (`Tensor`): - Denoised audio latents. - audio_num_frames (`int`): - Number of audio latent frames, used to unpack the audio latent sequence. + Video latents to upsample, of shape [B, C, F, H, W] (normalized, not packed). Outputs: - videos (`list`): - The generated videos. - audio (`Tensor`): - The generated audio waveform. + latents (`Tensor`): + Upsampled video latents of shape [B, C, F, 2H, 2W] (normalized, not packed). + height (`int`): + Height of the upsampled latents, in pixels. + width (`int`): + Width of the upsampled latents, in pixels. """ model_name = "ltx2.5" - # Mirrors `LTX2AutoDecoderStep` on the same `base_token_count` trigger; only the video decoder differs. - block_classes = [LTX25ConditionDecoderStep, LTX25DecoderStep] - block_names = ["condition", "default"] - block_trigger_inputs = ["base_token_count", None] + block_classes = [LTX2LatentUpsampleStep] + block_names = ["latent_upsample"] @property def description(self): return ( - "Auto decode block for LTX-2.5 that selects the decoder based on inputs.\n" - " - `LTX25ConditionDecoderStep` when `base_token_count` is present, i.e. the denoised sequence carries " - "appended keyframe / reference tokens (condition, in-context).\n" - " - `LTX25DecoderStep` otherwise (text-to-video, image-to-video)." + "Upsample stage of the two-stage recipe: `LTX2LatentUpsampleStep` as an LTX-2.5 block, so that popped " + "into its own pipeline it resolves the LTX-2.5 latent statistics when no autoencoder is loaded." ) - @property - def outputs(self): - return [ - OutputParam.template("videos"), - OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform."), - ] - # auto_docstring -class LTX25AutoBlocks(SequentialPipelineBlocks): +class LTX25TwoStageBlocks(SequentialPipelineBlocks): """ - Auto blocks for LTX-2.5 supporting text-to-video, image-to-video, condition-to-video and in-context (IC-LoRA) - generation (joint video + audio). Identical to `LTX2AutoBlocks` except that the video decoder is - `LTX2DiffusionVaeDecoderStep`, since the diffusion decoder is the native default from LTX-2.5 on. To decode with - the convolutional VAE instead, swap the decode block: `blocks.sub_blocks["decode"] = LTX2AutoDecoderStep()`. + Blocks for the LTX-2.5 distilled two-stage recipe (joint video + audio) in one call, for every workflow + `LTX25AutoBlocks` supports: a first pass at the requested `height` / `width`, a 2x latent upsample, and a second + pass that refines at the upsampled resolution, with the diffusion decoder at the end -- so the output is twice the + size asked for, as with the standard pipelines. `stage_1` is the same auto denoise step as `LTX25AutoBlocks`; + `stage_2` selects the workflow's second-pass group, which re-noises the upsampled latents on `stage_2_sigmas`. + Image and frame conditions are re-encoded at the upsampled resolution ahead of the second pass, as the standard + pipelines do on their second call. The text conditioning is expanded by `num_videos_per_prompt` once, by the + `input` step, so `stage_1` / `upsample` / `stage_2` / `decode` can each be popped and run as their own pipeline: + `stage_1` followed by `decode` previews the first pass, and popping `stage_1` and `upsample` leaves a standalone + second pass that takes `latents` / `audio_latents`. Supported workflows: - `text2video`: requires `prompt` @@ -229,8 +1627,9 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) transformer (`LTX2VideoTransformer3DModel`) - scheduler (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) - audio_guider (`LTX2Guidance`) diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) vocoder (`LTX2Vocoder`) + scheduler (`FlowMatchEulerDiscreteScheduler`) guider (`LTX2Guidance`) audio_guider (`LTX2Guidance`) + latent_upsampler (`LTX2LatentUpsamplerModel`) diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) audio_vae + (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) Inputs: prompt (`str`, *optional*): @@ -238,8 +1637,6 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): conditions (`list`, *optional*): `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the generated video. - enable_prompt_enhancement (`bool`, *optional*, defaults to False): - Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. system_prompt (`str`, *optional*): System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. @@ -254,10 +1651,15 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): Torch generator for deterministic generation. image (`Image | list`, *optional*): Reference image(s) for denoising. Can be a single image or list of images. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. negative_prompt (`str`, *optional*): The prompt or prompts not to guide the image generation. max_sequence_length (`int`, *optional*, defaults to 1024): Maximum sequence length for prompt encoding. + num_frames (`int`, *optional*): + The number of frames in the generated video. Omit to have this step predict it with the `duration_head`; + the denoise blocks then take the predicted count. min_seconds (`float`, *optional*, defaults to 1.0): Lower bound on the auto-predicted duration. max_seconds (`float`, *optional*, defaults to 20.0): @@ -270,66 +1672,52 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): The width in pixels of the generated image. image_crf (`int`, *optional*): H.264 CRF used to re-compress the conditioning `image` before VAE encode, matching the compression the - model was trained against. `None` (default) resolves from the text-encoder generation (33 through - LTX-2.3, 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when - re-compression runs. - num_frames (`int`, *optional*): - The number of frames in the generated video. Omit to auto-predict via the `duration_head` (see - `LTX2AutoDurationStep`). + model was trained against. `None` (default) uses the pipeline's `default_image_crf` (33 through LTX-2.3, + 18 for LTX-2.5). Pass `0` to skip re-compression. Requires a `PIL.Image.Image` when re-compression runs. reference_conditions (`list`, *optional*): `LTX2ReferenceCondition` (or list of them) whose videos are encoded into extra latent tokens the IC-LoRA adapter attends to. reference_downscale_factor (`int`, *optional*, defaults to 1): Ratio between the target and reference resolutions; 2 means the reference is preprocessed at half the - target resolution. Spatial coordinates are scaled by this factor so the reference tokens land in the - target coordinate space. Must match the factor the IC-LoRA was trained with. - conditioning_attention_strength (`float`, *optional*, defaults to 1.0): - Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 - (default) leaves attention unmasked. - conditioning_attention_mask (`Tensor`, *optional*): - Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying - attention strength. Downsampled to the reference's latent grid and multiplied by - `conditioning_attention_strength`. + target resolution. Must match the factor the IC-LoRA was trained with. num_videos_per_prompt (`int`, *optional*, defaults to 1): The number of images to generate per prompt. condition_latents (`list`, *optional*): - Per-condition normalized VAE latents of shape [1, C, F, H, W]. + Per-condition VAE latents of shape [1, C, F, H, W] (normalized, not packed). condition_strengths (`list`, *optional*): Per-condition conditioning strengths. condition_indices (`list`, *optional*): Per-condition latent frame index at which the condition is applied. condition_pixel_frames (`list`, *optional*): Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. - reference_latents (`Tensor`, *optional*): - Packed reference tokens of shape [1, total_reference_tokens, C], or `None` when no reference conditions - were supplied (`LTX2AutoReferenceEncoderStep` is skipped). - reference_coords (`Tensor`, *optional*): - RoPE coordinates for the reference tokens. - reference_token_counts (`list`, *optional*): - Per-reference token counts, in `reference_conditions` order. - latents (`Tensor`): - Pre-generated noisy latents for image generation. + reference_latents (`list`, *optional*): + Per-reference VAE latents of shape [1, C, F, H, W] (normalized, not packed) from + `LTX2ReferenceEncoderStep`, or `None` when no reference conditions were supplied + (`LTX2AutoReferenceEncoderStep` is skipped). noise_scale (`float`, *optional*): Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom `sigmas` are supplied, else 1.0. - sigmas (`list`, *optional*): + sigmas (`list`, *optional*, defaults to [1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875]): Custom sigmas for the denoising process. - reference_cross_mask (`Tensor`, *optional*): - Per-reference-token noisy<->reference attention strengths of shape [1, num_ref_tokens]. - num_inference_steps (`int`): - The number of denoising steps. + conditioning_attention_strength (`float`, *optional*, defaults to 1.0): + Scalar in [0, 1] controlling how strongly the noisy tokens and reference tokens attend to each other. 1.0 + (default) leaves attention unmasked. + conditioning_attention_mask (`Tensor`, *optional*): + Optional pixel-space mask of shape (1, 1, F, H, W) with values in [0, 1] giving spatially varying + attention strength. Downsampled to each reference's latent grid and multiplied by + `conditioning_attention_strength`. timesteps (`Tensor`): Timesteps for the denoising process. - audio_latents (`Tensor`): - Optional pre-encoded audio latents; random noise is used when not provided. **denoiser_input_fields (`None`, *optional*): conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. - use_cross_timestep (`bool`, *optional*, defaults to True): - Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). attention_kwargs (`dict`, *optional*): Additional kwargs for attention processors. image_latents (`Tensor`, *optional*): VAE-encoded reference-image latents used for image-to-video conditioning. + stage_2_sigmas (`list`, *optional*, defaults to [0.909375, 0.725, 0.421875]): + Custom sigmas for the denoising process. + stage_2_timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. @@ -340,16 +1728,21 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): The generated audio waveform. """ - model_name = "ltx2.5" + model_name = "ltx2.5-two-stage" block_classes = [ - LTX2AutoPromptEnhancerStep, - LTX2TextConditioningStep, - LTX2AutoDurationStep, - LTX2AutoVaeEncoderStep, - LTX2AutoConditionEncoderStep, - LTX2AutoReferenceEncoderStep, - LTX2AutoCoreDenoiseStep, - LTX25AutoDecoderStep, + LTX25AutoPromptEnhancerStep, + LTX25TextConditioningStep, + LTX25AutoDurationStep, + LTX25AutoVaeEncoderStep, + LTX25AutoConditionEncoderStep, + LTX25AutoReferenceEncoderStep, + LTX2TextInputStep, + LTX25AutoCoreDenoiseStep, + LTX25UpsampleStep, + LTX25AutoVaeEncoderStep, + LTX25AutoConditionEncoderStep, + LTX25AutoStage2CoreDenoiseStep, + LTX25DecoderStep, ] block_names = [ "prompt_enhancer", @@ -358,12 +1751,14 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): "vae_encoder", "condition_encoder", "reference_encoder", - "denoise", + "input", + "stage_1", + "upsample", + "stage_2_vae_encoder", + "stage_2_condition_encoder", + "stage_2", "decode", ] - # `num_frames` on `in_context` is a requirement, not a trigger: the in-context checkpoints ship without a - # `duration_head`, so the workflow drops `LTX2AutoDurationStep` and `LTX2ConditionEncoderStep` raises if - # `num_frames` is still `None`. Matches `LTX2InContextBlocks`, which omits the duration step outright. _workflow_map = { "text2video": {"prompt": True}, "image2video": {"image": True, "prompt": True}, @@ -374,11 +1769,17 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): @property def description(self): return ( - "Auto blocks for LTX-2.5 supporting text-to-video, image-to-video, condition-to-video and in-context " - "(IC-LoRA) generation (joint video + audio). Identical to `LTX2AutoBlocks` except that the video decoder " - "is `LTX2DiffusionVaeDecoderStep`, since the diffusion decoder is the native default from LTX-2.5 on. To " - 'decode with the convolutional VAE instead, swap the decode block: `blocks.sub_blocks["decode"] = ' - "LTX2AutoDecoderStep()`." + "Blocks for the LTX-2.5 distilled two-stage recipe (joint video + audio) in one call, for every workflow " + "`LTX25AutoBlocks` supports: a first pass at the requested `height` / `width`, a 2x latent upsample, and " + "a second pass that refines at the upsampled resolution, with the diffusion decoder at the end -- so the " + "output is twice the size asked for, as with the standard pipelines. `stage_1` is the same auto denoise " + "step as `LTX25AutoBlocks`; `stage_2` selects the workflow's second-pass group, which re-noises the " + "upsampled latents on `stage_2_sigmas`. Image and frame conditions are re-encoded at the upsampled " + "resolution ahead of the second pass, as the standard pipelines do on their second call. The text " + "conditioning is expanded by `num_videos_per_prompt` once, by the `input` step, so `stage_1` / " + "`upsample` / `stage_2` / `decode` can each be popped and run as their own pipeline: `stage_1` followed " + "by `decode` previews the first pass, and popping `stage_1` and `upsample` leaves a standalone second " + "pass that takes `latents` / `audio_latents`." ) @property diff --git a/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py b/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py index c2e3409bbf37..eb28342bb107 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py @@ -15,9 +15,19 @@ import torch from ...loaders import LTX2LoraLoaderMixin +from ...pipelines.ltx2.utils import DEFAULT_IMAGE_CRF, LTX2_5_IMAGE_CRF from ...utils import logging from ..modular_pipeline import ModularPipeline -from .utils import LTX2_AUDIO_LATENTS_MEAN, LTX2_AUDIO_LATENTS_STD, LTX2_LATENTS_MEAN, LTX2_LATENTS_STD +from .utils import ( + LTX2_AUDIO_LATENTS_MEAN, + LTX2_AUDIO_LATENTS_STD, + LTX2_LATENTS_MEAN, + LTX2_LATENTS_STD, + LTX25_AUDIO_LATENTS_MEAN, + LTX25_AUDIO_LATENTS_STD, + LTX25_LATENTS_MEAN, + LTX25_LATENTS_STD, +) logger = logging.get_logger(__name__) @@ -34,21 +44,26 @@ class LTX2ModularPipeline( default_blocks_name = "LTX2AutoBlocks" + # The video latent geometry and statistics live on whichever video autoencoder the checkpoint ships: the conv + # `vae`, or the `diffusion_decoder`, which carries the same ratios and buffers and is the only one registered when + # a checkpoint decodes with it and does not need the conv encoder -- or when only the decode blocks are run. + # The literal fallbacks are the values of `Lightricks/LTX-2`. @property def vae_spatial_compression_ratio(self): if getattr(self, "vae", None) is not None: return self.vae.spatial_compression_ratio + if getattr(self, "diffusion_decoder", None) is not None: + return self.diffusion_decoder.spatial_compression_ratio return 32 @property def vae_temporal_compression_ratio(self): if getattr(self, "vae", None) is not None: return self.vae.temporal_compression_ratio + if getattr(self, "diffusion_decoder", None) is not None: + return self.diffusion_decoder.temporal_compression_ratio return 8 - # The video latent statistics live on whichever video autoencoder the checkpoint ships: the conv `vae`, or the - # `diffusion_decoder`, which carries the same buffers and is the only one registered when a checkpoint decodes - # with it and does not need the conv encoder. `LTX2_LATENTS_*` are the values of `Lightricks/LTX-2`. @property def vae_scaling_factor(self): if getattr(self, "vae", None) is not None: @@ -97,6 +112,19 @@ def audio_vae_temporal_compression_ratio(self): return self.audio_vae.temporal_compression_ratio return 4 + @property + def audio_latent_channels(self): + if getattr(self, "audio_vae", None) is not None: + return self.audio_vae.config.latent_channels + return 8 + + @property + def audio_latent_mel_bins(self): + # Mel bins of the audio latent grid: the audio VAE's mel bins after its mel compression. + if getattr(self, "audio_vae", None) is not None: + return self.audio_vae.config.mel_bins // self.audio_vae.mel_compression_ratio + return 16 + @property def audio_latents_mean(self): if getattr(self, "audio_vae", None) is not None: @@ -121,6 +149,29 @@ def audio_hop_length(self): return self.audio_vae.config.mel_hop_length return 160 + # Whether the transformer's cross-modality attention is modulated by the other modality's sigma: the LTX-2.3 / + # LTX-2.5 behaviour. LTX-2.0 checkpoints used the modality's own timestep (`False`). + @property + def use_cross_timestep(self) -> bool: + return True + + # Default H.264 CRF used to re-compress conditioning images before VAE encoding (LTX-2: 33, LTX-2.5: 18). + @property + def default_image_crf(self) -> int: + return DEFAULT_IMAGE_CRF + + # Whether the pipeline will require unconditional (negative-prompt) embeddings: only when a guider is registered + # and running classifier-free guidance. + @property + def requires_unconditional_embeds(self): + for name in ("guider", "audio_guider"): + guider = getattr(self, name, None) + if guider is None or not guider._enabled: + continue + if guider.is_cfg_enabled() if hasattr(guider, "is_cfg_enabled") else guider.num_conditions > 1: + return True + return False + class LTX25ModularPipeline(LTX2ModularPipeline): """ @@ -132,3 +183,50 @@ class LTX25ModularPipeline(LTX2ModularPipeline): """ default_blocks_name = "LTX25AutoBlocks" + + @property + def default_image_crf(self) -> int: + return LTX2_5_IMAGE_CRF + + # The LTX-2.5 autoencoders' latent statistics, for blocks that run without one loaded (a popped `upsample` or + # `stage_2` pipeline); a loaded `vae` / `diffusion_decoder` / `audio_vae` still wins. + @property + def latents_mean(self): + if getattr(self, "vae", None) is not None: + return self.vae.latents_mean + if getattr(self, "diffusion_decoder", None) is not None: + return self.diffusion_decoder.latents_mean + return torch.tensor(LTX25_LATENTS_MEAN) + + @property + def latents_std(self): + if getattr(self, "vae", None) is not None: + return self.vae.latents_std + if getattr(self, "diffusion_decoder", None) is not None: + return self.diffusion_decoder.latents_std + return torch.tensor(LTX25_LATENTS_STD) + + @property + def audio_latents_mean(self): + if getattr(self, "audio_vae", None) is not None: + return self.audio_vae.latents_mean + return torch.tensor(LTX25_AUDIO_LATENTS_MEAN) + + @property + def audio_latents_std(self): + if getattr(self, "audio_vae", None) is not None: + return self.audio_vae.latents_std + return torch.tensor(LTX25_AUDIO_LATENTS_STD) + + +class LTX25TwoStageModularPipeline(LTX25ModularPipeline): + """ + A ModularPipeline for the LTX-2.5 distilled two-stage recipe (joint video + audio generation): a first pass at half + the target resolution, a 2x latent upsample, and a second pass that refines at the target resolution. + + Identical to [`LTX25ModularPipeline`] except for its default blocks, [`LTX25TwoStageBlocks`]. A checkpoint routes + here through `modular_model_index.json`. + + """ + + default_blocks_name = "LTX25TwoStageBlocks" diff --git a/src/diffusers/modular_pipelines/ltx2/utils.py b/src/diffusers/modular_pipelines/ltx2/utils.py index 814af7ccdff4..792944615aa5 100644 --- a/src/diffusers/modular_pipelines/ltx2/utils.py +++ b/src/diffusers/modular_pipelines/ltx2/utils.py @@ -96,3 +96,526 @@ 1.3125, 1.289062, 1.296875, 1.242188, 1.234375, 1.21875, 1.226562, 1.054688, ] # fmt: on + +# Latent statistics of the `Lightricks/LTX-2.5-Diffusers` video / audio VAEs (the `latents_mean` / `latents_std` +# buffers), the fallback of `LTX25ModularPipeline` for blocks that run without an autoencoder loaded. +LTX25_LATENTS_MEAN = [ + 0.022216796875, + 0.0028228759765625, + -0.01214599609375, + -0.01318359375, + 0.0137939453125, + -0.022705078125, + -0.01275634765625, + -0.01116943359375, + 0.01416015625, + -0.011962890625, + -7.772445678710938e-05, + -0.00127410888671875, + 0.01220703125, + -0.05029296875, + 0.0003566741943359375, + 0.0026397705078125, + 0.006439208984375, + 0.005767822265625, + 0.004180908203125, + 0.00634765625, + -0.0269775390625, + 0.00653076171875, + 0.0031890869140625, + 0.0361328125, + -0.05810546875, + -0.0024566650390625, + 0.045166015625, + -0.00140380859375, + 0.0174560546875, + -0.00194549560546875, + -0.00982666015625, + 0.00567626953125, + 0.000247955322265625, + -0.000820159912109375, + 0.005218505859375, + -0.0103759765625, + 0.006561279296875, + 0.01165771484375, + -0.0084228515625, + 0.00098419189453125, + 0.00433349609375, + 0.051513671875, + 0.0001850128173828125, + -0.0113525390625, + 0.00823974609375, + -0.330078125, + -0.005615234375, + 0.00982666015625, + -0.00677490234375, + 0.00118255615234375, + -0.017333984375, + -0.0164794921875, + -0.001983642578125, + -0.00147247314453125, + 0.0020599365234375, + 0.01904296875, + 0.00848388671875, + -0.10888671875, + -0.008056640625, + 0.003936767578125, + -0.0091552734375, + 0.000926971435546875, + 0.01373291015625, + 0.166015625, + -0.01055908203125, + 0.027587890625, + 0.000881195068359375, + 0.00775146484375, + 0.03466796875, + -0.00157928466796875, + 0.0030670166015625, + -0.0020751953125, + -0.10205078125, + 0.00118255615234375, + 0.01165771484375, + 0.427734375, + 0.001373291015625, + 0.004730224609375, + -0.00653076171875, + -0.00811767578125, + -0.02294921875, + -0.01318359375, + 0.004150390625, + 0.056640625, + 0.0021514892578125, + 0.002197265625, + -0.09814453125, + -0.0145263671875, + -0.0036163330078125, + -0.00482177734375, + 0.002044677734375, + -0.0234375, + -0.0498046875, + 0.0059814453125, + -0.006439208984375, + 0.0020904541015625, + 0.006866455078125, + 0.2294921875, + -0.044677734375, + -0.00738525390625, + 0.0045166015625, + -0.0118408203125, + 0.0010986328125, + -0.5546875, + 0.01361083984375, + -0.031494140625, + 0.0026702880859375, + -0.00970458984375, + 0.11572265625, + -0.0654296875, + -0.006744384765625, + -0.010498046875, + 0.00201416015625, + -0.000820159912109375, + -0.013427734375, + -0.00102996826171875, + 0.306640625, + 0.0078125, + -0.00823974609375, + 0.01141357421875, + -0.00396728515625, + -0.035400390625, + 0.00384521484375, + 0.00531005859375, + -0.000896453857421875, + -0.00019168853759765625, + 0.02978515625, + -0.0059814453125, +] +LTX25_LATENTS_STD = [ + 0.23828125, + 0.0751953125, + 0.103515625, + 0.0947265625, + 0.123046875, + 0.1396484375, + 0.2236328125, + 0.365234375, + 0.1572265625, + 0.1865234375, + 0.09228515625, + 0.140625, + 0.166015625, + 0.251953125, + 0.08837890625, + 0.12353515625, + 0.138671875, + 0.10986328125, + 0.10693359375, + 0.138671875, + 0.197265625, + 0.103515625, + 0.1328125, + 0.208984375, + 0.1796875, + 0.1474609375, + 0.162109375, + 0.1298828125, + 0.1181640625, + 0.0771484375, + 0.146484375, + 0.0849609375, + 0.08544921875, + 0.0751953125, + 0.125, + 0.095703125, + 0.125, + 0.1376953125, + 0.09375, + 0.09130859375, + 0.07763671875, + 0.177734375, + 0.08447265625, + 0.09521484375, + 0.1240234375, + 0.1904296875, + 0.09375, + 0.08544921875, + 0.15234375, + 0.09326171875, + 0.1416015625, + 0.1298828125, + 0.08203125, + 0.12109375, + 0.10009765625, + 0.193359375, + 0.0947265625, + 0.6953125, + 0.087890625, + 0.08056640625, + 0.11865234375, + 0.08251953125, + 0.109375, + 0.53125, + 0.2119140625, + 0.2255859375, + 0.109375, + 0.138671875, + 0.15625, + 0.0966796875, + 0.09521484375, + 0.09423828125, + 0.486328125, + 0.07421875, + 0.11669921875, + 0.75390625, + 0.12451171875, + 0.09033203125, + 0.12353515625, + 0.10498046875, + 0.16796875, + 0.1435546875, + 0.1650390625, + 0.16796875, + 0.1005859375, + 0.07470703125, + 0.208984375, + 0.1396484375, + 0.1005859375, + 0.126953125, + 0.09423828125, + 0.1650390625, + 0.23828125, + 0.126953125, + 0.1416015625, + 0.1533203125, + 0.2021484375, + 0.59765625, + 0.19140625, + 0.087890625, + 0.1259765625, + 0.17578125, + 0.09912109375, + 0.9140625, + 0.109375, + 0.154296875, + 0.1005859375, + 0.1279296875, + 0.322265625, + 0.2080078125, + 0.345703125, + 0.12890625, + 0.16015625, + 0.150390625, + 0.1865234375, + 0.2158203125, + 0.7109375, + 0.0986328125, + 0.10009765625, + 0.12060546875, + 0.107421875, + 0.171875, + 0.1240234375, + 0.107421875, + 0.08447265625, + 0.09423828125, + 0.1826171875, + 0.11669921875, +] +LTX25_AUDIO_LATENTS_MEAN = [ + 1.34375, + 2.734375, + 2.296875, + 1.8984375, + 1.4609375, + 1.203125, + 0.98828125, + 0.859375, + 0.6328125, + 0.39453125, + 0.177734375, + -0.140625, + -0.51171875, + -0.86328125, + -1.1953125, + -0.79296875, + 0.59375, + -0.56640625, + -0.37890625, + -0.23046875, + -0.0277099609375, + 0.115234375, + 0.2265625, + 0.2734375, + 0.333984375, + 0.431640625, + 0.494140625, + 0.62890625, + 0.7265625, + 0.98828125, + 0.58203125, + 0.7734375, + 0.75, + 0.337890625, + 0.33203125, + 0.11669921875, + -0.035888671875, + -0.185546875, + -0.2294921875, + -0.3046875, + -0.4453125, + -0.494140625, + -0.65625, + -0.7578125, + -0.9375, + -1.1328125, + -1.0, + -1.0625, + -0.06787109375, + -0.4140625, + -0.130859375, + -0.26953125, + -0.375, + -0.515625, + -0.57421875, + -0.58203125, + -0.6484375, + -0.7265625, + -0.73046875, + -0.83203125, + -0.85546875, + -0.9609375, + -1.078125, + -0.498046875, + 1.6015625, + 1.6875, + 1.84375, + 1.7890625, + 1.75, + 1.703125, + 1.6640625, + 1.6328125, + 1.609375, + 1.5625, + 1.515625, + 1.453125, + 1.4140625, + 1.2265625, + 1.3046875, + 0.76171875, + -0.09423828125, + -1.90625, + -1.296875, + -1.1875, + -0.9609375, + -0.7578125, + -0.6640625, + -0.5703125, + -0.49609375, + -0.384765625, + -0.279296875, + -0.1357421875, + 0.021240234375, + 0.25, + 0.41796875, + 0.2001953125, + 0.8359375, + 0.5078125, + 0.609375, + 0.333984375, + 0.16015625, + 0.06396484375, + -0.0089111328125, + -0.047119140625, + -0.10546875, + -0.1708984375, + -0.2421875, + -0.34375, + -0.484375, + -0.734375, + -0.86328125, + -1.546875, + 0.0712890625, + 0.0262451171875, + 0.2099609375, + -0.0537109375, + -0.2890625, + -0.48046875, + -0.6015625, + -0.7265625, + -0.8125, + -0.97265625, + -1.09375, + -1.28125, + -1.4453125, + -1.65625, + -1.765625, + -1.8984375, +] +LTX25_AUDIO_LATENTS_STD = [ + 1.875, + 2.046875, + 2.03125, + 2.03125, + 1.9765625, + 1.953125, + 1.9453125, + 1.9140625, + 1.875, + 1.8515625, + 1.8359375, + 1.8125, + 1.796875, + 1.796875, + 1.859375, + 1.65625, + 1.0390625, + 1.1796875, + 1.171875, + 1.1328125, + 1.1015625, + 1.0625, + 1.03125, + 1.0, + 0.96875, + 0.9296875, + 0.91796875, + 0.89453125, + 0.875, + 0.85546875, + 0.80078125, + 0.77734375, + 1.3046875, + 1.3203125, + 1.296875, + 1.2265625, + 1.171875, + 1.125, + 1.1015625, + 1.0625, + 1.015625, + 0.99609375, + 0.98046875, + 0.96875, + 0.94140625, + 0.9453125, + 0.984375, + 1.03125, + 1.3984375, + 1.171875, + 1.125, + 1.0625, + 1.015625, + 0.97265625, + 0.9296875, + 0.91796875, + 0.8828125, + 0.8359375, + 0.8359375, + 0.8203125, + 0.80859375, + 0.79296875, + 0.7734375, + 0.8828125, + 0.76171875, + 1.078125, + 1.046875, + 1.0546875, + 1.03125, + 1.015625, + 0.98046875, + 0.94921875, + 0.90234375, + 0.859375, + 0.828125, + 0.79296875, + 0.765625, + 0.75, + 0.75390625, + 0.7421875, + 1.2421875, + 1.328125, + 1.25, + 1.25, + 1.21875, + 1.1796875, + 1.1796875, + 1.15625, + 1.125, + 1.109375, + 1.09375, + 1.0703125, + 1.0703125, + 1.046875, + 1.0703125, + 1.125, + 1.390625, + 1.234375, + 1.21875, + 1.2109375, + 1.15625, + 1.140625, + 1.109375, + 1.1171875, + 1.09375, + 1.0703125, + 1.0625, + 1.046875, + 1.03125, + 1.0390625, + 1.09375, + 1.359375, + 1.15625, + 1.5234375, + 1.4453125, + 1.4296875, + 1.3828125, + 1.3359375, + 1.3203125, + 1.2734375, + 1.2421875, + 1.2109375, + 1.1875, + 1.15625, + 1.1484375, + 1.1328125, + 1.140625, + 0.99609375, +] diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 3aa2c854dfe6..737b1e897b51 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -154,6 +154,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("ltx", _create_default_map_fn("LTXModularPipeline")), ("ltx2", _create_default_map_fn("LTX2ModularPipeline")), ("ltx2.5", _create_default_map_fn("LTX25ModularPipeline")), + ("ltx2.5-two-stage", _create_default_map_fn("LTX25TwoStageModularPipeline")), ("minimax-h3", _create_default_map_fn("MiniMaxH3ModularPipeline")), ("minimax-music3", _create_default_map_fn("MiniMaxMusic3ModularPipeline")), ("ernie-image", _create_default_map_fn("ErnieImageModularPipeline")), diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 376596d632ea..9620d606b03a 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -542,6 +542,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class LTX25TwoStageBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class LTX25TwoStageModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class LTXAutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py index 2ee330ece20a..2e2a811354b5 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2.py @@ -47,6 +47,7 @@ ("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2PrepareCoordsStep"), ("denoise.denoise", "LTX2DenoiseStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2VaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -62,6 +63,7 @@ ("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2PrepareCoordsStep"), ("denoise.denoise", "LTX2Image2VideoDenoiseStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2VaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -76,7 +78,8 @@ ("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"), ("denoise.denoise", "LTX2ConditionDenoiseStep"), - ("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2VaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -91,7 +94,8 @@ ("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"), ("denoise.denoise", "LTX2ConditionDenoiseStep"), - ("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2VaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -105,7 +109,7 @@ class LTX2ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_blocks_class = LTX2AutoBlocks pretrained_model_name_or_path = LTX2_REPO_ID batch_params = frozenset(["prompt"]) - optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) + optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt"]) expected_workflow_blocks = LTX2_WORKFLOWS output_name = "videos" diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py index 69080095c627..9d3f1c40dc84 100644 --- a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25.py @@ -18,7 +18,7 @@ import pytest import torch -from diffusers.modular_pipelines import LTX2AutoBlocks, LTX25AutoBlocks, LTX25ModularPipeline +from diffusers.modular_pipelines import LTX25AutoBlocks, LTX25ModularPipeline from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition from diffusers.pipelines.ltx2.pipeline_ltx2_ic_lora import LTX2ReferenceCondition @@ -38,12 +38,13 @@ ("text_encoder.text_encoder", "LTX2TextEncoderStep"), ("text_encoder.connectors", "LTX2TextConnectorStep"), ("duration", "LTX2DurationStep"), - ("denoise.input", "LTX2TextInputStep"), + ("input", "LTX2TextInputStep"), ("denoise.set_timesteps", "LTX2SetTimestepsStep"), ("denoise.prepare_latents", "LTX2PrepareLatentsStep"), ("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2PrepareCoordsStep"), ("denoise.denoise", "LTX2DenoiseStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2DiffusionVaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -52,13 +53,14 @@ ("text_encoder.connectors", "LTX2TextConnectorStep"), ("duration", "LTX2DurationStep"), ("vae_encoder", "LTX2VaeEncoderStep"), - ("denoise.input", "LTX2TextInputStep"), + ("input", "LTX2TextInputStep"), ("denoise.set_timesteps", "LTX2SetTimestepsStep"), ("denoise.prepare_latents", "LTX2PrepareLatentsStep"), ("denoise.prepare_i2v_latents", "LTX2Image2VideoPrepareLatentsStep"), ("denoise.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2PrepareCoordsStep"), ("denoise.denoise", "LTX2Image2VideoDenoiseStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2DiffusionVaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -67,13 +69,14 @@ ("text_encoder.connectors", "LTX2TextConnectorStep"), ("duration", "LTX2DurationStep"), ("condition_encoder", "LTX2ConditionEncoderStep"), - ("denoise.input", "LTX2TextInputStep"), + ("input", "LTX2TextInputStep"), ("denoise.prepare_latents", "LTX2ConditionPrepareLatentsStep"), ("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"), ("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"), ("denoise.denoise", "LTX2ConditionDenoiseStep"), - ("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2DiffusionVaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -82,13 +85,14 @@ ("text_encoder.connectors", "LTX2TextConnectorStep"), ("condition_encoder", "LTX2ConditionEncoderStep"), ("reference_encoder", "LTX2ReferenceEncoderStep"), - ("denoise.input", "LTX2TextInputStep"), + ("input", "LTX2TextInputStep"), ("denoise.prepare_latents", "LTX2InContextPrepareLatentsStep"), ("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"), ("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"), ("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"), ("denoise.denoise", "LTX2ConditionDenoiseStep"), - ("decode.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.trim_condition_tokens", "LTX2TrimConditionTokensStep"), + ("denoise.unpack", "LTX2UnpackLatentsStep"), ("decode.video_decode", "LTX2DiffusionVaeDecoderStep"), ("decode.audio_decode", "LTX2AudioDecoderStep"), ], @@ -102,7 +106,7 @@ class LTX25ModularPipelineTesterConfig(BaseModularPipelineTesterConfig): pipeline_blocks_class = LTX25AutoBlocks pretrained_model_name_or_path = LTX25_REPO_ID batch_params = frozenset(["prompt"]) - optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) + optional_params = frozenset(["num_videos_per_prompt"]) expected_workflow_blocks = LTX25_WORKFLOWS output_name = "videos" @@ -111,7 +115,7 @@ def get_dummy_inputs(self, seed=0): "prompt": "a robot dancing", "negative_prompt": "", "generator": self.get_generator(seed), - "num_inference_steps": 2, + "sigmas": [1.0, 0.5], "height": 32, "width": 32, "num_frames": 5, @@ -151,24 +155,6 @@ def test_diffusion_decoder_output(self): assert audio.shape[1] == pipe.vocoder.config.out_channels assert torch.isnan(audio).sum() == 0 - def test_graph_matches_ltx2_except_video_decode(self): - # `LTX25AutoBlocks` restates the `LTX2AutoBlocks` graph rather than subclassing it, so nothing but this - # test keeps the two in step: a stage added to one and not the other passes both blocksets' own - # `expected_workflow_blocks`. - ltx2_blocks, ltx25_blocks = LTX2AutoBlocks(), LTX25AutoBlocks() - assert ltx2_blocks.available_workflows == ltx25_blocks.available_workflows - - for workflow_name in ltx25_blocks.available_workflows: - expected = [ - (name, "LTX2DiffusionVaeDecoderStep" if name == "decode.video_decode" else type(block).__name__) - for name, block in ltx2_blocks.get_workflow(workflow_name).sub_blocks.items() - ] - actual = [ - (name, type(block).__name__) - for name, block in ltx25_blocks.get_workflow(workflow_name).sub_blocks.items() - ] - assert actual == expected, f"Workflow '{workflow_name}' diverges from `LTX2AutoBlocks`" - def test_auto_duration_predicts_a_grid_valid_frame_count(self): pipe = self.get_pipeline().to("cpu") diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25_two_stage.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25_two_stage.py new file mode 100644 index 000000000000..71c14387804c --- /dev/null +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx25_two_stage.py @@ -0,0 +1,250 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import pytest +import torch + +from diffusers.modular_pipelines import LTX25TwoStageBlocks, LTX25TwoStageModularPipeline +from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel + +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularLoadingTesterMixin, + ModularPipelineTesterMixin, + ModularWorkflowTesterMixin, +) + + +LTX25_REPO_ID = "hf-internal-testing/tiny-ltx2-5-modular-pipe" + +LTX25_TWO_STAGE_WORKFLOWS = { + "text2video": [ + ("text_encoder.text_encoder", "LTX2TextEncoderStep"), + ("text_encoder.connectors", "LTX2TextConnectorStep"), + ("duration", "LTX2DurationStep"), + ("input", "LTX2TextInputStep"), + ("stage_1.set_timesteps", "LTX2SetTimestepsStep"), + ("stage_1.prepare_latents", "LTX2PrepareLatentsStep"), + ("stage_1.prepare_audio_latents", "LTX2PrepareAudioLatentsStep"), + ("stage_1.prepare_coords", "LTX2PrepareCoordsStep"), + ("stage_1.denoise", "LTX2DenoiseStep"), + ("stage_1.unpack", "LTX2UnpackLatentsStep"), + ("upsample.latent_upsample", "LTX2LatentUpsampleStep"), + ("stage_2.prepare_latents", "LTX2Stage2PrepareLatentsStep"), + ("stage_2.set_timesteps", "LTX2SetTimestepsStep"), + ("stage_2.prepare_audio_latents", "LTX2Stage2PrepareAudioLatentsStep"), + ("stage_2.prepare_coords", "LTX2PrepareCoordsStep"), + ("stage_2.denoise", "LTX2DenoiseStep"), + ("stage_2.unpack", "LTX2UnpackLatentsStep"), + ("decode.video_decode", "LTX2DiffusionVaeDecoderStep"), + ("decode.audio_decode", "LTX2AudioDecoderStep"), + ], +} + + +class LTX25TwoStageModularPipelineTesterConfig(BaseModularPipelineTesterConfig): + pipeline_class = LTX25TwoStageModularPipeline + pipeline_blocks_class = LTX25TwoStageBlocks + pretrained_model_name_or_path = LTX25_REPO_ID + params = frozenset(["prompt", "height", "width", "num_frames"]) + batch_params = frozenset(["prompt"]) + optional_params = frozenset(["num_videos_per_prompt"]) + # Each pass runs its own fixed sigma schedule, so there is no step count to set; and the first pass always + # starts from noise, so there are no pre-generated `latents` to pass -- the upsample step produces them. + not_params = frozenset(["num_inference_steps", "latents"]) + expected_workflow_blocks = LTX25_TWO_STAGE_WORKFLOWS + output_name = "videos" + + def get_pipeline(self, components_manager=None, dtype=torch.float32): + # The LTX-2.5 fixture ships no `latent_upsampler`, so the stage bridge is a seeded tiny one built here. Its + # weights are fixed by the seed, which keeps outputs comparable across the pipelines a test class builds. + pipe = super().get_pipeline(components_manager=components_manager, dtype=dtype) + torch.manual_seed(0) + latent_upsampler = LTX2LatentUpsamplerModel( + in_channels=pipe.transformer.config.in_channels, mid_channels=32, num_blocks_per_stage=1 + ) + pipe.update_components(latent_upsampler=latent_upsampler.to(dtype)) + return pipe + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "a robot dancing", + "negative_prompt": "", + "generator": self.get_generator(seed), + "sigmas": [1.0, 0.5], + "stage_2_sigmas": [0.5, 0.25], + "height": 32, + "width": 32, + "num_frames": 5, + "frame_rate": 25.0, + "max_sequence_length": 16, + "output_type": "pt", + } + + +class TestLTX25TwoStageModularPipelineFast(LTX25TwoStageModularPipelineTesterConfig, ModularPipelineTesterMixin): + @pytest.mark.skip(reason="num_videos_per_prompt") + def test_num_images_per_prompt(self): + pass + + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(expected_max_diff=1e-3) + + def test_output_is_twice_the_requested_resolution(self): + # As with the standard pipelines, `height` / `width` are the first pass's resolution and the upsample + # doubles them. + pipe = self.get_pipeline().to("cpu") + + output = pipe(**self.get_dummy_inputs(), output=["videos", "audio"]) + videos, audio = output["videos"], output["audio"] + + assert videos.shape == (1, 5, 3, 64, 64) + assert audio.shape[0] == 1 + assert torch.isnan(audio).sum() == 0 + + def test_first_pass_decodes_on_its_own(self): + # Popping `upsample` and `stage_2` leaves `stage_1 -> decode`: a preview of the first pass at the requested + # resolution, since the pass leaves packed latents in state like any other core denoise. + pipe = self.get_pipeline().to("cpu") + blocks = LTX25TwoStageBlocks() + blocks.sub_blocks.pop("upsample") + blocks.sub_blocks.pop("stage_2") + preview_pipe = blocks.init_pipeline(LTX25_REPO_ID) + preview_pipe.update_components(**{name: getattr(pipe, name) for name in pipe.pretrained_component_names}) + + videos = preview_pipe(**self.get_dummy_inputs(), output="videos") + + assert videos.shape == (1, 5, 3, 32, 32) + + def test_stage_2_follows_the_workflow(self): + # `stage_2` picks the workflow's second-pass group, and the image / frame conditions are re-encoded at the + # upsampled resolution ahead of it, so an image-to-video or condition run refines under its conditioning. + blocks = LTX25TwoStageBlocks() + + def selected(**inputs): + names = { + name: type(block).__name__ for name, block in blocks.get_execution_blocks(**inputs).sub_blocks.items() + } + return ( + names["stage_2.prepare_latents"], + "stage_2.prepare_i2v_latents" in names, + "stage_2_vae_encoder" in names, + "stage_2_condition_encoder" in names, + ) + + assert selected(prompt=True) == ("LTX2Stage2PrepareLatentsStep", False, False, False) + assert selected(prompt=True, image=True, image_latents=True) == ( + "LTX2Stage2PrepareLatentsStep", + True, + True, + False, + ) + assert selected(prompt=True, conditions=True, condition_latents=True) == ( + "LTX2ConditionStage2PrepareLatentsStep", + False, + False, + True, + ) + + def test_split_stages_match_the_single_call(self): + # The point of the blockset: `stage_1` / `upsample` / `stage_2` are each usable as their own pipeline, and + # chaining them by hand -- one generator threaded through, as in the standard two-stage recipe -- is the + # same computation as the one call. + pipe = self.get_pipeline().to("cpu") + reference = pipe(**self.get_dummy_inputs(), output="videos") + + blocks = LTX25TwoStageBlocks() + stage_2 = blocks.sub_blocks.pop("stage_2") + upsample = blocks.sub_blocks.pop("upsample") + decode = blocks.sub_blocks.pop("decode") + components = {name: getattr(pipe, name) for name in pipe.pretrained_component_names} + # The stage blocks declare no autoencoder, so a standalone stage pipeline resolves the latent geometry and + # statistics to the `LTX2ModularPipeline` fallbacks -- the production values, which the tiny fixture VAEs do + # not have. Pin the fixture's values on those pipelines; a real run has the checkpoint's. + vae_values = { + name: getattr(pipe, name) + for name in ( + "vae_spatial_compression_ratio", + "vae_temporal_compression_ratio", + "vae_scaling_factor", + "latents_mean", + "latents_std", + "audio_latent_channels", + "audio_latent_mel_bins", + "audio_vae_mel_compression_ratio", + "audio_vae_temporal_compression_ratio", + "audio_sampling_rate", + "audio_hop_length", + "audio_latents_mean", + "audio_latents_std", + ) + } + pipes = [] + for stage_blocks in (blocks, upsample, stage_2, decode): + stage_pipe = stage_blocks.init_pipeline(LTX25_REPO_ID) + stage_pipe.update_components(**{k: v for k, v in components.items() if k in stage_pipe.components}) + if "diffusion_decoder" not in stage_pipe.components: + stage_pipe.__class__ = type( + type(stage_pipe).__name__, + (type(stage_pipe),), + {name: property(lambda self, value=value: value) for name, value in vae_values.items()}, + ) + pipes.append(stage_pipe) + stage_1_pipe, upsample_pipe, stage_2_pipe, decode_pipe = pipes + + inputs = self.get_dummy_inputs() + + def carry(to_pipe, *from_states): + # Hand a downstream pipeline whatever it declares as an input, from the upstream states (later wins). + carried = {} + for from_state in from_states: + carried.update( + { + name: from_state.get(name) + for name in to_pipe.blocks.input_names + if from_state.get(name) is not None + } + ) + return carried + + stage_1_state = stage_1_pipe(**inputs) + upsample_state = upsample_pipe(**carry(upsample_pipe, stage_1_state)) + stage_2_inputs = carry(stage_2_pipe, stage_1_state, upsample_state) + stage_2_inputs.update(stage_2_sigmas=inputs["stage_2_sigmas"], generator=inputs["generator"]) + stage_2_state = stage_2_pipe(**stage_2_inputs) + decode_inputs = carry(decode_pipe, stage_2_state) + decode_inputs.update(generator=inputs["generator"], output_type=inputs["output_type"]) + videos = decode_pipe(**decode_inputs, output="videos") + + assert torch.allclose(videos, reference, atol=1e-4) + + +class TestLTX25TwoStageModularPipelineLoading(LTX25TwoStageModularPipelineTesterConfig, ModularLoadingTesterMixin): + @pytest.mark.skip(reason="the fixture ships no `latent_upsampler`, so a reloaded pipeline cannot run") + def test_save_from_pretrained(self): + pass + + +class TestLTX25TwoStageModularPipelineWorkflow(LTX25TwoStageModularPipelineTesterConfig, ModularWorkflowTesterMixin): + # `ModularPipeline.from_pretrained(repo, workflow=...)` resolves the blocks from the repo's + # `modular_model_index.json`, which for the LTX-2.5 fixture is `LTX25AutoBlocks`. These two tests need a fixture + # whose index names `LTX25TwoStageBlocks` and ships a `latent_upsampler`. + @pytest.mark.skip(reason="the fixture routes to `LTX25AutoBlocks`") + def test_from_pretrained_workflow(self): + pass + + @pytest.mark.skip(reason="the fixture routes to `LTX25AutoBlocks`") + def test_load_components_workflow(self): + pass