[LTX-2.5] Refactor LTX-2.5 Diffusion Decoder Forward Methods

docs/source/en/api/models/ltx2_diffusion_decoder.md

@@ -19,17 +19,31 @@ consumes unchanged, so latents are interchangeable between the convolutional dec
 itself a diffusion model it is driven by [`LTX2VideoDiffusionDecodePipeline`] rather than being passed as a
 pipeline's `vae`: run any LTX-2 pipeline with `output_type="latent"`, then decode.
 
+`forward` is a single denoising step, like any other denoiser in the library. The loop over steps, the scheduler
+that drives it, and the tiling wrapped around it live in [`LTX2VideoDiffusionDecodePipeline`], so this model is not
+called directly in normal use — the two `encode_context_*` methods build the conditioning the step consumes.
+
 ```python
 import torch
-from diffusers import LTX2Pipeline, LTX2VideoDiffusionDecodePipeline, LTX2VideoDiffusionDecoderModel
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    LTX2Pipeline,
+    LTX2VideoDiffusionDecodePipeline,
+    LTX2VideoDiffusionDecoderModel,
+)
 
 pipe = LTX2Pipeline.from_pretrained("Lightricks/LTX-2.5-Diffusers", dtype=torch.bfloat16).to("cuda")  # or "mps", "xpu", "cpu"
 latents = pipe(prompt="a potter shaping a clay vase", output_type="latent").frames
 
 decoder = LTX2VideoDiffusionDecoderModel.from_pretrained(
     "Lightricks/LTX-2.5-Diffusers", subfolder="diffusion_decoder", dtype=torch.bfloat16
 ).to("cuda")
-decode_pipe = LTX2VideoDiffusionDecodePipeline(diffusion_decoder=decoder, scheduler=pipe.scheduler)
+# The decoder's own scheduler, not `pipe.scheduler`: the transformer's is resolution-shifted and would need
+# a `mu` this pipeline does not compute.
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    "Lightricks/LTX-2.5-Diffusers", subfolder="diffusion_decoder_scheduler"
+)
+decode_pipe = LTX2VideoDiffusionDecodePipeline(diffusion_decoder=decoder, scheduler=scheduler)
 
 # `denormalize=False`: `output_type="latent"` already applied the latent statistics, so applying them
 # again here would scale every channel by its std a second time.
@@ -71,18 +85,20 @@ accepts the `BlockMask`. Use the NATTEN processor above instead.
 
 ## Tiling
 
-`decoder.enable_tiling()` decodes in overlapping tiles that are blended back together, bounding peak memory by the
-tile size instead of the video size. The cheap early upsampling stages still see the full latent — only the last
-upsampling stage and the diffusion stage, which dominate decode memory, run per tile — so tiling changes the output
-only near tile borders. Because the diffusion stage denoises each tile separately, a tiled decode does not
-reproduce the untiled result exactly; the default tile and overlap sizes match the reference implementation's.
-Neighborhood attention rejects any grid smaller than its kernel, so a trailing remnant tile is merged into its
-neighbor rather than decoded on its own.
+`decoder.enable_tiling()` decodes in overlapping tiles that are blended back together, bounding peak memory by the tile
+size instead of the video size. It sets the tile sizes; the decoder works out where the cuts fall, and walking them is
+[`LTX2VideoDiffusionDecodePipeline`]'s job, because tiles are cut in the *middle* of the decoder and each one runs its
+own denoising loop. Tiling engages only once the video exceeds one tile. The cheap early upsampling stages still see the
+full latent — only the last upsampling stage and the diffusion stage, which dominate decode memory, run per tile — so
+tiling changes the output only near tile borders. Because the diffusion stage denoises each tile separately, a tiled
+decode does not reproduce the untiled result exactly; the default tile and overlap sizes match the reference
+implementation's. Neighborhood attention rejects any grid smaller than its kernel, so a trailing remnant tile is merged
+into its neighbor rather than decoded on its own.
 
 ## LTX2VideoDiffusionDecoderModel
 
 [[autodoc]] LTX2VideoDiffusionDecoderModel
-    - decode
+    - forward
     - enable_tiling
     - disable_tiling
     - all

docs/source/en/api/pipelines/ltx2.md

@@ -680,15 +680,20 @@ The upsample step and the stage 2 call itself are unchanged from the distilled r
 LTX-2.5 ships two video decoders over the same latent space, so latents are interchangeable between them:
 
 - `vae/` — the convolutional VAE ([`AutoencoderKLLTX2Video`]). It is what the pipelines decode with, so every snippet above already uses it, and it is the only one of the two that tiles (`pipe.vae.enable_tiling()`), which is usually what makes a high resolution fit.
-- `diffusion_decoder/` — [`LTX2VideoDiffusionDecoderModel`]. It is a diffusion model in its own right rather than a pipeline component, so it is not passed as a `vae`: run the pipeline with `output_type="latent"` and hand the latents to [`LTX2VideoDiffusionDecodePipeline`].
+- `diffusion_decoder/` — [`LTX2VideoDiffusionDecoderModel`]. It is a diffusion model in its own right rather than a pipeline component, so it is not passed as a `vae`: run the pipeline with `output_type="latent"` and hand the latents to [`LTX2VideoDiffusionDecodePipeline`]. Because it denoises it needs a scheduler of its own, kept in `diffusion_decoder_scheduler/` — the repo's top-level `scheduler/` is the transformer's and is shifted for the transformer's sequence lengths.
 
 Encoding always goes through `vae/`, so image and video conditioning are unaffected by the choice.
 
 Two things change when you decode with the diffusion decoder. `output_type="latent"` also skips the vocoder, so the audio comes back as latents and has to be finished by hand, and the NATTEN processor is effectively required at video resolutions:
 
 ```py
 import torch
