diff --git a/docs/source/en/api/models/ltx2_diffusion_decoder.md b/docs/source/en/api/models/ltx2_diffusion_decoder.md index 9a4059166267..d2b3938c6993 100644 --- a/docs/source/en/api/models/ltx2_diffusion_decoder.md +++ b/docs/source/en/api/models/ltx2_diffusion_decoder.md @@ -19,9 +19,18 @@ 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 @@ -29,7 +38,12 @@ latents = pipe(prompt="a potter shaping a clay vase", output_type="latent").fram 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 diff --git a/docs/source/en/api/pipelines/ltx2.md b/docs/source/en/api/pipelines/ltx2.md index e69b158bdf8a..3fa2e7abcd7d 100644 --- a/docs/source/en/api/pipelines/ltx2.md +++ b/docs/source/en/api/pipelines/ltx2.md @@ -680,7 +680,7 @@ 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. @@ -688,7 +688,12 @@ Two things change when you decode with the diffusion decoder. `output_type="late ```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. diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py index 33b91790ef1c..be05882968d6 100644 --- a/scripts/convert_ltx2_to_diffusers.py +++ b/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: diff --git a/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py b/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py index 41388991e0b4..28228ce89599 100644 --- a/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py +++ b/src/diffusers/models/autoencoders/ltx2_diffusion_decoder.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. import math +from dataclasses import dataclass import torch import torch.nn as nn @@ -21,12 +22,12 @@ from ...utils import is_kernels_available, logging from ...utils.accelerate_utils import apply_forward_hook from ...utils.constants import DIFFUSERS_DISABLE_REMOTE_CODE -from ...utils.torch_utils import maybe_adjust_dtype_for_device, randn_tensor +from ...utils.torch_utils import maybe_adjust_dtype_for_device from ..attention import AttentionMixin, AttentionModuleMixin from ..attention_dispatch import dispatch_attention_fn from ..embeddings import PixArtAlphaCombinedTimestepSizeEmbeddings +from ..modeling_outputs import Transformer2DModelOutput from ..modeling_utils import ModelMixin -from .vae import DecoderOutput logger = logging.get_logger(__name__) # pylint: disable=invalid-name @@ -476,10 +477,8 @@ def forward(self, hidden_states: torch.Tensor, drop_leading_frame: bool = True) class LTX2VideoDiffusionDecoder3d(nn.Module): """The LTX-2.5 diffusion video decoder. - Stages 1-4 deterministically upsample the latent into a context volume with neighborhood-attention blocks. Stage 5 - then denoises patchified pixels, conditioned on that context through AdaLN-Zero scale/shift. With - `model_output_type="x0"` and a single step — how LTX-2.5 ships — stage 5 runs once and its prediction *is* the - output; more steps add reverse Euler updates. + Stages 1-4 deterministically upsample the latent into a context volume with neighborhood-attention blocks; that + volume conditions stage 5, which is an ordinary diffusion transformer over patchified pixels. """ def __init__( @@ -573,13 +572,13 @@ def __init__( self.norm_out = nn.RMSNorm(stage5_channels, eps=1e-6) self.conv_out = nn.Linear(stage5_channels, noised_pixel_channels, bias=True) - def forward_stages_1_to_3(self, hidden_states: torch.Tensor) -> torch.Tensor: - """All deterministic stages but the last: latent `(B, C, T, H, W)` to a channels-last feature volume. + def encode_context_stages_1_to_3(self, hidden_states: torch.Tensor) -> torch.Tensor: + """All deterministic context stages but the last: latent `(B, C, T, H, W)` to a channels-last feature volume. - The trailing ghost frames added for NATTEN's border shift stay in the output; [`forward_stage_4`] crops them. - The split at this point exists for tiled decoding: these stages are cheap enough to run on the full volume, - while stage 4 and the diffusion stage — where the grid and the channel-hidden products get large — run per - tile. + The trailing ghost frames added for NATTEN's border shift stay in the output; [`encode_context_stage_4`] crops + them. The split at this point exists for tiled decoding: these stages are cheap enough to run on the full + volume, while stage 4 and the diffusion stage — where the grid and the channel-hidden products get large — run + per tile. """ num_pad = self.trailing_pad_latent_frames if num_pad > 0: @@ -596,10 +595,10 @@ def forward_stages_1_to_3(self, hidden_states: torch.Tensor) -> torch.Tensor: hidden_states = upsample(hidden_states) return hidden_states - def forward_stage_4( + def encode_context_stage_4( self, hidden_states: torch.Tensor, drop_leading_frame: bool = True, crop_trailing_ghost: bool = True ) -> torch.Tensor: - """Last deterministic stage: [`forward_stages_1_to_3`] output to context `(B, T_5, H_5, W_5, C_5)`. + """Last deterministic stage: [`encode_context_stages_1_to_3`] output to context `(B, T_5, H_5, W_5, C_5)`. The defaults describe the untiled decode. A tiled decode overrides them per temporal tile: only the tile containing t=0 drops the upsample's duplicate leading frame, and only the tile containing the video end carries @@ -616,10 +615,25 @@ def forward_stage_4( hidden_states = hidden_states[:, : -num_pad * self.temporal_compression_ratio] return hidden_states - def forward_diffusion_step( - self, latent_context: torch.Tensor, x_t: torch.Tensor, timestep: torch.Tensor + def forward( + self, hidden_states: torch.Tensor, latent_context: torch.Tensor, timestep: torch.Tensor ) -> torch.Tensor: - """One stage-5 step. Returns the model's prediction in pixel space, `(B, C, F, H, W)`.""" + r""" + One stage-5 denoising step. + + Args: + hidden_states (`torch.Tensor`): + Noised pixels of shape `(B, C, F, H, W)`. + latent_context (`torch.Tensor`): + The conditioning volume from [`encode_context_stage_4`], of shape `(B, F, H // patch_size, W // + patch_size, C_5)`. It is projected into the residual stream of every block rather than cross-attended, + and shares the token grid with `hidden_states`. + timestep (`torch.Tensor`): + Noise level in `[0, 1]`, of shape `(B,)`. + + Returns: + `torch.Tensor`: the model's prediction in pixel space, `(B, C, F, H, W)`. + """ t_emb = self.t_embedder( self.timestep_scale_multiplier * timestep, resolution=None, @@ -629,7 +643,7 @@ def forward_diffusion_step( ) modulation = self.shared_adaln(t_emb) - hidden_states = _patchify(x_t, self.patch_size).permute(0, 2, 3, 4, 1) + hidden_states = _patchify(hidden_states, self.patch_size).permute(0, 2, 3, 4, 1) hidden_states = self.conv_in_x_t(hidden_states) block_mask = self.diff_blocks[0].attn.build_block_mask(hidden_states) for block in self.diff_blocks: @@ -640,48 +654,6 @@ def forward_diffusion_step( hidden_states = hidden_states.permute(0, 4, 1, 2, 3).contiguous() return _unpatchify(hidden_states, self.patch_size) - def denoise(self, latent_context: torch.Tensor, x_t: torch.Tensor, num_inference_steps: int) -> torch.Tensor: - """Denoise `x_t` `(B, C, F, H, W)` through the stage-5 diffusion loop, conditioned on `latent_context`.""" - batch_size = latent_context.shape[0] - timesteps = torch.linspace( - 1.0, 1.0 / num_inference_steps, num_inference_steps, device=latent_context.device, dtype=torch.float32 - ) - - if num_inference_steps == 1 and self.model_output_type == "x0": - return self.forward_diffusion_step(latent_context, x_t, timesteps[:1].expand(batch_size)) - - for step_idx in range(num_inference_steps): - t_now = timesteps[step_idx].expand(batch_size) - t_next = timesteps[step_idx + 1] if step_idx + 1 < num_inference_steps else torch.zeros_like(t_now) - model_out = self.forward_diffusion_step(latent_context, x_t, t_now).float() - x_t_fp32 = x_t.float() - if self.model_output_type == "x0": - sigma = t_now.view(-1, *([1] * (x_t.ndim - 1))) - model_out = (x_t_fp32 - model_out) / sigma - dt = (t_now - t_next).view(-1, *([1] * (x_t.ndim - 1))) - x_t = (x_t_fp32 - dt * model_out).to(x_t.dtype) - return x_t - - def forward( - self, - hidden_states: torch.Tensor, - generator: torch.Generator | None = None, - num_inference_steps: int | None = None, - ) -> torch.Tensor: - num_inference_steps = num_inference_steps or self.default_num_inference_steps - latent_context = self.forward_stage_4(self.forward_stages_1_to_3(hidden_states)) - # The context grid is the stage-5 token grid, so the pixel canvas is its shape times the patch size — - # temporally that is the causal (T - 1) * ratio + 1 mapping of the LTX-2 latent space. - pixel_shape = ( - hidden_states.shape[0], - self.out_channels, - latent_context.shape[1], - latent_context.shape[2] * self.patch_size, - latent_context.shape[3] * self.patch_size, - ) - x_t = randn_tensor(pixel_shape, generator=generator, device=hidden_states.device, dtype=hidden_states.dtype) - return self.denoise(latent_context, x_t, num_inference_steps) - def _tile_intervals(length: int, tile_size: int, stride: int, min_size: int) -> list[tuple[int, int]]: """Overlapping `[start, end)` tiles covering `[0, length)`, with starts spaced `stride` apart. @@ -697,6 +669,70 @@ def _tile_intervals(length: int, tile_size: int, stride: int, min_size: int) -> return [(start, min(start + tile_size, length)) for start in starts[:-1]] + [(starts[-1], length)] +@dataclass(frozen=True) +class LTX2VideoDiffusionDecoderTileSchedule: + """Where a tiled decode cuts, and where each tile's pixels land. + + Tiles are expressed in cells of the grid entering the last deterministic stage — the only place the decoder can be + split, since everything before it is one attention neighbourhood over the whole volume. This holds the arithmetic + that mapping implies (cell-to-pixel scales, the causal frame offsets, the ghost frames NATTEN's border shift leaves + behind) so that the pipeline driving the tiles can stay a plain loop. + """ + + temporal: list[tuple[int, int]] + height: list[tuple[int, int]] + width: list[tuple[int, int]] + scales: tuple[int, int, int] + """Output pixels per cell, as (frames, height, width).""" + cell_strides: tuple[int, int, int] + """Distance between consecutive tile starts, in cells.""" + blend: tuple[int, int, int] + """Overlap to blend across a seam, in pixels.""" + num_frames: int + """Temporal cells of real video, i.e. excluding the trailing ghost frames.""" + total_frames: int + """Temporal cells including the ghost frames.""" + + @property + def num_tiles(self) -> int: + return len(self.temporal) * len(self.height) * len(self.width) + + @property + def pixel_shape(self) -> tuple[int, int, int]: + """The full decoded canvas, in pixels.""" + scale_t, scale_h, scale_w = self.scales + return ( + self.pixel_frames(self.num_frames, is_origin=True), + self.height[-1][1] * scale_h, + self.width[-1][1] * scale_w, + ) + + def pixel_frames(self, num_cells: int, is_origin: bool) -> int: + """Pixel frames `num_cells` cells decode to. + + The causal mapping spends the first cell on a single frame rather than `scale_t` of them, so a run that starts + at t=0 is one frame shorter than the cell count suggests. + """ + scale_t = self.scales[0] + return num_cells * scale_t - (1 if is_origin and scale_t == 2 else 0) + + def pixel_origin(self, t0: int) -> int: + """Where the tile starting at cell `t0` begins on the pixel canvas: exactly where the cells before it end. + + Which is one pixel frame earlier than `t0 * scale_t`, because the run those cells form starts at the origin and + so spends its first cell on a single frame. The origin tile itself starts at 0. + """ + return self.pixel_frames(t0, is_origin=True) if t0 else 0 + + def feature_end(self, t1: int) -> int: + """Where to stop slicing the feature volume for a tile ending at cell `t1`. + + Only the tile holding the end of the video carries the ghost frames into the last stage, since that is the only + place their border shift can still affect real output. + """ + return self.total_frames if t1 == self.num_frames else t1 + + class LTX2VideoDiffusionDecoderModel(ModelMixin, AttentionMixin, ConfigMixin): r""" The LTX-2 diffusion video decoder, introduced in LTX-2.5. @@ -707,7 +743,10 @@ class LTX2VideoDiffusionDecoderModel(ModelMixin, AttentionMixin, ConfigMixin): It is also a diffusion model rather than a deterministic decoder — it denoises pixels conditioned on a context volume built from the latents — which is why it is driven by [`LTX2VideoDiffusionDecodePipeline`] rather than being - passed as a pipeline's `vae`. + passed as a pipeline's `vae`. [`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 all live in that pipeline. What + stays here is the model itself, plus the tile *sizes* — [`enable_tiling`] configures the pipeline's tiling the way + `vae.enable_tiling()` does everywhere else. The latent statistics are carried here as buffers so the decode pipeline can denormalize without loading a second autoencoder just for two vectors. @@ -762,12 +801,10 @@ def __init__( default_num_inference_steps=decoder_num_inference_steps, ) - self.spatial_compression_ratio = spatial_compression_ratio - self.temporal_compression_ratio = temporal_compression_ratio - # When decoding a large enough video, the memory-dominant stages (the last deterministic stage and the # stage-5 diffusion blocks) can run on overlapping tiles that are blended back together. The earlier - # stages always see the full latent, so tiling changes the output only near tile borders. + # stages always see the full latent, so tiling changes the output only near tile borders. The decode + # pipeline reads the settings below; nothing here acts on them. self.use_tiling = False # The tile size and the distance between the starts of two consecutive tiles, in pixels/frames of the @@ -798,6 +835,8 @@ def enable_tiling( (they run at low resolution and are cheap); the last stage and the stage-5 diffusion blocks — which dominate decode memory — run on overlapping tiles whose seams are blended linearly. + These are settings, not behaviour: [`LTX2VideoDiffusionDecodePipeline`] reads them when it decodes. + Args: tile_sample_min_height (`int`, *optional*): The height of one decoded tile, in pixels. @@ -825,215 +864,114 @@ def disable_tiling(self) -> None: r"""Disable tiled decoding, returning to decoding the whole video in one pass.""" self.use_tiling = False - # Copied from diffusers.models.autoencoders.autoencoder_kl_ltx2.AutoencoderKLLTX2Video.blend_v - def blend_v(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: - blend_extent = min(a.shape[3], b.shape[3], blend_extent) - for y in range(blend_extent): - b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( - y / blend_extent - ) - return b - - # Copied from diffusers.models.autoencoders.autoencoder_kl_ltx2.AutoencoderKLLTX2Video.blend_h - def blend_h(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: - blend_extent = min(a.shape[4], b.shape[4], blend_extent) - for x in range(blend_extent): - b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( - x / blend_extent - ) - return b - - # Copied from diffusers.models.autoencoders.autoencoder_kl_ltx2.AutoencoderKLLTX2Video.blend_t - def blend_t(self, a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: - blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) - for x in range(blend_extent): - b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( - x / blend_extent - ) - return b - - def tiled_decode( - self, - z: torch.Tensor, - generator: torch.Generator | None = None, - num_inference_steps: int | None = None, - ) -> torch.Tensor: - r"""Decode a batch of latents with the last deterministic stage and the diffusion stage running per tile. + def get_tile_schedule( + self, feature_shape: torch.Size | tuple[int, int, int] + ) -> "LTX2VideoDiffusionDecoderTileSchedule": + """Plan a tiled decode over `feature_shape`, the `(T, H, W)` of an [`encode_context_stages_1_to_3`] output. - Tiles live on the grid entering the last deterministic stage, where one cell maps to a fixed block of output - pixels; the `tile_sample_*` sizes are converted to that grid, so they should be multiples of the cell size (8 - px spatially and 2 frames temporally for the production config). Temporal tiles follow the causal frame - mapping: the tile containing t=0 drops the temporal upsample's duplicate leading frame and only the tile - containing the video end carries the NATTEN border padding. + The cut is derived here, next to the stages whose geometry decides it, rather than in the pipeline that walks + it: how many pixels a cell covers, how small a tile the remaining attention kernels tolerate, and how many + ghost frames the border shift left on the end. The `tile_sample_*` sizes are in output pixels/frames and are + converted to cells, so they should be multiples of the cell size — 8 px and 2 frames for the production config. """ - decoder = self.decoder - num_inference_steps = num_inference_steps or decoder.default_num_inference_steps - batch_size = z.shape[0] - patch_size = decoder.patch_size - - # Pixels per cell of the tiling grid: the last upsample's stride times the stage-5 patch size. - upsample_stride = decoder.upsamples[-1].stride - scale_t, scale_h, scale_w = ( - upsample_stride[0], - upsample_stride[1] * patch_size, - upsample_stride[2] * patch_size, - ) - tile_t, stride_t = self.tile_sample_min_num_frames // scale_t, self.tile_sample_stride_num_frames // scale_t - tile_h, stride_h = self.tile_sample_min_height // scale_h, self.tile_sample_stride_height // scale_h - tile_w, stride_w = self.tile_sample_min_width // scale_w, self.tile_sample_stride_width // scale_w - # Every tile must satisfy both remaining neighborhood-attention kernels: the last deterministic stage - # sees the tile as-is, stage 5 sees it scaled by the upsample stride. + config = self.config + patch_size = config.patch_size + # One cell of this grid covers the last upsample's stride times the diffusion stage's patch size. + upsample_stride = config.decoder_upsample_strides[-1] + scales = (upsample_stride[0], upsample_stride[1] * patch_size, upsample_stride[2] * patch_size) + # Every tile must satisfy both remaining neighborhood-attention kernels: the last deterministic stage sees + # the tile as-is, the diffusion stage sees it scaled by the upsample stride. min_sizes = [ max(kernel_4, -(-kernel_5 // stride)) for kernel_4, kernel_5, stride in zip( - self.config.decoder_stage_kernels[-1], self.config.decoder_stage5_kernel, upsample_stride + config.decoder_stage_kernels[-1], config.decoder_stage5_kernel, upsample_stride ) ] - - features = decoder.forward_stages_1_to_3(z) # The trailing ghost frames replicate through the earlier stages' temporal upsamples, whose composed # mapping is affine with slope equal to the product of their strides. - ghost_frames = decoder.trailing_pad_latent_frames * math.prod(up.stride[0] for up in decoder.upsamples[:-1]) - num_frames = features.shape[1] - ghost_frames - height, width = features.shape[2], features.shape[3] - - temporal_tiles = _tile_intervals(num_frames, tile_t, stride_t, min_sizes[0]) - height_tiles = _tile_intervals(height, tile_h, stride_h, min_sizes[1]) - width_tiles = _tile_intervals(width, tile_w, stride_w, min_sizes[2]) - blend_frames = (tile_t - stride_t) * scale_t - blend_height = (tile_h - stride_h) * scale_h - blend_width = (tile_w - stride_w) * scale_w - - # A single-step x0 decode predicts pixels from pure noise, so each tile draws its own; a multi-step - # decode integrates its noise across steps, so overlapping tiles must start from the same canvas. - single_step_x0 = num_inference_steps == 1 and decoder.model_output_type == "x0" - x_t_full = None - if not single_step_x0: - pixel_frames = num_frames * scale_t - (1 if scale_t == 2 else 0) - x_t_full = randn_tensor( - (batch_size, decoder.out_channels, pixel_frames, height * scale_h, width * scale_w), - generator=generator, - device=z.device, - dtype=z.dtype, - ) + ghost_frames = self.decoder.trailing_pad_latent_frames * math.prod( + stride[0] for stride in config.decoder_upsample_strides[:-1] + ) - frame_groups = [] - for t0, t1 in temporal_tiles: - is_origin = t0 == 0 - is_trailing = t1 == num_frames - # The tile containing the video end takes the ghost frames with it into stage 4. - feature_t1 = features.shape[1] if is_trailing else t1 - rows = [] - for h0, h1 in height_tiles: - row = [] - for w0, w1 in width_tiles: - context = decoder.forward_stage_4( - features[:, t0:feature_t1, h0:h1, w0:w1], - drop_leading_frame=is_origin, - crop_trailing_ghost=is_trailing, - ) - tile_pixel_shape = ( - batch_size, - decoder.out_channels, - context.shape[1], - context.shape[2] * patch_size, - context.shape[3] * patch_size, - ) - if single_step_x0: - x_t = randn_tensor(tile_pixel_shape, generator=generator, device=z.device, dtype=z.dtype) - else: - # A non-origin tile keeps the duplicate leading frame, placing its first cell one pixel - # frame earlier than `t0 * scale_t` — the causal 1-then-`scale_t` frame mapping. - pixel_t0 = t0 * scale_t - (1 if not is_origin and scale_t == 2 else 0) - x_t = x_t_full[ - :, - :, - pixel_t0 : pixel_t0 + tile_pixel_shape[2], - h0 * scale_h : h0 * scale_h + tile_pixel_shape[3], - w0 * scale_w : w0 * scale_w + tile_pixel_shape[4], - ] - row.append(decoder.denoise(context, x_t, num_inference_steps)) - rows.append(row) - - result_rows = [] - for i, row in enumerate(rows): - result_row = [] - for j, tile in enumerate(row): - # blend the above tile and the left tile to the current tile and add the current tile to - # the result row - if i > 0: - tile = self.blend_v(rows[i - 1][j], tile, blend_height) - if j > 0: - tile = self.blend_h(row[j - 1], tile, blend_width) - # The last tile can extend past the stride grid (a short remnant is merged into it), so it - # keeps its full extent instead of being cropped to the stride. - keep_height = stride_h * scale_h if i < len(rows) - 1 else tile.shape[3] - keep_width = stride_w * scale_w if j < len(row) - 1 else tile.shape[4] - result_row.append(tile[:, :, :, :keep_height, :keep_width]) - result_rows.append(torch.cat(result_row, dim=4)) - frame_groups.append(torch.cat(result_rows, dim=3)) - - result = [] - for k, group in enumerate(frame_groups): - if k > 0: - group = self.blend_t(frame_groups[k - 1], group, blend_frames) - if k < len(frame_groups) - 1: - # The origin group is one frame short of `stride * scale`: its first cell decodes to a single - # pixel frame under the causal mapping. - keep_frames = stride_t * scale_t - (1 if k == 0 and scale_t == 2 else 0) - group = group[:, :, :keep_frames] - result.append(group) - return torch.cat(result, dim=2) + total_frames, height, width = feature_shape[0], feature_shape[1], feature_shape[2] + num_frames = total_frames - ghost_frames + tiles = (self.tile_sample_min_num_frames, self.tile_sample_min_height, self.tile_sample_min_width) + strides = ( + self.tile_sample_stride_num_frames, + self.tile_sample_stride_height, + self.tile_sample_stride_width, + ) + cell_tiles = tuple(tile // scale for tile, scale in zip(tiles, scales)) + cell_strides = tuple(stride // scale for stride, scale in zip(strides, scales)) + + return LTX2VideoDiffusionDecoderTileSchedule( + temporal=_tile_intervals(num_frames, cell_tiles[0], cell_strides[0], min_sizes[0]), + height=_tile_intervals(height, cell_tiles[1], cell_strides[1], min_sizes[1]), + width=_tile_intervals(width, cell_tiles[2], cell_strides[2], min_sizes[2]), + scales=scales, + cell_strides=cell_strides, + blend=tuple((tile - stride) * scale for tile, stride, scale in zip(cell_tiles, cell_strides, scales)), + num_frames=num_frames, + total_frames=total_frames, + ) + # `@apply_forward_hook` on both context stages: accelerate's offload hooks fire on `forward`, and the + # decode pipeline calls these before it ever calls one, so without it a CPU-offloaded model stays on the + # CPU here. @apply_forward_hook - def decode( - self, - z: torch.Tensor, - generator: torch.Generator | None = None, - num_inference_steps: int | None = None, - return_dict: bool = True, - ) -> DecoderOutput | torch.Tensor: - """Decode a batch of latents. + def encode_context_stages_1_to_3(self, hidden_states: torch.Tensor) -> torch.Tensor: + r"""All deterministic context stages but the last: latent `(B, C, T, H, W)` to a channels-last feature volume. - `z` is expected to be denormalized already (the pipeline applies `latents_mean` / `latents_std`), matching - [`AutoencoderKLLTX2Video`]. This decoder denoises, so pass `generator` for reproducibility. + The trailing ghost frames added for NATTEN's border shift stay in the output; [`encode_context_stage_4`] crops + them. The split exists for tiled decoding: these stages are cheap enough to run on the full volume, while stage + 4 and the diffusion stage — where the grid and the channel-hidden products get large — run per tile. """ - tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio - tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio - tile_latent_min_num_frames = self.tile_sample_min_num_frames // self.temporal_compression_ratio - if self.use_tiling and ( - z.shape[2] > tile_latent_min_num_frames - or z.shape[3] > tile_latent_min_height - or z.shape[4] > tile_latent_min_width - ): - decoded = self.tiled_decode(z, generator=generator, num_inference_steps=num_inference_steps) - else: - decoded = self.decoder(z, generator=generator, num_inference_steps=num_inference_steps) + return self.decoder.encode_context_stages_1_to_3(hidden_states) - if not return_dict: - return (decoded,) - return DecoderOutput(sample=decoded) + @apply_forward_hook + def encode_context_stage_4( + self, hidden_states: torch.Tensor, drop_leading_frame: bool = True, crop_trailing_ghost: bool = True + ) -> torch.Tensor: + r"""Last deterministic stage: [`encode_context_stages_1_to_3`] output to context `(B, T_5, H_5, W_5, C_5)`. + + The defaults describe an untiled decode. A tiled decode overrides them per temporal tile: only the tile + containing t=0 drops the upsample's duplicate leading frame, and only the tile containing the video end carries + the trailing ghost frames to crop. + """ + return self.decoder.encode_context_stage_4( + hidden_states, drop_leading_frame=drop_leading_frame, crop_trailing_ghost=crop_trailing_ghost + ) def forward( self, - z: torch.Tensor, - generator: torch.Generator | None = None, - num_inference_steps: int | None = None, + hidden_states: torch.Tensor, + latent_context: torch.Tensor, + timestep: torch.Tensor, return_dict: bool = True, - ) -> DecoderOutput | tuple[torch.Tensor]: + ) -> Transformer2DModelOutput | tuple[torch.Tensor]: r""" + One denoising step. The loop over steps, and the tiling around it, belong to + [`LTX2VideoDiffusionDecodePipeline`]. + Args: - z (`torch.Tensor`): - Latents of shape `(B, C, F, H, W)`, expected to be denormalized already (the pipeline applies - `latents_mean` / `latents_std`), matching [`AutoencoderKLLTX2Video`]. - generator (`torch.Generator`, *optional*): - This decoder denoises, so pass a generator to make decoding reproducible. - num_inference_steps (`int`, *optional*): - Number of denoising steps. Defaults to the decoder's `decoder_num_inference_steps` config value. + hidden_states (`torch.Tensor`): + Noised pixels of shape `(B, C, F, H, W)`. + latent_context (`torch.Tensor`): + The conditioning volume from [`encode_context_stage_4`], of shape `(B, F, H // patch_size, W // + patch_size, C_5)`. It is projected into the residual stream of every block rather than cross-attended, + and shares the token grid with `hidden_states`. + timestep (`torch.Tensor`): + Noise level in `[0, 1]`, of shape `(B,)`. return_dict (`bool`, *optional*, defaults to `True`): - Whether to return a [`~models.autoencoders.vae.DecoderOutput`] instead of a plain tuple. + Whether to return a [`~models.modeling_outputs.Transformer2DModelOutput`] instead of a plain tuple. Returns: - [`~models.autoencoders.vae.DecoderOutput`] or `tuple` + [`~models.modeling_outputs.Transformer2DModelOutput`] or `tuple`: the model's prediction in pixel space, + `(B, C, F, H, W)`. Whether that is the denoised sample or the velocity is set by the + `decoder_model_output_type` config value. """ - return self.decode(z, generator=generator, num_inference_steps=num_inference_steps, return_dict=return_dict) + sample = self.decoder(hidden_states, latent_context, timestep) + + if not return_dict: + return (sample,) + return Transformer2DModelOutput(sample=sample) diff --git a/src/diffusers/modular_pipelines/ltx2/decoders.py b/src/diffusers/modular_pipelines/ltx2/decoders.py index fc957a3f9925..344e4d140c8a 100644 --- a/src/diffusers/modular_pipelines/ltx2/decoders.py +++ b/src/diffusers/modular_pipelines/ltx2/decoders.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext from typing import Any import torch @@ -25,6 +26,7 @@ # `src/diffusers/models/` and re-export from `diffusers.models` before this lands. Imported from the # pipelines path here only so the draft is runnable; switch to the models path once moved. from ...pipelines.ltx2.vocoder import LTX2Vocoder +from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor @@ -85,6 +87,258 @@ def _unpack_audio_latents( return latents +# The diffusion decoder's denoising loop and its mid-network tiling live on +# `LTX2VideoDiffusionDecodePipeline`. Modular blocks must not import from `diffusers.pipelines.*` +# (modular.md gotcha #1), so they are copied here and kept in sync by `make fix-copies`. +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._blend_v +def _blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + """Blend `b`'s top edge into `a`'s bottom edge with a linear ramp. See `AutoencoderKLLTX2Video.blend_v`.""" + blend_extent = min(a.shape[3], b.shape[3], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._blend_h +def _blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + """Blend `b`'s left edge into `a`'s right edge with a linear ramp. See `AutoencoderKLLTX2Video.blend_h`.""" + blend_extent = min(a.shape[4], b.shape[4], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._blend_t +def _blend_t(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + """Blend `b`'s first frames into `a`'s last frames with a linear ramp. See `AutoencoderKLLTX2Video.blend_t`.""" + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._check_scheduler +def _check_scheduler(scheduler: FlowMatchEulerDiscreteScheduler) -> None: + """Reject only what this pipeline cannot drive, which is resolution-dependent shifting. + + Nothing else about the scheduler is checked. The sigmas the decoder was distilled on are a default, not a + requirement: a shift, a terminal shift or a stochastic update are all legitimate choices, and a finetune may well + want them — moving the loop onto a scheduler is what makes them possible. + """ + if scheduler.config.use_dynamic_shifting: + raise ValueError( + f"{scheduler.__class__.__name__} has `use_dynamic_shifting=True`, which needs a resolution-derived " + "`mu` that this pipeline does not compute. Use a scheduler with `use_dynamic_shifting=False` — the " + "converted checkpoints ship one in a `diffusion_decoder_scheduler` subfolder. Note that a " + "transformer's scheduler usually has it on, so it cannot be reused here as-is." + ) + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._progress_bar +def _progress_bar(progress_bar, total: int): + """`progress_bar(total=total)` when the caller has one, an inert context otherwise.""" + return progress_bar(total=total) if progress_bar is not None else nullcontext() + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._decoder_sigmas +def _decoder_sigmas(decoder: LTX2VideoDiffusionDecoderModel, num_inference_steps: int | None = None) -> list[float]: + """The decoder's sigma schedule: `linspace(1, 1 / num_inference_steps, num_inference_steps)`. + + This has to be handed to `set_timesteps` explicitly rather than left to the scheduler: its own default walks + `linspace(sigma_max, sigma_min, n)` with `sigma_min = 1 / num_train_timesteps`, i.e. 0.001 rather than `1 / n`, so + the two agree only at n=1 and no static config reconciles them. `num_inference_steps` defaults to what the decoder + was distilled for. + """ + if num_inference_steps is None: + num_inference_steps = decoder.config.decoder_num_inference_steps + return torch.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps, dtype=torch.float32).tolist() + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._denoise +def _denoise( + decoder: LTX2VideoDiffusionDecoderModel, + scheduler: FlowMatchEulerDiscreteScheduler, + latent_context: torch.Tensor, + x_t: torch.Tensor, + sigmas: list[float], + progress_bar=None, +) -> torch.Tensor: + """Denoise `x_t` `(B, C, F, H, W)` through the decoder's diffusion stage, conditioned on `latent_context`.""" + model_output_type = decoder.config.decoder_model_output_type + batch_size, dtype = latent_context.shape[0], x_t.dtype + + # Once per call, not once per decode: a tiled decode runs this loop from the top for every tile, and + # `set_timesteps` is what rewinds the scheduler's step index between them. + scheduler.set_timesteps(sigmas=sigmas, device=x_t.device) + num_inference_steps = len(scheduler.timesteps) + + for i, t in enumerate(scheduler.timesteps): + # The decoder takes the noise level in [0, 1] and scales it itself, so hand it the sigma rather than the + # scheduler's `sigma * num_train_timesteps` timestep. + sigma = scheduler.sigmas[i] + model_output = decoder(x_t, latent_context, sigma.expand(batch_size), return_dict=False)[0] + + # An x0 prediction at the last step *is* the sample: the Euler update to t=0 reduces to + # `x_t - t * (x_t - prediction) / t`. Returning it skips a full-canvas float32 round trip. + if model_output_type == "x0" and i == num_inference_steps - 1: + if progress_bar is not None: + progress_bar.update() + return model_output + + model_output = model_output.float() + if model_output_type == "x0": + # The scheduler integrates a velocity, so turn the sample prediction into one. + model_output = (x_t.float() - model_output) / sigma + # `step` returns in `model_output`'s dtype, float32 above to keep the update off the canvas dtype; the + # canvas itself stays in the dtype its noise was drawn in. + x_t = scheduler.step(model_output, t, x_t, return_dict=False)[0].to(dtype) + if progress_bar is not None: + progress_bar.update() + return x_t + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._tiled_decode +def _tiled_decode( + decoder: LTX2VideoDiffusionDecoderModel, + scheduler: FlowMatchEulerDiscreteScheduler, + z: torch.Tensor, + generator: torch.Generator | None, + sigmas: list[float], + progress_bar=None, +) -> torch.Tensor: + """Decode with the last deterministic stage and the diffusion stage running per tile. + + This tiles unconditionally; [`LTX2VideoDiffusionDecodePipeline.__call__`] is what consults `use_tiling` and the + video size before routing here. The cut itself comes from [`LTX2VideoDiffusionDecoderModel.get_tile_schedule`] — it + is a fact about the decoder's grid, not about sampling. What is here is the part that has to be: each tile runs its + own denoising loop, so the loop over tiles necessarily wraps the loop over steps. + """ + config = decoder.config + batch_size = z.shape[0] + patch_size = config.patch_size + + features = decoder.encode_context_stages_1_to_3(z) + schedule = decoder.get_tile_schedule(features.shape[1:4]) + scale_t, scale_h, scale_w = schedule.scales + stride_t, stride_h, stride_w = schedule.cell_strides + blend_frames, blend_height, blend_width = schedule.blend + + # A single-step x0 decode predicts pixels from pure noise, so each tile draws its own; a multi-step decode + # integrates its noise across steps, so overlapping tiles must start from the same canvas. + single_step_x0 = len(sigmas) == 1 and config.decoder_model_output_type == "x0" + x_t_full = None + if not single_step_x0: + x_t_full = randn_tensor( + (batch_size, config.out_channels, *schedule.pixel_shape), + generator=generator, + device=z.device, + dtype=z.dtype, + ) + + frame_groups = [] + with _progress_bar(progress_bar, schedule.num_tiles * len(sigmas)) as bar: + for t0, t1 in schedule.temporal: + rows = [] + for h0, h1 in schedule.height: + row = [] + for w0, w1 in schedule.width: + context = decoder.encode_context_stage_4( + features[:, t0 : schedule.feature_end(t1), h0:h1, w0:w1], + drop_leading_frame=t0 == 0, + crop_trailing_ghost=t1 == schedule.num_frames, + ) + tile_pixel_shape = ( + batch_size, + config.out_channels, + context.shape[1], + context.shape[2] * patch_size, + context.shape[3] * patch_size, + ) + if single_step_x0: + x_t = randn_tensor(tile_pixel_shape, generator=generator, device=z.device, dtype=z.dtype) + else: + pixel_t0 = schedule.pixel_origin(t0) + x_t = x_t_full[ + :, + :, + pixel_t0 : pixel_t0 + tile_pixel_shape[2], + h0 * scale_h : h0 * scale_h + tile_pixel_shape[3], + w0 * scale_w : w0 * scale_w + tile_pixel_shape[4], + ] + row.append(_denoise(decoder, scheduler, context, x_t, sigmas, progress_bar=bar)) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile to the current tile and add the current tile to the + # result row + if i > 0: + tile = _blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = _blend_h(row[j - 1], tile, blend_width) + # The last tile can extend past the stride grid (a short remnant is merged into it), so it + # keeps its full extent instead of being cropped to the stride. + keep_height = stride_h * scale_h if i < len(rows) - 1 else tile.shape[3] + keep_width = stride_w * scale_w if j < len(row) - 1 else tile.shape[4] + result_row.append(tile[:, :, :, :keep_height, :keep_width]) + result_rows.append(torch.cat(result_row, dim=4)) + frame_groups.append(torch.cat(result_rows, dim=3)) + + result = [] + for k, group in enumerate(frame_groups): + if k > 0: + group = _blend_t(frame_groups[k - 1], group, blend_frames) + if k < len(frame_groups) - 1: + group = group[:, :, : schedule.pixel_frames(stride_t, is_origin=k == 0)] + result.append(group) + return torch.cat(result, dim=2) + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._should_tile +def _should_tile(decoder: LTX2VideoDiffusionDecoderModel, z: torch.Tensor) -> bool: + """Whether tiling is on *and* the video is big enough for the schedule to actually split it.""" + config = decoder.config + return decoder.use_tiling and ( + z.shape[2] > decoder.tile_sample_min_num_frames // config.temporal_compression_ratio + or z.shape[3] > decoder.tile_sample_min_height // config.spatial_compression_ratio + or z.shape[4] > decoder.tile_sample_min_width // config.spatial_compression_ratio + ) + + +# Copied from diffusers.pipelines.ltx2.pipeline_ltx2_diffusion_decode._untiled_decode +def _untiled_decode( + decoder: LTX2VideoDiffusionDecoderModel, + scheduler: FlowMatchEulerDiscreteScheduler, + z: torch.Tensor, + generator: torch.Generator | None, + sigmas: list[float], + progress_bar=None, +) -> torch.Tensor: + """Decode denormalized latents in one pass: every stage sees the whole volume.""" + config = decoder.config + latent_context = decoder.encode_context_stage_4(decoder.encode_context_stages_1_to_3(z)) + # The context grid is the diffusion stage's token grid, so the pixel canvas is its shape times the patch size — + # temporally that is the causal (T - 1) * ratio + 1 mapping of the LTX-2 latent space. + pixel_shape = ( + z.shape[0], + config.out_channels, + latent_context.shape[1], + latent_context.shape[2] * config.patch_size, + latent_context.shape[3] * config.patch_size, + ) + x_t = randn_tensor(pixel_shape, generator=generator, device=z.device, dtype=z.dtype) + with _progress_bar(progress_bar, len(sigmas)) as bar: + return _denoise(decoder, scheduler, latent_context, x_t, sigmas, progress_bar=bar) + + class LTX2TrimConditionTokensStep(ModularPipelineBlocks): model_name = "ltx2" @@ -134,13 +388,31 @@ def description(self) -> str: "Step that unpacks and decodes the denoised video latents with the LTX-2 diffusion decoder (or returns " "latents). Swap this in for `LTX2VaeDecoderStep` on checkpoints that ship the diffusion decoder, which " "from LTX-2.5 on is the native default. The decoder denoises rather than deterministically decoding, so " - "it draws its own noise from `generator` and takes no decode timestep." + "it draws its own noise from `generator`, runs its own denoising loop over " + "`decode_num_inference_steps`, and needs its own scheduler." ) @property def expected_components(self) -> list[ComponentSpec]: return [ ComponentSpec("diffusion_decoder", LTX2VideoDiffusionDecoderModel), + # Not `scheduler`: that name is the transformer's, and its config has `use_dynamic_shifting` on, which + # the decoder cannot satisfy — it walks a plain uniform sigma schedule. Created from config by default + # so a repo whose `modular_model_index.json` predates this component still loads. + ComponentSpec( + "diffusion_decoder_scheduler", + FlowMatchEulerDiscreteScheduler, + config=FrozenDict( + { + "num_train_timesteps": 1000, + "shift": 1.0, + "use_dynamic_shifting": False, + "shift_terminal": None, + "stochastic_sampling": False, + } + ), + default_creation_method="from_config", + ), ComponentSpec( "video_processor", VideoProcessor, @@ -161,6 +433,14 @@ def inputs(self) -> list[tuple[str, Any]]: ), InputParam.template("generator"), InputParam.template("dtype", required=True), + InputParam( + "decode_num_inference_steps", + type_hint=int, + description=( + "Number of denoising steps the diffusion decoder takes. Separate from `num_inference_steps`, " + "which belongs to the transformer's loop. Defaults to what the decoder was distilled for." + ), + ), ] @property @@ -197,8 +477,13 @@ def __call__(self, components, state: PipelineState) -> PipelineState: latents, components.latents_mean, components.latents_std, components.vae_scaling_factor ) latents = latents.to(decoder.dtype) - # It samples the noise it denoises, so pass the generator to keep decoding reproducible. - video = decoder.decode(latents, generator=block_state.generator, return_dict=False)[0] + # The decoder's `forward` is a single denoising step, so the loop (and the tiling around it) is driven + # from here. It samples the noise it denoises, so pass the generator to keep decoding reproducible. + scheduler = components.diffusion_decoder_scheduler + _check_scheduler(scheduler) + sigmas = _decoder_sigmas(decoder, block_state.decode_num_inference_steps) + decode = _tiled_decode if _should_tile(decoder, latents) else _untiled_decode + video = decode(decoder, scheduler, latents, block_state.generator, sigmas) block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type) self.set_block_state(state, block_state) diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py index 7c77aba94a74..c11feea99224 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx25.py @@ -39,8 +39,9 @@ class LTX25DecoderStep(SequentialPipelineBlocks): returns latents). Components: - diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) video_processor (`VideoProcessor`) audio_vae - (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) + diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) diffusion_decoder_scheduler + (`FlowMatchEulerDiscreteScheduler`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) + vocoder (`LTX2Vocoder`) Inputs: latents (`Tensor`): @@ -57,6 +58,9 @@ class LTX25DecoderStep(SequentialPipelineBlocks): Torch generator for deterministic generation. dtype (`dtype`): The dtype of the model inputs, can be generated in input step. + decode_num_inference_steps (`int`, *optional*): + Number of denoising steps the diffusion decoder takes. Separate from `num_inference_steps`, which belongs + to the transformer's loop. Defaults to what the decoder was distilled for. audio_latents (`Tensor`): Denoised audio latents. audio_num_frames (`int`): @@ -95,8 +99,9 @@ class LTX25ConditionDecoderStep(SequentialPipelineBlocks): latents with the diffusion decoder and vocodes the audio latents (or returns latents). Components: - diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) video_processor (`VideoProcessor`) audio_vae - (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) + diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) diffusion_decoder_scheduler + (`FlowMatchEulerDiscreteScheduler`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) + vocoder (`LTX2Vocoder`) Inputs: latents (`Tensor`): @@ -115,6 +120,9 @@ class LTX25ConditionDecoderStep(SequentialPipelineBlocks): Torch generator for deterministic generation. dtype (`dtype`): The dtype of the model inputs, can be generated in input step. + decode_num_inference_steps (`int`, *optional*): + Number of denoising steps the diffusion decoder takes. Separate from `num_inference_steps`, which belongs + to the transformer's loop. Defaults to what the decoder was distilled for. audio_latents (`Tensor`): Denoised audio latents. audio_num_frames (`int`): @@ -156,8 +164,9 @@ class LTX25AutoDecoderStep(AutoPipelineBlocks): - `LTX25DecoderStep` otherwise (text-to-video, image-to-video). Components: - diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) video_processor (`VideoProcessor`) audio_vae - (`AutoencoderKLLTX2Audio`) vocoder (`LTX2Vocoder`) + diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) diffusion_decoder_scheduler + (`FlowMatchEulerDiscreteScheduler`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) + vocoder (`LTX2Vocoder`) Inputs: latents (`Tensor`): @@ -176,6 +185,9 @@ class LTX25AutoDecoderStep(AutoPipelineBlocks): Torch generator for deterministic generation. dtype (`dtype`): The dtype of the model inputs, can be generated in input step. + decode_num_inference_steps (`int`, *optional*): + Number of denoising steps the diffusion decoder takes. Separate from `num_inference_steps`, which belongs + to the transformer's loop. Defaults to what the decoder was distilled for. audio_latents (`Tensor`): Denoised audio latents. audio_num_frames (`int`): @@ -230,7 +242,8 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) transformer (`LTX2VideoTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) guider (`LTX2Guidance`) - audio_guider (`LTX2Guidance`) diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) vocoder (`LTX2Vocoder`) + audio_guider (`LTX2Guidance`) diffusion_decoder (`LTX2VideoDiffusionDecoderModel`) + diffusion_decoder_scheduler (`FlowMatchEulerDiscreteScheduler`) vocoder (`LTX2Vocoder`) Inputs: prompt (`str`, *optional*): @@ -332,6 +345,9 @@ class LTX25AutoBlocks(SequentialPipelineBlocks): VAE-encoded reference-image latents used for image-to-video conditioning. output_type (`str`, *optional*, defaults to pil): Output format: 'pil', 'np', 'pt'. + decode_num_inference_steps (`int`, *optional*): + Number of denoising steps the diffusion decoder takes. Separate from `num_inference_steps`, which belongs + to the transformer's loop. Defaults to what the decoder was distilled for. Outputs: videos (`list`): diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_diffusion_decode.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_diffusion_decode.py index 2f1137830a57..11a5f7011d7a 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_ltx2_diffusion_decode.py +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_diffusion_decode.py @@ -12,10 +12,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +from contextlib import nullcontext + import torch from ...models.autoencoders import AutoencoderKLLTX2Video, LTX2VideoDiffusionDecoderModel +from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging +from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline from .pipeline_output import LTX2VideoDecodeOutput @@ -24,6 +28,245 @@ logger = logging.get_logger(__name__) # pylint: disable=invalid-name +def _blend_v(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + """Blend `b`'s top edge into `a`'s bottom edge with a linear ramp. See `AutoencoderKLLTX2Video.blend_v`.""" + blend_extent = min(a.shape[3], b.shape[3], blend_extent) + for y in range(blend_extent): + b[:, :, :, y, :] = a[:, :, :, -blend_extent + y, :] * (1 - y / blend_extent) + b[:, :, :, y, :] * ( + y / blend_extent + ) + return b + + +def _blend_h(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + """Blend `b`'s left edge into `a`'s right edge with a linear ramp. See `AutoencoderKLLTX2Video.blend_h`.""" + blend_extent = min(a.shape[4], b.shape[4], blend_extent) + for x in range(blend_extent): + b[:, :, :, :, x] = a[:, :, :, :, -blend_extent + x] * (1 - x / blend_extent) + b[:, :, :, :, x] * ( + x / blend_extent + ) + return b + + +def _blend_t(a: torch.Tensor, b: torch.Tensor, blend_extent: int) -> torch.Tensor: + """Blend `b`'s first frames into `a`'s last frames with a linear ramp. See `AutoencoderKLLTX2Video.blend_t`.""" + blend_extent = min(a.shape[-3], b.shape[-3], blend_extent) + for x in range(blend_extent): + b[:, :, x, :, :] = a[:, :, -blend_extent + x, :, :] * (1 - x / blend_extent) + b[:, :, x, :, :] * ( + x / blend_extent + ) + return b + + +def _check_scheduler(scheduler: FlowMatchEulerDiscreteScheduler) -> None: + """Reject only what this pipeline cannot drive, which is resolution-dependent shifting. + + Nothing else about the scheduler is checked. The sigmas the decoder was distilled on are a default, not a + requirement: a shift, a terminal shift or a stochastic update are all legitimate choices, and a finetune may well + want them — moving the loop onto a scheduler is what makes them possible. + """ + if scheduler.config.use_dynamic_shifting: + raise ValueError( + f"{scheduler.__class__.__name__} has `use_dynamic_shifting=True`, which needs a resolution-derived " + "`mu` that this pipeline does not compute. Use a scheduler with `use_dynamic_shifting=False` — the " + "converted checkpoints ship one in a `diffusion_decoder_scheduler` subfolder. Note that a " + "transformer's scheduler usually has it on, so it cannot be reused here as-is." + ) + + +def _progress_bar(progress_bar, total: int): + """`progress_bar(total=total)` when the caller has one, an inert context otherwise.""" + return progress_bar(total=total) if progress_bar is not None else nullcontext() + + +def _decoder_sigmas(decoder: LTX2VideoDiffusionDecoderModel, num_inference_steps: int | None = None) -> list[float]: + """The decoder's sigma schedule: `linspace(1, 1 / num_inference_steps, num_inference_steps)`. + + This has to be handed to `set_timesteps` explicitly rather than left to the scheduler: its own default walks + `linspace(sigma_max, sigma_min, n)` with `sigma_min = 1 / num_train_timesteps`, i.e. 0.001 rather than `1 / n`, so + the two agree only at n=1 and no static config reconciles them. `num_inference_steps` defaults to what the decoder + was distilled for. + """ + if num_inference_steps is None: + num_inference_steps = decoder.config.decoder_num_inference_steps + return torch.linspace(1.0, 1.0 / num_inference_steps, num_inference_steps, dtype=torch.float32).tolist() + + +def _denoise( + decoder: LTX2VideoDiffusionDecoderModel, + scheduler: FlowMatchEulerDiscreteScheduler, + latent_context: torch.Tensor, + x_t: torch.Tensor, + sigmas: list[float], + progress_bar=None, +) -> torch.Tensor: + """Denoise `x_t` `(B, C, F, H, W)` through the decoder's diffusion stage, conditioned on `latent_context`.""" + model_output_type = decoder.config.decoder_model_output_type + batch_size, dtype = latent_context.shape[0], x_t.dtype + + # Once per call, not once per decode: a tiled decode runs this loop from the top for every tile, and + # `set_timesteps` is what rewinds the scheduler's step index between them. + scheduler.set_timesteps(sigmas=sigmas, device=x_t.device) + num_inference_steps = len(scheduler.timesteps) + + for i, t in enumerate(scheduler.timesteps): + # The decoder takes the noise level in [0, 1] and scales it itself, so hand it the sigma rather than the + # scheduler's `sigma * num_train_timesteps` timestep. + sigma = scheduler.sigmas[i] + model_output = decoder(x_t, latent_context, sigma.expand(batch_size), return_dict=False)[0] + + # An x0 prediction at the last step *is* the sample: the Euler update to t=0 reduces to + # `x_t - t * (x_t - prediction) / t`. Returning it skips a full-canvas float32 round trip. + if model_output_type == "x0" and i == num_inference_steps - 1: + if progress_bar is not None: + progress_bar.update() + return model_output + + model_output = model_output.float() + if model_output_type == "x0": + # The scheduler integrates a velocity, so turn the sample prediction into one. + model_output = (x_t.float() - model_output) / sigma + # `step` returns in `model_output`'s dtype, float32 above to keep the update off the canvas dtype; the + # canvas itself stays in the dtype its noise was drawn in. + x_t = scheduler.step(model_output, t, x_t, return_dict=False)[0].to(dtype) + if progress_bar is not None: + progress_bar.update() + return x_t + + +def _tiled_decode( + decoder: LTX2VideoDiffusionDecoderModel, + scheduler: FlowMatchEulerDiscreteScheduler, + z: torch.Tensor, + generator: torch.Generator | None, + sigmas: list[float], + progress_bar=None, +) -> torch.Tensor: + """Decode with the last deterministic stage and the diffusion stage running per tile. + + This tiles unconditionally; [`LTX2VideoDiffusionDecodePipeline.__call__`] is what consults `use_tiling` and the + video size before routing here. The cut itself comes from [`LTX2VideoDiffusionDecoderModel.get_tile_schedule`] — it + is a fact about the decoder's grid, not about sampling. What is here is the part that has to be: each tile runs its + own denoising loop, so the loop over tiles necessarily wraps the loop over steps. + """ + config = decoder.config + batch_size = z.shape[0] + patch_size = config.patch_size + + features = decoder.encode_context_stages_1_to_3(z) + schedule = decoder.get_tile_schedule(features.shape[1:4]) + scale_t, scale_h, scale_w = schedule.scales + stride_t, stride_h, stride_w = schedule.cell_strides + blend_frames, blend_height, blend_width = schedule.blend + + # A single-step x0 decode predicts pixels from pure noise, so each tile draws its own; a multi-step decode + # integrates its noise across steps, so overlapping tiles must start from the same canvas. + single_step_x0 = len(sigmas) == 1 and config.decoder_model_output_type == "x0" + x_t_full = None + if not single_step_x0: + x_t_full = randn_tensor( + (batch_size, config.out_channels, *schedule.pixel_shape), + generator=generator, + device=z.device, + dtype=z.dtype, + ) + + frame_groups = [] + with _progress_bar(progress_bar, schedule.num_tiles * len(sigmas)) as bar: + for t0, t1 in schedule.temporal: + rows = [] + for h0, h1 in schedule.height: + row = [] + for w0, w1 in schedule.width: + context = decoder.encode_context_stage_4( + features[:, t0 : schedule.feature_end(t1), h0:h1, w0:w1], + drop_leading_frame=t0 == 0, + crop_trailing_ghost=t1 == schedule.num_frames, + ) + tile_pixel_shape = ( + batch_size, + config.out_channels, + context.shape[1], + context.shape[2] * patch_size, + context.shape[3] * patch_size, + ) + if single_step_x0: + x_t = randn_tensor(tile_pixel_shape, generator=generator, device=z.device, dtype=z.dtype) + else: + pixel_t0 = schedule.pixel_origin(t0) + x_t = x_t_full[ + :, + :, + pixel_t0 : pixel_t0 + tile_pixel_shape[2], + h0 * scale_h : h0 * scale_h + tile_pixel_shape[3], + w0 * scale_w : w0 * scale_w + tile_pixel_shape[4], + ] + row.append(_denoise(decoder, scheduler, context, x_t, sigmas, progress_bar=bar)) + rows.append(row) + + result_rows = [] + for i, row in enumerate(rows): + result_row = [] + for j, tile in enumerate(row): + # blend the above tile and the left tile to the current tile and add the current tile to the + # result row + if i > 0: + tile = _blend_v(rows[i - 1][j], tile, blend_height) + if j > 0: + tile = _blend_h(row[j - 1], tile, blend_width) + # The last tile can extend past the stride grid (a short remnant is merged into it), so it + # keeps its full extent instead of being cropped to the stride. + keep_height = stride_h * scale_h if i < len(rows) - 1 else tile.shape[3] + keep_width = stride_w * scale_w if j < len(row) - 1 else tile.shape[4] + result_row.append(tile[:, :, :, :keep_height, :keep_width]) + result_rows.append(torch.cat(result_row, dim=4)) + frame_groups.append(torch.cat(result_rows, dim=3)) + + result = [] + for k, group in enumerate(frame_groups): + if k > 0: + group = _blend_t(frame_groups[k - 1], group, blend_frames) + if k < len(frame_groups) - 1: + group = group[:, :, : schedule.pixel_frames(stride_t, is_origin=k == 0)] + result.append(group) + return torch.cat(result, dim=2) + + +def _should_tile(decoder: LTX2VideoDiffusionDecoderModel, z: torch.Tensor) -> bool: + """Whether tiling is on *and* the video is big enough for the schedule to actually split it.""" + config = decoder.config + return decoder.use_tiling and ( + z.shape[2] > decoder.tile_sample_min_num_frames // config.temporal_compression_ratio + or z.shape[3] > decoder.tile_sample_min_height // config.spatial_compression_ratio + or z.shape[4] > decoder.tile_sample_min_width // config.spatial_compression_ratio + ) + + +def _untiled_decode( + decoder: LTX2VideoDiffusionDecoderModel, + scheduler: FlowMatchEulerDiscreteScheduler, + z: torch.Tensor, + generator: torch.Generator | None, + sigmas: list[float], + progress_bar=None, +) -> torch.Tensor: + """Decode denormalized latents in one pass: every stage sees the whole volume.""" + config = decoder.config + latent_context = decoder.encode_context_stage_4(decoder.encode_context_stages_1_to_3(z)) + # The context grid is the diffusion stage's token grid, so the pixel canvas is its shape times the patch size — + # temporally that is the causal (T - 1) * ratio + 1 mapping of the LTX-2 latent space. + pixel_shape = ( + z.shape[0], + config.out_channels, + latent_context.shape[1], + latent_context.shape[2] * config.patch_size, + latent_context.shape[3] * config.patch_size, + ) + x_t = randn_tensor(pixel_shape, generator=generator, device=z.device, dtype=z.dtype) + with _progress_bar(progress_bar, len(sigmas)) as bar: + return _denoise(decoder, scheduler, latent_context, x_t, sigmas, progress_bar=bar) + + class LTX2VideoDiffusionDecodePipeline(DiffusionPipeline): r""" Decode LTX-2 video latents with the diffusion decoder introduced in LTX-2.5. @@ -32,11 +275,20 @@ class LTX2VideoDiffusionDecodePipeline(DiffusionPipeline): context volume built from the latents, so it needs a scheduler and a generator. Pair it with any LTX-2 pipeline run with `output_type="latent"`, passing `denormalize=False` since that path already applied the latent statistics. + Because the decoder denoises, the tiling lives here rather than on the model: tiles are cut in the *middle* of the + decoder, on the grid entering its last deterministic stage, and each tile runs its own denoising loop before the + results are blended. Turn it on with `pipe.diffusion_decoder.enable_tiling()`, which sets the tile sizes this + pipeline reads. + Args: diffusion_decoder ([`LTX2VideoDiffusionDecoderModel`]): - The diffusion video decoder. + The diffusion video decoder. Its `forward` is a single denoising step; this pipeline owns the loop. scheduler ([`FlowMatchEulerDiscreteScheduler`]): - Scheduler driving the decoder's denoising steps. + Scheduler driving the decoder's denoising steps. Not the transformer's: that one normally has + `use_dynamic_shifting=True`, which needs a `mu` this pipeline does not compute. Checkpoints converted by + `convert_ltx2_to_diffusers.py` ship a matching one in a `diffusion_decoder_scheduler` subfolder, configured + for the uniform schedule the LTX-2.5 decoder was distilled on. Anything else the scheduler can express — a + shift, a different sigma schedule — is a supported choice, not a misconfiguration. vae ([`AutoencoderKLLTX2Video`], *optional*): Only consulted for the latent statistics used to denormalize. When omitted the pipeline falls back to the LTX-2 defaults, so a decode-only workflow does not have to load a second autoencoder. @@ -48,7 +300,7 @@ class LTX2VideoDiffusionDecodePipeline(DiffusionPipeline): def __init__( self, diffusion_decoder: LTX2VideoDiffusionDecoderModel, - scheduler, + scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKLLTX2Video = None, ): super().__init__() @@ -80,6 +332,8 @@ def _denormalize_latents( def __call__( self, latents: torch.Tensor, + num_inference_steps: int | None = None, + sigmas: list[float] | None = None, generator: torch.Generator | list[torch.Generator] | None = None, output_type: str = "pil", return_dict: bool = True, @@ -90,6 +344,13 @@ def __call__( latents (`torch.Tensor`): Latents of shape `(B, C, F, H, W)`. Note that an LTX-2 pipeline run with `output_type="latent"` returns latents that are *already* denormalized, so pass `denormalize=False` for those. + num_inference_steps (`int`, *optional*): + Number of denoising steps. Defaults to the decoder's `decoder_num_inference_steps` config value, which + is what the checkpoint was distilled for — 1 for LTX-2.5. + sigmas (`list[float]`, *optional*): + Custom sigma schedule, overriding `num_inference_steps`. The default is the uniform `linspace(1, 1 / + num_inference_steps, num_inference_steps)` the decoder was trained on. Whatever is passed still goes + through the scheduler, so a scheduler configured with a `shift` reshapes this too. generator (`torch.Generator`, *optional*): The decoder samples the noise it denoises, so pass a generator to make decoding reproducible. output_type (`str`, *optional*, defaults to `"pil"`): @@ -103,6 +364,10 @@ def __call__( Returns: [`~pipelines.ltx2.pipeline_output.LTX2VideoDecodeOutput`] or `tuple` """ + if sigmas is not None and num_inference_steps is not None: + raise ValueError("Only one of `num_inference_steps` or `sigmas` can be passed, not both.") + _check_scheduler(self.scheduler) + device = self._execution_device latents = latents.to(device) @@ -111,7 +376,13 @@ def __call__( latents = self._denormalize_latents(latents, latents_mean, latents_std, scaling_factor) latents = latents.to(self.diffusion_decoder.dtype) - video = self.diffusion_decoder.decode(latents, generator=generator, return_dict=False)[0] + if sigmas is None: + sigmas = _decoder_sigmas(self.diffusion_decoder, num_inference_steps) + + decoder, scheduler = self.diffusion_decoder, self.scheduler + # Tiling is worth its seams only once the video actually exceeds one tile. + decode = _tiled_decode if _should_tile(decoder, latents) else _untiled_decode + video = decode(decoder, scheduler, latents, generator, sigmas, self.progress_bar) video = self.video_processor.postprocess_video(video, output_type=output_type) self.maybe_free_model_hooks() diff --git a/tests/models/autoencoders/test_models_ltx2_diffusion_decoder.py b/tests/models/autoencoders/test_models_ltx2_diffusion_decoder.py index df4b9536af4a..7acf716ab7e1 100644 --- a/tests/models/autoencoders/test_models_ltx2_diffusion_decoder.py +++ b/tests/models/autoencoders/test_models_ltx2_diffusion_decoder.py @@ -22,7 +22,7 @@ from diffusers.utils import is_kernels_available from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, require_accelerator, require_torch_gpu, torch_device +from ...testing_utils import enable_full_determinism, require_torch_gpu, torch_device from ..testing_utils import ( AttentionTesterMixin, BaseModelTesterConfig, @@ -44,7 +44,7 @@ class LTX2VideoDiffusionDecoderModelTesterConfig(BaseModelTesterConfig): @property def main_input_name(self): - return "z" + return "hidden_states" @property def model_class(self): @@ -76,25 +76,30 @@ def get_init_dict(self): } def get_dummy_inputs(self): - # The decoder takes latents directly now: 2 latent frames decode to 9 pixel frames. - latents = randn_tensor((2, 8, 2, 3, 3), generator=self.generator, device=torch_device) - # The decoder denoises, so it draws noise on every call: without a seeded generator no two forward - # passes agree and every output comparison below would be meaningless. - return {"z": latents, "generator": self.generator} + """One denoising step's worth of input: noised pixels, the context conditioning them, and a noise level. + + The latent this corresponds to is `(2, 8, 2, 3, 3)`, which decodes to 9 pixel frames of 48x48. The context + shares the diffusion stage's token grid, so it is that canvas divided by `patch_size` with the last stage's + channel count. Building it directly rather than by running the context stages keeps this a pure input + fixture -- `LTX2VideoDiffusionDecodePipeline` is where the two are wired together. + """ + hidden_states = randn_tensor((2, 3, 9, 48, 48), generator=self.generator, device=torch_device) + latent_context = randn_tensor((2, 9, 24, 24, 16), generator=self.generator, device=torch_device) + timestep = torch.full((2,), 0.5, device=torch_device) + return {"hidden_states": hidden_states, "latent_context": latent_context, "timestep": timestep} + + def get_dummy_latents(self): + """A latent to feed the context stages, matching the canvas `get_dummy_inputs` describes.""" + return randn_tensor((2, 8, 2, 3, 3), generator=self.generator, device=torch_device) + + def encode_context(self, model, latents): + """The two deterministic halves, as `LTX2VideoDiffusionDecodePipeline` runs them.""" + return model.encode_context_stage_4(model.encode_context_stages_1_to_3(latents)) class TestLTX2VideoDiffusionDecoderModel(LTX2VideoDiffusionDecoderModelTesterConfig, ModelTesterMixin): base_precision = 1e-2 - @pytest.mark.skip( - "`forward` runs through the `apply_forward_hook`-decorated `decode`, and that decorator's " - "`pre_forward` call clears the input device accelerate's `AlignDevicesHook` recorded for the caller, so the " - "output comes back on the last device of the split rather than the input device and the comparison raises. " - "`test_cpu_offload` covers split placement instead — there every submodule executes on the same device." - ) - def test_model_parallelism(self, base_model_output, tmp_path, atol=1e-5, rtol=0): - pass - class TestLTX2VideoDiffusionDecoderModelSwiGLUTiling(LTX2VideoDiffusionDecoderModelTesterConfig): """The SwiGLU evaluates in token tiles to bound decode memory; that must not change the result.""" @@ -112,14 +117,12 @@ def test_token_tiled_swiglu_matches_untiled(self): """ model = self.model_class(**self.get_init_dict()).to(torch_device).eval() inputs = self.get_dummy_inputs() - latent = inputs["z"] def decode(): - # Re-seed per call: the decoder samples the noise it denoises, so a shared generator would - # hand the second call different noise and the comparison would be vacuous. - generator = torch.Generator(device=torch_device).manual_seed(0) + # A fixed input rather than sampled noise: the point is that tiling the MLP changes nothing, so + # both calls have to see the same tensors. with torch.no_grad(): - return model.decode(latent, generator=generator, return_dict=False)[0] + return model(**inputs, return_dict=False)[0] original = ltx2_diffusion_decoder._SWIGLU_TILE_SIZE try: @@ -136,74 +139,84 @@ def decode(): ) -class TestLTX2VideoDiffusionDecoderModelTiling(LTX2VideoDiffusionDecoderModelTesterConfig): - """Tiled decoding: the early stages run on the full latent, stages 4-5 run per tile with blending. +class TestLTX2VideoDiffusionDecoderModelTileSchedule(LTX2VideoDiffusionDecoderModelTesterConfig): + """Where a tiled decode cuts, independently of anything decoding it. - The latent is 3x4x5 (17x64x80 pixels) so every axis is large enough to split: the tiling grid — the - stage-4 input grid — is 9x16x20, and the tile sizes below cut it into three temporal and two/three - spatial tiles. + The pipeline walks this schedule; the arithmetic in it is the decoder's own -- cell-to-pixel scales, the + causal frame mapping, the ghost frames NATTEN's border shift leaves on the end. Pinning it here rather than + only through a decode means a mistake reads as a wrong number instead of a wrong picture. """ - def get_latent(self): - return randn_tensor((1, 8, 3, 4, 5), generator=self.generator, device=torch_device) - - def decode(self, model, latent, num_inference_steps=None): - # Re-seed per call: the decoder samples the noise it denoises, so outputs are only comparable - # across calls that drew from the same generator state. - generator = torch.Generator("cpu").manual_seed(0) - with torch.no_grad(): - return model.decode(latent, generator=generator, num_inference_steps=num_inference_steps)[0] - - @require_accelerator - def test_tiles_covering_the_video_match_untiled_exactly(self): - """A tile schedule with a single covering tile must reproduce the untiled decode bit for bit. - - This pins the per-tile plumbing — the ghost-frame carry/crop, the leading-frame drop, and the - stitching — because any offset in them shifts the single tile's output relative to the untiled path. - The default tile sizes are larger than the test video, so `tiled_decode` builds exactly one tile. + TILES = { + "tile_sample_min_num_frames": 8, + "tile_sample_stride_num_frames": 6, + "tile_sample_min_height": 32, + "tile_sample_stride_height": 24, + "tile_sample_min_width": 32, + "tile_sample_stride_width": 24, + } + + def get_schedule(self, feature_shape=(11, 16, 20), **tiles): + model = self.model_class(**self.get_init_dict()) + model.enable_tiling(**{**self.TILES, **tiles}) + return model.get_tile_schedule(feature_shape) + + def test_cells_map_to_pixels_by_the_last_upsample_and_the_patch_size(self): + """A cell is the last upsample's stride times the diffusion stage's patch size: (2, 2, 2) x 2 here.""" + schedule = self.get_schedule() + assert schedule.scales == (2, 4, 4) + # 8 frames / 32 px tiles over those scales, with 6 / 24 strides, so the overlap is 1 cell each way. + assert schedule.cell_strides == (3, 6, 6) + assert schedule.blend == (2, 8, 8) + + def test_ghost_frames_are_excluded_from_the_cut_but_kept_for_the_last_tile(self): + """The border-shift padding is real signal for the final tile's attention and nothing else. + + With a kernel of 3 the decoder pads 2 latent frames, and the earlier temporal upsamples (strides 1, 2, 2) + carry them to 8 cells -- so an 11-cell feature volume holds 3 cells of video. """ - model = self.model_class(**self.get_init_dict()).to(torch_device).eval() - latent = self.get_latent() + schedule = self.get_schedule(feature_shape=(11, 16, 20)) + assert (schedule.total_frames, schedule.num_frames) == (11, 3) + # Only the tile ending at the last real cell reaches past it, and it reaches all the way. + assert schedule.feature_end(schedule.num_frames) == 11 + assert schedule.feature_end(2) == 2 - for num_inference_steps in (None, 3): # None: the single-step x0 shortcut; 3: the Euler loop - untiled = self.decode(model, latent, num_inference_steps) - generator = torch.Generator("cpu").manual_seed(0) - with torch.no_grad(): - tiled = model.tiled_decode(latent, generator=generator, num_inference_steps=num_inference_steps) - assert torch.equal(tiled, untiled), ( - f"single-tile tiled decode diverged from untiled by {(tiled - untiled).abs().max().item():.3e} " - f"with num_inference_steps={num_inference_steps}" - ) - - def test_tiled_decode_with_splits(self): - """Actually-split tiles must reassemble to the untiled output shape, on both noise paths. - - Values legitimately differ from the untiled decode (each tile sees a truncated attention context at - its borders), so this asserts geometry, not closeness. The multi-step run additionally covers the - shared noise canvas that overlapping tiles slice from. - """ - model = self.model_class(**self.get_init_dict()).to(torch_device).eval() - latent = self.get_latent() - untiled = self.decode(model, latent) - - model.enable_tiling( - # Tiling-grid cells are 2 frames x 4 px x 4 px here (last upsample stride (2, 2, 2), patch 2), so - # this is a 4-cell tile with a 3-cell stride temporally and 8x8-cell tiles with 6-cell strides - # spatially: tiles (0, 4), (3, 7), (6, 9) over T and (0, 8), (6, 16|20) over H/W. - tile_sample_min_num_frames=8, - tile_sample_stride_num_frames=6, - tile_sample_min_height=32, - tile_sample_stride_height=24, - tile_sample_min_width=32, - tile_sample_stride_width=24, - ) - for num_inference_steps in (None, 3): - tiled = self.decode(model, latent, num_inference_steps) - assert tiled.shape == untiled.shape - assert torch.isfinite(tiled).all() + def test_the_causal_frame_mapping_places_tiles_without_gaps_or_overlap(self): + """The origin cell decodes to one frame, every later cell to `scale_t`, so tile starts are offset by one. - model.disable_tiling() - assert torch.equal(self.decode(model, latent), untiled) + This is the arithmetic a tiled decode is most easily wrong about: an off-by-one here still produces a + full-sized video, just one sampled from the wrong slice of the noise canvas. + """ + schedule = self.get_schedule(feature_shape=(20, 16, 20), tile_sample_min_num_frames=8) + assert len(schedule.temporal) > 1, "need a real temporal split for this to mean anything" + scale_t = schedule.scales[0] + + assert schedule.pixel_origin(0) == 0 + for t0, _ in schedule.temporal[1:]: + # A non-origin tile keeps the upsample's duplicate leading frame, so it starts one frame early. + assert schedule.pixel_origin(t0) == t0 * scale_t - 1 + # Each group contributes exactly the frames the next one starts after. + for index, (t0, _) in enumerate(schedule.temporal[:-1]): + kept = schedule.pixel_frames(schedule.cell_strides[0], is_origin=index == 0) + assert schedule.pixel_origin(t0) + kept == schedule.pixel_origin(schedule.temporal[index + 1][0]) + + def test_a_short_trailing_remnant_is_merged_into_its_neighbour(self): + """Neighborhood attention rejects a grid smaller than its kernel, so a stub tile cannot stand alone.""" + # A stride of 6 cells over 20 would start a final tile at 18, leaving 2 cells -- under the kernel of 3. + schedule = self.get_schedule(feature_shape=(28, 16, 20), tile_sample_min_num_frames=8) + assert all(end - start >= 3 for start, end in schedule.temporal), schedule.temporal + assert schedule.temporal[-1][1] == schedule.num_frames, "the tiles must still cover the whole video" + + def test_tiles_cover_every_axis_end_to_end(self): + schedule = self.get_schedule(feature_shape=(20, 30, 40)) + for axis, tiles, length in ( + ("t", schedule.temporal, schedule.num_frames), + ("h", schedule.height, 30), + ("w", schedule.width, 40), + ): + assert tiles[0][0] == 0 and tiles[-1][1] == length, (axis, tiles) + for (_, prev_end), (next_start, _) in zip(tiles, tiles[1:]): + assert next_start < prev_end, f"{axis} tiles leave a gap: {tiles}" class TestLTX2VideoDiffusionDecoderModelMemory(LTX2VideoDiffusionDecoderModelTesterConfig, MemoryTesterMixin): @@ -235,8 +248,12 @@ def test_natten_processor_decodes(self): isinstance(processor, LTX2VideoVaeNeighborhoodNattenProcessor) for processor in processors.values() ) + # Through the context stages as well as the denoising step: the deterministic stages are where most + # of the neighborhood attention lives, and they use a different kernel size per stage. + latents = self.get_dummy_latents() with torch.no_grad(): - output = model.decode(inputs["z"], generator=inputs["generator"], return_dict=False)[0] + latent_context = self.encode_context(model, latents) + output = model(inputs["hidden_states"], latent_context, inputs["timestep"], return_dict=False)[0] - assert output.shape == (inputs["z"].shape[0], *self.output_shape) + assert output.shape == (inputs["hidden_states"].shape[0], *self.output_shape) assert torch.isfinite(output).all(), "NATTEN decode produced NaN/inf values" diff --git a/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py b/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py index 89806c78f946..e507c90de0f3 100644 --- a/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py +++ b/tests/pipelines/ltx2/test_ltx2_diffusion_decode.py @@ -13,6 +13,7 @@ # See the License for the specific language governing permissions and # limitations under the License. +import pytest import torch from diffusers import ( @@ -20,8 +21,10 @@ LTX2VideoDiffusionDecodePipeline, LTX2VideoDiffusionDecoderModel, ) +from diffusers.pipelines.ltx2 import pipeline_ltx2_diffusion_decode as decode_module +from diffusers.utils.torch_utils import randn_tensor -from ...testing_utils import enable_full_determinism, torch_device +from ...testing_utils import enable_full_determinism, require_accelerator, torch_device from .testing_utils import get_dummy_vae @@ -45,9 +48,9 @@ } -def _build(with_vae: bool = False): +def _build(with_vae: bool = False, **config_overrides): torch.manual_seed(0) - decoder = LTX2VideoDiffusionDecoderModel(**DECODER_CONFIG).to(torch_device).eval() + decoder = LTX2VideoDiffusionDecoderModel(**{**DECODER_CONFIG, **config_overrides}).to(torch_device).eval() # Non-trivial statistics, so a run that skipped denormalization would not accidentally match. with torch.no_grad(): decoder.latents_mean.copy_(torch.linspace(-0.1, 0.1, DECODER_CONFIG["latent_channels"])) @@ -64,8 +67,17 @@ def _build(with_vae: bool = False): .eval() ) - return LTX2VideoDiffusionDecodePipeline( - diffusion_decoder=decoder, scheduler=FlowMatchEulerDiscreteScheduler(), vae=vae + return LTX2VideoDiffusionDecodePipeline(diffusion_decoder=decoder, scheduler=_scheduler(), vae=vae) + + +def _scheduler(): + """The decoder's scheduler, as `convert_ltx2_to_diffusers.py` saves it: a plain uniform sigma walk.""" + return FlowMatchEulerDiscreteScheduler( + num_train_timesteps=1000, + shift=1.0, + use_dynamic_shifting=False, + shift_terminal=None, + stochastic_sampling=False, ) @@ -73,6 +85,26 @@ def _latents(): return torch.randn(1, 8, 2, 3, 3, generator=torch.Generator().manual_seed(1)).to(torch_device) +def _sigmas(pipe, num_inference_steps=None): + return decode_module._decoder_sigmas(pipe.diffusion_decoder, num_inference_steps) + + +def _decode(pipe, latents, num_inference_steps=None, tiled=False, generator=None): + """Decode denormalized latents down one path, without the pre/post-processing `__call__` puts around it. + + The path is named rather than routed: which one `__call__` picks is what + `test_decode_skips_tiling_for_a_video_that_fits_in_one_tile` is for, and reproducing the gate here would make + that test circular. + """ + decoder, sigmas = pipe.diffusion_decoder, _sigmas(pipe, num_inference_steps) + decode = decode_module._tiled_decode if tiled else decode_module._untiled_decode + # Re-seed per call: the decoder samples the noise it denoises, so outputs are only comparable across calls + # that drew from the same generator state. + generator = generator if generator is not None else torch.Generator("cpu").manual_seed(0) + with torch.no_grad(): + return decode(decoder, pipe.scheduler, latents, generator, sigmas, pipe.progress_bar) + + def test_decode_without_vae(): """`vae` is optional: the pipeline must fall back to the decoder's own latent statistics.""" pipe = _build(with_vae=False) @@ -112,3 +144,297 @@ def test_denormalize_can_be_skipped(): latents, generator=torch.Generator(torch_device).manual_seed(0), output_type="pt", denormalize=False ).frames assert not torch.equal(normalized, raw) + + +def test_sigma_schedule_is_uniform(): + """The decoder walks `linspace(1, 1/n, n)`, not the scheduler's default `linspace(sigma_max, sigma_min, n)`. + + Nothing downstream would raise if the scheduler's default schedule were used instead -- it is the same length + and the same shape -- so the schedule itself is what has to be pinned. + """ + pipe = _build() + assert _sigmas(pipe, 1) == [1.0] + assert _sigmas(pipe, 4) == [1.0, 0.75, 0.5, 0.25] + # The default comes from the checkpoint, i.e. what the decoder was distilled for. + assert len(_sigmas(pipe)) == pipe.diffusion_decoder.config.decoder_num_inference_steps + + +def test_num_inference_steps_and_sigmas_are_exclusive(): + pipe = _build() + with pytest.raises(ValueError, match="Only one of"): + pipe(_latents(), num_inference_steps=2, sigmas=[1.0, 0.5]) + + +def test_multi_step_decode_runs_the_scheduler_loop(): + """More than one step must actually integrate: the extra steps have to change the result.""" + pipe, latents = _build(), _latents() + one = pipe( + latents, num_inference_steps=1, generator=torch.Generator(torch_device).manual_seed(0), output_type="pt" + ).frames + three = pipe( + latents, num_inference_steps=3, generator=torch.Generator(torch_device).manual_seed(0), output_type="pt" + ).frames + assert one.shape == three.shape + assert not torch.equal(one, three) + assert torch.isfinite(three).all() + + +class TestTiling: + """Tiled decoding: the early stages run on the full latent, the last stage and the diffusion loop run per tile. + + `latent` is 3x4x5 (17x64x80 pixels) so every axis is large enough to split: the tiling grid -- the grid + entering the last deterministic stage -- is 9x16x20, and the tile sizes below cut it into three temporal and + two/three spatial tiles. Tests that never want a split use `small_latent` instead. + """ + + SPLIT_TILES = { + # Tiling-grid cells are 2 frames x 4 px x 4 px here (last upsample stride (2, 2, 2), patch 2), so this is + # a 4-cell tile with a 3-cell stride temporally and 8x8-cell tiles with 6-cell strides spatially: tiles + # (0, 4), (3, 7), (6, 9) over T and (0, 8), (6, 16|20) over H/W. + "tile_sample_min_num_frames": 8, + "tile_sample_stride_num_frames": 6, + "tile_sample_min_height": 32, + "tile_sample_stride_height": 24, + "tile_sample_min_width": 32, + "tile_sample_stride_width": 24, + } + + def latent(self): + return torch.randn(1, 8, 3, 4, 5, generator=torch.Generator().manual_seed(2)).to(torch_device) + + def small_latent(self): + """2x3x3 (9x48x48 pixels): under the default tile sizes, over `SPLIT_TILES`. + + Stage 5 attends over the whole grid, so decode cost grows with the square of the video. A test that only + exercises the single-tile path has no use for a splittable video and should not pay for one. + """ + return torch.randn(1, 8, 2, 3, 3, generator=torch.Generator().manual_seed(2)).to(torch_device) + + def test_tiles_covering_the_video_match_untiled_exactly(self): + """A tile schedule with a single covering tile must reproduce the untiled decode bit for bit. + + This pins the per-tile plumbing -- the ghost-frame carry/crop, the leading-frame drop, and the stitching -- + because any offset in them shifts the single tile's output relative to the untiled path. The default tile + sizes are larger than the test video, so `tiled_decode` builds exactly one tile. + """ + pipe, latent = _build(), self.small_latent() + + for num_inference_steps in (None, 3): # None: the single-step x0 shortcut; 3: the Euler loop + untiled = _decode(pipe, latent, num_inference_steps) + tiled = _decode(pipe, latent, num_inference_steps, tiled=True) + assert torch.equal(tiled, untiled), ( + f"single-tile tiled decode diverged from untiled by {(tiled - untiled).abs().max().item():.3e} " + f"with num_inference_steps={num_inference_steps}" + ) + + def test_tiled_decode_with_splits(self): + """Actually-split tiles must reassemble to the untiled output shape, on both noise paths. + + Values legitimately differ from the untiled decode (each tile sees a truncated attention context at its + borders), so this asserts geometry, not closeness. The multi-step run additionally covers the shared noise + canvas that overlapping tiles slice from. + """ + pipe, latent = _build(), self.latent() + untiled = _decode(pipe, latent) + + pipe.diffusion_decoder.enable_tiling(**self.SPLIT_TILES) + for num_inference_steps in (None, 3): + tiled = _decode(pipe, latent, num_inference_steps, tiled=True) + assert tiled.shape == untiled.shape + assert torch.isfinite(tiled).all() + + def test_tiled_decode_tiles_even_when_tiling_is_disabled(self): + """`tiled_decode` tiles on its own terms; `use_tiling` only gates whether `decode` routes to it. + + `disable_tiling` flips the routing flag and leaves the configured tile sizes alone, so a direct + `tiled_decode` call still has a split schedule to honor. This counts last-stage invocations rather than + comparing outputs because output comparison cannot see the failure: a `tiled_decode` that quietly fell back + to one full-grid tile would reproduce the untiled decode exactly and pass every other test in this class. + """ + pipe, latent = _build(), self.latent() + decoder = pipe.diffusion_decoder + decoder.enable_tiling(**self.SPLIT_TILES) + decoder.disable_tiling() + assert not decoder.use_tiling + + # The last deterministic stage runs once per tile, so its call count is the tile count. + stage_4_calls = [] + original_stage_4 = decoder.encode_context_stage_4 + + def counting_stage_4(hidden_states, *args, **kwargs): + stage_4_calls.append(tuple(hidden_states.shape[1:4])) + return original_stage_4(hidden_states, *args, **kwargs) + + decoder.encode_context_stage_4 = counting_stage_4 + try: + _decode(pipe, latent, tiled=True) + tiled_call_count = len(stage_4_calls) + + stage_4_calls.clear() + _decode(pipe, latent) + untiled_call_count = len(stage_4_calls) + finally: + del decoder.encode_context_stage_4 + + assert tiled_call_count > 1, ( + f"tiled_decode ran the last stage {tiled_call_count} time(s) with use_tiling=False; it must tile " + "regardless of the flag" + ) + assert untiled_call_count == 1, ( + f"decode ran the last stage {untiled_call_count} times with use_tiling=False; it must not tile" + ) + + def test_call_skips_tiling_for_a_video_that_fits_in_one_tile(self, monkeypatch): + """`__call__` sizes the latent up before routing, so tiling only engages when it would split. + + The two outcomes are indistinguishable from the output alone: a video below the tile size that reaches the + tiled path anyway gets a single-tile schedule, which decodes to the same pixels. So this asserts the routing + directly -- the tiled path is never entered -- and separately pins the contract that matters to callers, + that turning tiling on cannot change a small video's output. + """ + pipe, latent = _build(), self.small_latent() + decoder = pipe.diffusion_decoder + + def run(): + return pipe( + latent, generator=torch.Generator(torch_device).manual_seed(0), output_type="pt", denormalize=False + ).frames + + decoder.disable_tiling() + untiled = run() + + calls = [] + original = decode_module._tiled_decode + + def counting_tiled_decode(*args, **kwargs): + calls.append(1) + return original(*args, **kwargs) + + monkeypatch.setattr(decode_module, "_tiled_decode", counting_tiled_decode) + + # Default tile sizes are far larger than this 9x48x48 video, so the gate declines to tile. + decoder.enable_tiling() + fits_in_one_tile = run() + assert not calls, "__call__ routed to the tiled path for a video that fits in a single tile" + assert torch.equal(fits_in_one_tile, untiled), ( + "enabling tiling changed the output of a video below the tile size by " + f"{(fits_in_one_tile - untiled).abs().max().item():.3e}" + ) + + # Shrink the tiles below the video and the same latent must now route. + decoder.enable_tiling(**self.SPLIT_TILES) + run() + assert calls, "__call__ did not route to the tiled path for a video larger than the tile size" + + +@require_accelerator +def test_model_cpu_offload_decodes(): + """Offloading must survive the pipeline reaching into the decoder for its context stages. + + Accelerate's offload hook fires on `forward`, and this pipeline calls `encode_context_stages_1_to_3` and + `encode_context_stage_4` before it ever calls one -- so those carry `@apply_forward_hook`. Without it the + weights stay on the CPU and the first matmul raises a device mismatch, on both the tiled and untiled paths. + """ + for tiled in (False, True): + pipe = _build() + if tiled: + pipe.diffusion_decoder.enable_tiling(**TestTiling.SPLIT_TILES) + pipe.enable_model_cpu_offload(device=torch_device) + latents = torch.randn(1, 8, 3, 4, 5, generator=torch.Generator().manual_seed(2)) + frames = pipe(latents, generator=torch.Generator(torch_device).manual_seed(0), output_type="pt").frames + assert frames.shape == (1, 17, 3, 64, 80) + assert torch.isfinite(frames).all() + + +@pytest.mark.parametrize("model_output_type", ["v", "x0"]) +def test_scheduler_step_matches_the_closed_form_euler_update(model_output_type): + """The scheduler must integrate exactly what the decoder's own solver did, on both prediction types. + + This is the regression the move to a scheduler is most exposed to: `step` would still return a plausible + tensor if the sign of `dt` flipped, if the x0-to-velocity conversion used the wrong sigma, or if the sigma + handed to the model were the scheduler's `sigma * num_train_timesteps` timestep instead. So the loop is + recomputed here in closed form -- `x - (sigma - sigma_next) * v` -- and compared bit for bit. + """ + steps = 3 + pipe = _build(decoder_model_output_type=model_output_type, decoder_num_inference_steps=steps) + decoder, latents = pipe.diffusion_decoder, _latents() + sigmas = _sigmas(pipe) + assert len(sigmas) == steps + + with torch.no_grad(): + context = decoder.encode_context_stage_4(decoder.encode_context_stages_1_to_3(latents)) + pixel_shape = ( + latents.shape[0], + DECODER_CONFIG["out_channels"], + context.shape[1], + context.shape[2] * DECODER_CONFIG["patch_size"], + context.shape[3] * DECODER_CONFIG["patch_size"], + ) + # Same draw the pipeline makes, so both loops start from the same canvas. + x_t = randn_tensor( + pixel_shape, + generator=torch.Generator(torch_device).manual_seed(0), + device=latents.device, + dtype=latents.dtype, + ) + + # float32 scalars, matching the dtype the scheduler holds its sigmas in: a Python float would divide + # and subtract in double and leave a few ulps of difference that say nothing about the update rule. + sigma_values = torch.tensor(sigmas + [0.0], dtype=torch.float32, device=torch_device) + for i in range(steps): + sigma, sigma_next = sigma_values[i], sigma_values[i + 1] + prediction = decoder(x_t, context, sigma.expand(latents.shape[0]), return_dict=False)[0] + if model_output_type == "x0": + if i == steps - 1: + # The x0 shortcut: `x - sigma * (x - x0) / sigma` is the prediction itself, and taking it + # directly is what keeps the common one-step decode off a full-canvas float32 round trip. + expected = prediction + break + velocity = (x_t.float() - prediction.float()) / sigma + else: + velocity = prediction.float() + x_t = (x_t.float() - (sigma - sigma_next) * velocity).to(x_t.dtype) + else: + expected = x_t + + actual = _decode(pipe, latents, steps, generator=torch.Generator(torch_device).manual_seed(0)) + + assert torch.equal(actual, expected), ( + f"scheduler-driven decode diverged from the closed-form Euler update by " + f"{(actual - expected).abs().max().item():.3e} for model_output_type={model_output_type!r}" + ) + + +def test_a_reshaped_sigma_schedule_is_honoured(): + """A scheduler configured away from the shipped defaults must take effect, not be second-guessed. + + The uniform schedule is what the LTX-2.5 checkpoint was distilled on, but it is a default, not a law: a + finetune may prefer a shift, and driving the loop from a scheduler is what makes that expressible. So `shift` + has to reach the sigmas and change the decode, with no warning and no correction. + """ + pipe, latents = _build(), _latents() + shipped = pipe(latents, generator=torch.Generator(torch_device).manual_seed(0), output_type="pt").frames + + pipe.scheduler = FlowMatchEulerDiscreteScheduler(**{**_scheduler().config, "shift": 5.0}) + sigmas = _sigmas(pipe, 3) + shifted = pipe( + latents, sigmas=sigmas, generator=torch.Generator(torch_device).manual_seed(0), output_type="pt" + ).frames + + # `shift` bends the schedule the pipeline handed in, rather than being ignored or overridden. + assert pipe.scheduler.sigmas[:-1].tolist() != sigmas + assert not torch.equal(shipped, shifted) + assert torch.isfinite(shifted).all() + + +def test_dynamic_shifting_is_rejected_with_an_actionable_error(): + """The one scheduler setting this pipeline cannot drive: it never computes `mu`, so the decode cannot run. + + Worth its own error because it is what `scheduler=pipe.scheduler` gives you — a transformer's scheduler + normally has dynamic shifting on — and the scheduler's own complaint (`mu` must be passed) does not say where + to get a working one. + """ + pipe = _build() + pipe.scheduler = FlowMatchEulerDiscreteScheduler(**{**_scheduler().config, "use_dynamic_shifting": True}) + with pytest.raises(ValueError, match="diffusion_decoder_scheduler"): + pipe(_latents(), generator=torch.Generator(torch_device).manual_seed(0), output_type="pt")