Skip to content

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

Open
yiyixuxu wants to merge 6 commits into
mainfrom
modular-ltx25-two-stage
Open

[modular] LTX-2.5: two-stage generation as one pipeline#14612
yiyixuxu wants to merge 6 commits into
mainfrom
modular-ltx25-two-stage

Conversation

@yiyixuxu

@yiyixuxuyiyixuxu commented Aug 27, 2026

Copy link
Copy Markdown
Collaborator

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

importtorchfromdiffusersimportComponentsManager, ModularPipelinefromdiffusers.modular_pipelinesimportLTX25TwoStageBlocksfromdiffusers.models.autoencoders.ltx2_diffusion_decoderimportLTX2VideoVaeNeighborhoodNattenProcessorfromdiffusers.pipelines.ltx2.latent_upsamplerimportLTX2LatentUpsamplerModelfromdiffusers.utilsimportencode_videodevice="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()
defsave(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 processorpipe.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.

inputdefault
prompt(required)
negative_prompt, max_sequence_lengthNone, 1024
num_frames, min_seconds, max_seconds, frame_rateNone (auto), 1.0, 20.0, 24.0
sigmas, timestepsDISTILLED_SIGMA_VALUES, None
height, width512, 704
num_videos_per_prompt, generator, attention_kwargs, output_type1, 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 explicitlypipe.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 1408x1024state=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:

inputdefault
prompt(required)
negative_prompt, max_sequence_lengthNone, 1024
num_frames, min_seconds, max_seconds, frame_rateNone (auto), 1.0, 20.0, 24.0
sigmas, timestepsDISTILLED_SIGMA_VALUES, None
height, width512, 704 (the first pass; the output is 2x)
stage_2_sigmas, stage_2_timestepsSTAGE_2_DISTILLED_SIGMA_VALUES, None
noise_scaleNonestage_2_sigmas[0], the level the upsampled latents are re-noised to
num_videos_per_prompt, generator, attention_kwargs, output_type1, 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 managerstage_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:

inputdefault
latents, audio_latentsfrom stage_1 (latents through upsample)(required)
connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_maskthe text conditioning from text_encoder + input(required)
negative_connector_* (3)only under classifier-free guidanceNone
batch_size, dtypefrom input(required)
stage_2_sigmas, stage_2_timestepsSTAGE_2_DISTILLED_SIGMA_VALUES, None
noise_scaleNonestage_2_sigmas[0]
frame_rate, num_videos_per_prompt, generator, attention_kwargs24.0, 1, None, None
image_latents / condition_latents, condition_strengths, condition_indices, condition_pixel_framesi2v / condition workflows only, from the stage_2_* encodersNone

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

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 <noreply@anthropic.com>
@github-actionsgithub-actionsBot added size/L PR with diff > 200 LOC documentation Improvements or additions to documentation tests modular-pipelines utils and removed size/L PR with diff > 200 LOC labels Aug 27, 2026
@github-actionsgithub-actionsBot added the size/L PR with diff > 200 LOC label Sep 6, 2026
@HuggingFaceDocBuilderDev

Copy link
Copy Markdown

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

yiyixuxuand others added 2 commits September 8, 2026 02:28
…ted patterns
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
… the official recipe
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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 = []

@yiyixuxuyiyixuxuSep 8, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LTX2ReferenceEncoderStep did these 4 things:
(1) encode each reference -> reference latent (normalized)
(2) pack + concat
(2) compute reference_coords
(3) build reference_cross_mask

we only keep the step(1) here in encoder step, move the packing and the rest of the stuff to core denoise blocks (i.e.LTX2InContextPrepareLatentsStep)

return noise_scale * noise + (1 - noise_scale) * latents


def _downsample_mask_to_latent(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

return latents


def _normalize_audio_latents(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

We only need to accept normalized latents now - so this function is not needed here
norm/denorm should be part of vae blocks (encoders.py and decoders.py)

this changes what stage-2 accepts: it no longer takes denormalized latents (the standard pipeline output for output_type="latents")



# auto_docstring
class LTX2AutoBuildVideoSelfAttentionMaskStep(ConditionalPipelineBlocks):

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is absorbed into LTX2InContextPrepareLatentsStep

return latent_mask.reshape(b, latent_num_frames * latent_height * latent_width)


def _build_video_self_attention_mask(

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

this is the old LTX2BuildVideoSelfAttentionMaskStep + the cross-mask computation moved from reference encoder; now called from LTX2InContextPrepareLatentsStep

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

@yiyixuxuyiyixuxuSep 8, 2026

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

split LTX2PrepareLatentsStep into two blocks:

this one is only for stage1 now, so basically just generate the initial noise, "latents" input now means the pre-generated initial noise consistent with how we define this input in our other pipelines

it no longer takes the stage1 output as latents and re-noises it (that's moved to LTX2Stage2PrepareLatentsStep

yiyixuxuand others added 2 commits September 8, 2026 10:01
…latents input
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01HCdbvRpL9fv3h3WwSUPpfS
return components, state


class LTX2Stage2PrepareAudioLatentsStep(ModularPipelineBlocks):

Copy link
Copy Markdown
CollaboratorAuthor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LTX2PrepareAudioLatentsStep also split into stage1 + stage2

Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

documentationImprovements or additions to documentationmodular-pipelinessize/LPR with diff > 200 LOCtestsutils

Projects

Status: In Progress

Development

Successfully merging this pull request may close these issues.

3 participants

@yiyixuxu@HuggingFaceDocBuilderDev@DN6