-from diffusers import LTX2Pipeline, LTX2VideoDiffusionDecodePipeline, LTX2VideoDiffusionDecoderModel
+from diffusers import (
+    FlowMatchEulerDiscreteScheduler,
+    LTX2Pipeline,
+    LTX2VideoDiffusionDecodePipeline,
+    LTX2VideoDiffusionDecoderModel,
+)
 from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoVaeNeighborhoodNattenProcessor
 from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES
 from diffusers.utils import encode_video
@@ -733,7 +738,10 @@ decoder.set_attn_processor(LTX2VideoVaeNeighborhoodNattenProcessor())
 # Decode in overlapping tiles so peak memory scales with the tile size rather than the video size.
 decoder.enable_tiling()
 
-decode_pipe = LTX2VideoDiffusionDecodePipeline(diffusion_decoder=decoder, scheduler=pipe.scheduler)
+# The decoder's own scheduler, not `pipe.scheduler`: the transformer's is resolution-shifted and would need
+# a `mu` this pipeline does not compute.
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(model_path, subfolder="diffusion_decoder_scheduler")
+decode_pipe = LTX2VideoDiffusionDecodePipeline(diffusion_decoder=decoder, scheduler=scheduler)
 
 # `denormalize=False`: `output_type="latent"` already applied the latent statistics, so applying them
 # again would rescale every channel by its std a second time. The decoder draws the noise it denoises,
