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

.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
 

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

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,

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,

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