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

.ai/references/modular.md

@@ -315,7 +315,15 @@ 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 possible transformations sit between a VAE and a transformer: *normalize/denormalize* (the VAE's latent statistics / `scaling_factor`) and *pack/unpack* (`[B, C, F, H, W]` <-> a token sequence `[B, S, D]`). Each must be applied and undone in mirrored pairs, and a core denoise group must hand `latents` back in the same form it received them -- so the `latents` in the state are always in one consistent form that any block (a decoder, a latent upsampler, a second denoise group) can consume, and no block outside the group needs `height` / `width` / `num_frames` just to interpret tokens. Two acceptable patterns follow this principle:
+    - **Pattern 1 (the most common one)** -- norm/denorm on the unpacked form, on the VAE blocks; pack/unpack inside the core denoise group: `encode -> norm -> pack -> denoise -> unpack -> denorm -> decode`. The pack/unpack pair lives either at block level (the prepare-latents step packs, a dedicated after-denoise step unpacks -- `QwenImageAfterDenoiseStep`, `LTX2UnpackLatentsStep`) or inside the transformer's `forward` (`wan`, `stable_diffusion_3`: the model patchifies/unpatchifies internally and blocks never pack at all).
+    - **Pattern 2 (packed-space statistics)** -- the VAE's statistics are defined over the *packed* channels, so norm/denorm can only run on packed tensors and pack/unpack move to the VAE blocks as well: `encode -> pack -> norm -> denoise -> denorm -> unpack -> decode` (`ernie_image`; `ideogram4` stays on pattern 1 instead by tiling the stats onto the unpacked channels in its decoder).
+    - A few older pipelines are inconsistent with this and are being fixed: https://github.com/huggingface/diffusers/issues/14730.
+    - 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.
 
 ## 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. The flow: a first pass at the requested `height` / `width`; a 2x latent upsample that doubles them; then — in the image-to-video and condition workflows — the image and frame conditions are re-encoded at the upsampled resolution; and a second pass refines the upsampled latents under that conditioning, re-noising them on `stage_2_sigmas` (the distilled stage-2 schedule by default) instead of sampling fresh noise. As with the standard pipelines, the resolution you pass is the first pass's, and the output is twice that size. `LTX25AutoBlocks` maps to `stage_1` only; `stage_2` is a separate auto step that selects the workflow's second-pass group, so to run two-stage, use this blockset (or add the stage-2 blocks to an assembly of your own).
+
+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`]).
 
 ### Diffusion Fidelity Rendering (DFR) for LTX-2.5
@@ -1332,3 +1360,11 @@ video = decode_pipe(latents=out3.frames, denormalize=False, output_type="np", re
 ## 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",
@@ -1403,6 +1405,8 @@
             LTX2ModularPipeline,
             LTX25AutoBlocks,
             LTX25ModularPipeline,
+            LTX25TwoStageBlocks,
+            LTX25TwoStageModularPipeline,
             LTXAutoBlocks,
             LTXModularPipeline,
             MiniMaxH3Blocks,

src/diffusers/modular_pipelines/__init__.py

@@ -126,8 +126,10 @@
     _import_structure["ltx2"] = [
         "LTX2AutoBlocks",
         "LTX25AutoBlocks",
+        "LTX25TwoStageBlocks",
         "LTX2ModularPipeline",
         "LTX25ModularPipeline",
+        "LTX25TwoStageModularPipeline",
     ]
     _import_structure["minimax_h3"] = [
         "MiniMaxH3Blocks",
@@ -191,7 +193,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