@@ -753,7 +761,9 @@ encode_video(
 
 To combine this with [two-stage generation](#two-stage-generation-for-ltx-25), ask *stage 2* for `output_type="latent"` and decode that.
 
-`decoder.enable_tiling()` is what keeps a high resolution decode in memory, the same way `pipe.vae.enable_tiling()` does for the convolutional VAE. The memory-dominant part of the decode — the last upsampling stage and the diffusion stage — then runs on overlapping tiles that are blended back together, so peak memory is bounded by the tile size instead of the video size. Tiling only kicks in once the latent exceeds one tile, and the tile and overlap sizes can be tuned via the `tile_sample_min_*` / `tile_sample_stride_*` arguments (defaults match the reference implementation). Since the diffusion stage denoises each tile separately, a tiled decode does not reproduce the untiled result exactly.
+`decoder.enable_tiling()` is what keeps a high resolution decode in memory, the same way `pipe.vae.enable_tiling()` does for the convolutional VAE. The memory-dominant part of the decode — the last upsampling stage and the diffusion stage — then runs on overlapping tiles that are blended back together, so peak memory is bounded by the tile size instead of the video size. Tiling only kicks in once the latent exceeds one tile, and the tile and overlap sizes can be tuned via the `tile_sample_min_*` / `tile_sample_stride_*` arguments (defaults match the reference implementation). Since the diffusion stage denoises each tile separately, a tiled decode does not reproduce the untiled result exactly. The call sets the sizes and [`LTX2VideoDiffusionDecodePipeline`] does the tiling, since each tile runs its own denoising loop — unlike the convolutional VAE, where the tiling is entirely inside the model.
+
+The decode is a denoising loop like any other, so `num_inference_steps` (or an explicit `sigmas` schedule) is a `__call__` argument. It defaults to what the checkpoint was distilled for, which is a single step on LTX-2.5; more steps cost proportionally more time and are rarely worth it on a distilled decoder. The scheduler is a real component rather than a formality — reshaping the schedule through its config (a `shift`, say) reaches the decode, which is the point of driving the loop from one. The only setting the pipeline cannot honour is `use_dynamic_shifting`, since it never computes a `mu`; that is why a transformer's scheduler cannot be reused here.
 
 On a single card it is also worth moving the pipeline out of the way before decoding (`pipe.to("cpu")` and `torch.cuda.empty_cache()`, after capturing `pipe.scheduler` and the vocoder's `output_sampling_rate`), since the decoder needs its own headroom. See [`LTX2VideoDiffusionDecoderModel`] for the attention backends, the tiling details, and the rest of the decoder's behaviour.
 
@@ -1247,14 +1257,18 @@ The two axes are seamed differently. Neither side of a spatial border holds a kn
 **Decoding with the diffusion decoder.** For maximum detail fidelity, stay on `output_type="latent"` and hand the (already denormalized, possibly `trim_canvas`'d) latents to [`LTX2VideoDiffusionDecodePipeline`].
 
 ```py
-from diffusers import LTX2VideoDiffusionDecodePipeline
+from diffusers import FlowMatchEulerDiscreteScheduler, LTX2VideoDiffusionDecodePipeline
 from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel
 
 decoder = LTX2VideoDiffusionDecoderModel.from_pretrained(
     "Lightricks/LTX-2.5-Diffusers", subfolder="diffusion_decoder", dtype=torch.bfloat16
 )
+# The decoder's own scheduler, not `pipe.scheduler`, which this pipeline cannot drive.
+scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained(
+    "Lightricks/LTX-2.5-Diffusers", subfolder="diffusion_decoder_scheduler"
+)
 decode_pipe = LTX2VideoDiffusionDecodePipeline(
-    diffusion_decoder=decoder, scheduler=pipe.scheduler, vae=pipe.vae
+    diffusion_decoder=decoder, scheduler=scheduler, vae=pipe.vae
 )
 decode_pipe.enable_model_cpu_offload()
 # `denormalize=False`: the `output_type="latent"` path already applied the latent statistics.

scripts/convert_ltx2_to_diffusers.py

@@ -892,6 +892,23 @@ def get_ltx2_diffusion_video_vae_config(version: str) -> tuple[dict[str, Any], d
     return config, LTX_2_3_VIDEO_VAE_RENAME_DICT, LTX_2_0_VAE_SPECIAL_KEYS_REMAP
 
 
+def get_ltx2_diffusion_decoder_scheduler() -> FlowMatchEulerDiscreteScheduler:
+    """The scheduler for the LTX-2.5 diffusion decoder's denoising loop.
+
+    This is deliberately *not* the transformer's scheduler. The decoder walks a plain uniform schedule,
+    `linspace(1, 1 / num_inference_steps, num_inference_steps)`, so every knob that bends the sigmas has to be off:
+    resolution-dependent shifting expects a `mu` the decoder has no sequence length to derive, and a terminal shift
+    would end the walk at `shift_terminal` instead of the last sigma the decoder was distilled on.
+    """
+    return FlowMatchEulerDiscreteScheduler(
+        num_train_timesteps=1000,
+        shift=1.0,
+        use_dynamic_shifting=False,
+        shift_terminal=None,
+        stochastic_sampling=False,
+    )
+
+
 def convert_ltx2_diffusion_video_vae(original_state_dict: dict[str, Any], version: str) -> dict[str, Any]:
     config, rename_dict, special_keys_remap = get_ltx2_diffusion_video_vae_config(version)
     diffusers_config = config["diffusers_config"]
@@ -1414,7 +1431,8 @@ def none_or_str(value: str):
         help=(
             "Whether to convert the LTX-2.5 diffusion decoder, saved to a `diffusion_decoder` subfolder — the "
             "component name `LTX2VideoDiffusionDecodePipeline` and the modular blocks resolve it by — so "
-            "`from_pretrained` keeps returning the conv decoder in `vae` by default"
+            "`from_pretrained` keeps returning the conv decoder in `vae` by default. Its denoising schedule "
+            "differs from the transformer's, so a `diffusion_decoder_scheduler` subfolder is written too"
         ),
     )
     parser.add_argument("--audio_vae", action="store_true", help="Whether to convert the audio VAE model")
@@ -1551,6 +1569,12 @@ def main(args):
         # from the subfolder named after it, so this folder name must match the `diffusion_decoder` component
         # of `LTX2VideoDiffusionDecodePipeline` (and the modular `ComponentSpec`) for those loads to work.
         diffusion_vae.to(vae_dtype).save_pretrained(os.path.join(args.output_path, "diffusion_decoder"))
+        # The decoder denoises, so it needs a scheduler, and the repo's top-level `scheduler` is the transformer's:
+        # it has `use_dynamic_shifting` on, which the decoder cannot satisfy. Save the decoder's own alongside the
+        # weights so `LTX2VideoDiffusionDecodePipeline` can be assembled entirely from this repo.
+        get_ltx2_diffusion_decoder_scheduler().save_pretrained(
+            os.path.join(args.output_path, "diffusion_decoder_scheduler")
+        )
 
     if args.audio_vae or args.full_pipeline:
         if args.audio_vae_filename is not None: