[modular] LTX-2.5: two-stage generation as one pipeline

[yiyixuxu]

Adds LTX25TwoStageBlocks the distilled two-stage recipe as a single modular pipeline:

first pass → 2x latent upsample → second pass → diffusion decode

for every workflow LTX25AutoBlocks supports (t2v / i2v / condition / in-context). It is assembled from the same leaf blocks as LTX25AutoBlocks, and the stages are ordinary blocks, so you can pop them and run a pass on its own.

setup

import torch
from diffusers import ComponentsManager, ModularPipeline
from diffusers.modular_pipelines import LTX25TwoStageBlocks
from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor
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."
cm = ComponentsManager()


def save(pipe, state, path):
    video, audio = state.get("videos"), state.get("audio")
    encode_video(
        video[0],
        fps=24.0,
        audio=audio[0].float().cpu(),
        audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
        output_path=path,
    )

example usage1: Single stage, everything at its default

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")
# NATTEN needs `pip install kernels`; omit to decode with the Flex Attention processor
pipe.diffusion_decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor())
pipe.diffusion_decoder.enable_tiling()

state = pipe(prompt=prompt, generator=torch.Generator(device).manual_seed(42), output_type="np")
save(pipe, state, "ltx2_5_single.mp4")

pipe.blocks.get_workflow("text2video").inputs — the distilled schedule is the default, num_frames is predicted by the duration head when omitted, and there is no num_inference_steps: the checkpoint runs a fixed sigma schedule, so there is no step count to choose.

input	default
prompt	(required)
negative_prompt, max_sequence_length	None, 1024
num_frames, min_seconds, max_seconds, frame_rate	None (auto), 1.0, 20.0, 24.0
sigmas, timesteps	DISTILLED_SIGMA_VALUES, None
height, width	512, 704
num_videos_per_prompt, generator, attention_kwargs, output_type	1, None, None, "pil"

example usage 2: Two stages as one call

pipe = LTX25TwoStageBlocks().init_pipeline(model_path, components_manager=cm)
pipe.load_components(dtype=torch.bfloat16)
# the repo's index does not list the upsampler yet, so load it explicitly
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")
pipe.diffusion_decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor())
pipe.diffusion_decoder.enable_tiling()

# first pass at 704x512 (the default), output at 1408x1024
state = pipe(prompt=prompt, generator=torch.Generator(device).manual_seed(42), output_type="np")
save(pipe, state, "ltx2_5_two_stage.mp4")

pipe.blocks.get_workflow("text2video").inputs — everything at its default again. The second pass reads its schedule under its own names (stage_2_*) so both passes can sit in one pipeline, and takes its height / width / num_frames from the upsampled latents rather than as inputs:

input	default
prompt	(required)
negative_prompt, max_sequence_length	None, 1024
num_frames, min_seconds, max_seconds, frame_rate	None (auto), 1.0, 20.0, 24.0
sigmas, timesteps	DISTILLED_SIGMA_VALUES, None
height, width	512, 704 (the first pass; the output is 2x)
stage_2_sigmas, stage_2_timesteps	STAGE_2_DISTILLED_SIGMA_VALUES, None
noise_scale	None → stage_2_sigmas[0], the level the upsampled latents are re-noised to
num_videos_per_prompt, generator, attention_kwargs, output_type	1, None, None, "pil"

example usage3: Two stages separately

you can pop each stage into their own pipelines and hand the state along. For instance, preview the first pass (and re-run it as many times as you like) before spending the second pass on it:

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 their components through the manager

stage_1_pipe = blocks.init_pipeline(model_path, components_manager=cm)
stage_1_pipe.load_components(dtype=torch.bfloat16)
upsample_pipe = upsample.init_pipeline(model_path, components_manager=cm)
upsample_pipe.update_components(
    latent_upsampler=LTX2LatentUpsamplerModel.from_pretrained(model_path, subfolder="latent_upsampler", dtype=torch.bfloat16)
)
stage_2_pipe = stage_2.init_pipeline(model_path, components_manager=cm)
stage_2_pipe.load_components(dtype=torch.bfloat16)
decode_pipe = decode.init_pipeline(model_path, components_manager=cm)
decode_pipe.load_components(dtype=torch.bfloat16)
cm.enable_auto_cpu_offload(device=device, memory_reserve_margin="20GB")
decode_pipe.diffusion_decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor())
decode_pipe.diffusion_decoder.enable_tiling()

# first pass, decoded as a preview (704x512). The decoder is called with the latents rather than the state so
# the preview's own generator does not replace the one `state` carries for the second pass.
generator = torch.Generator(device).manual_seed(42)
state = stage_1_pipe(prompt=prompt, generator=generator)
preview = decode_pipe(
    latents=state.get("latents"),
    audio_latents=state.get("audio_latents"),
    generator=torch.Generator(device).manual_seed(0),
    output_type="np",
)
save(decode_pipe, preview, "ltx2_5_stage_1_preview.mp4")

# happy with it: upsample and refine (1408x1024), then decode. (`upsample` and `stage_2` could also stay in one
# pipeline; they are split here only to show that each is a pipeline of its own.)
state = upsample_pipe(state=state)
state = stage_2_pipe(state=state)
save(decode_pipe, decode_pipe(state=state, output_type="np"), "ltx2_5_two_stage_split.mp4")

stage_2.inputs — what the popped second pass takes on its own (this is LTX25AutoStage2CoreDenoiseStep, so the inputs are the union of its t2v / i2v / condition branches). Everything comes from the first pass's state; its own settings are at their defaults:

input		default
latents, audio_latents	from stage_1 (latents through upsample)	(required)
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask	the text conditioning from text_encoder + input	(required)
negative_connector_* (3)	only under classifier-free guidance	None
batch_size, dtype	from input	(required)
stage_2_sigmas, stage_2_timesteps		STAGE_2_DISTILLED_SIGMA_VALUES, None
noise_scale		None → stage_2_sigmas[0]
frame_rate, num_videos_per_prompt, generator, attention_kwargs		24.0, 1, None, None
image_latents / condition_latents, condition_strengths, condition_indices, condition_pixel_frames	i2v / condition workflows only, from the stage_2_* encoders	None

No height / width / num_frames: the second pass reads them off the latents. (upsample alone takes just latents; decode takes latents, audio_latents, generator, output_type.)