diff --git a/docs/source/en/api/pipelines/ltx2.md b/docs/source/en/api/pipelines/ltx2.md index e73689b4af52..662d5c25c250 100644 --- a/docs/source/en/api/pipelines/ltx2.md +++ b/docs/source/en/api/pipelines/ltx2.md @@ -1030,6 +1030,117 @@ encode_video( You can see the supported workflows in the docs for each blockset (e.g. [`LTX2AutoBlocks`], [`LTX25AutoBlocks`]). +### Diffusion Fidelity Rendering (DFR) for LTX-2.5 + +`LTX2DFRPipeline` trades wall-clock time for detail fidelity. It generates at a fraction of the requested resolution *plus* extra single-pixel-frame **keyframe slots**, then re-renders at the full resolution seeded from both. A slot costs a full latent frame of tokens to buy one pixel frame, which relaxes the effective temporal compression at that position — so the surrounding video is conditioned on genuinely new frames instead of interpolated ones. Slot positions come from a segment grid aligned to the VAE's temporal border (24 or 32 pixel frames, whichever pads the request less); the canvas is padded to a whole number of segments internally and trimmed back before decoding. + +This needs a transformer whose config sets `use_keyframes_abs_pos_embedding`, which marks single-pixel-frame latents with a learned embedding. LTX-2.5 checkpoints ship it; the pipeline raises on anything older rather than spending the token budget on tokens the model cannot interpret. + +Budget for the extra tokens: each slot adds one latent frame's worth, so stage 2 runs a longer sequence than the equivalent two-stage distilled pass — +31% at 1024x1536 / 121 frames (24576 -> 32256 tokens, 5 slots on a 24-frame segment grid). Peak activation memory scales with that, so a resolution that just fits the plain distilled recipe may need `enable_sequential_cpu_offload`, `vae.enable_tiling()`, or a smaller canvas under DFR. + +The full recipe below is the one worth starting from: 1088x1920 image-to-video, one x2 temporal refine round, +and the x2 spatial detailing IC-LoRA on stage 2. The individual knobs are explained after it. + +```py +import torch +from diffusers import LTX2DFRPipeline +from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition +from diffusers.utils import encode_video, load_image + +# The published `model_index.json` has no upsamplers. The spatial one is a subfolder of the repo; the x2 temporal +# one is not published — convert it with `--temporal_latent_upsampler` (see "Temporal refinement" below). +latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + "Lightricks/LTX-2.5-Diffusers", subfolder="latent_upsampler", torch_dtype=torch.bfloat16 +) +temporal_latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + "path/to/converted/temporal_latent_upsampler", torch_dtype=torch.bfloat16 +) +pipe = LTX2DFRPipeline.from_pretrained( + "Lightricks/LTX-2.5-Diffusers", + latent_upsampler=latent_upsampler, + temporal_latent_upsampler=temporal_latent_upsampler, + torch_dtype=torch.bfloat16, +) + +# Stage 2 attends to the stage-1 latent through this IC-LoRA. It is calibrated for strength 0.5. +pipe.load_lora_weights("Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", adapter_name="detailing") +pipe.set_adapters(["detailing"], weights=[0.5]) +pipe.enable_model_cpu_offload() + +image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png") +frame_rate = 24.0 +video, audio = pipe( + prompt="A tabby cat stretching in a sunlit window, dust motes drifting in the light", + conditions=[LTX2VideoCondition(frames=image, index=0, strength=1.0)], + height=1088, + width=1920, + frame_rate=frame_rate, + # `num_frames` omitted: the `duration_head` predicts it from the prompt. + temporal_upscalings=1, + detailing_lora_adapter_name="detailing", + generator=torch.Generator(device="cuda").manual_seed(0), + output_type="np", + return_dict=False, +) + +# One refine round doubles the playback rate. +playback_fps = frame_rate * 2**1 +encode_video( + video[0], + fps=playback_fps, + audio=audio[0].float().cpu(), + audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + output_path="ltx2_5_dfr.mp4", +) +``` + +`height` and `width` are the *output* resolution and must be divisible by `2 ** spatial_upscalings` times the VAE's spatial compression ratio (64 for LTX-2.5 at the default `spatial_upscalings=1`, 128 at `spatial_upscalings=2`), since every stage below the output runs at a halved resolution. This is why a 4K run is **3840x2176**, not 3840x2160: UHD height has to be snapped up onto the 128 grid. 1080p is likewise 1920x1088, not 1920x1080. Each stage runs its own fixed distilled schedule (`stage_1_sigmas`, `stage_2_sigmas`), so there is no `num_inference_steps`; the distilled schedules are trained to be used without guidance, so there is no `negative_prompt` or `guidance_scale` either. The shipped audio is stage 1's — stage 2 and the temporal refine tiles still run an audio pass so the video branch has cross-modal attention; the waveform itself is not refined after stage 1. + +**Spatial detailing.** Load the 2x spatial detailing IC-LoRA under a named adapter and pass that name as `detailing_lora_adapter_name`. Stage 2 then runs with the adapter active *and* attends to the stage-1 half-resolution latent as an in-context reference; stage 1 and the temporal rounds run with it deactivated, and whatever adapters were active on entry are restored before the call returns. + +`detailing_reference_downscale_factor` (default `2`) scales the reference tokens' spatial coordinates into the target's coordinate space, and should match the LoRA's `reference_downscale_factor` metadata. The pipeline only activates and deactivates the adapter — its weight stays whatever you set, so give it the 0.5 it is calibrated for. + +**Temporal refinement.** `temporal_upscalings` (0-2) adds rounds that each double the frame rate: the canvas is temporally upsampled, split into `2 ** round` tiles that meet at shared keyframes, given fresh mid-segment slots, and densified with ancestral Euler. Each tile is handed the slice of the frozen stage-1 audio covering its own playback window, resampled to the tile's token count, so both sides of a seam densify against the same sound. Rounds need the optional `temporal_latent_upsampler` component, which is not part of the base repo — convert it with `--temporal_latent_upsampler` and pass it to `from_pretrained`, as in the example above. + +You always get `(num_frames - 1) * 2 ** temporal_upscalings + 1` frames back, and the returned video plays at `frame_rate * 2 ** temporal_upscalings`. Conditioning fps is 60 whenever playback is above 30, independently of muxing: RoPE time is `pixel_frame / fps`, so a 120 fps time base would halve every token's temporal span versus the trained distribution, and 48 fps would stretch it. Both lie that they are 60 and treat the decoded frames at the playback rate. + +**A third spatial stage.** `spatial_upscalings=2` starts the base canvas one more factor of two down and adds a detailing pass at the output resolution *after* the temporal rounds. A full-resolution forward pass over the refined canvas does not fit in one sequence, so the pass denoises the whole canvas in a single loop and tiles the *transformer call* inside it — two ways on each spatial axis and `2 ** temporal_upscalings` ways in time. Every Euler step therefore steps a canvas whose tiles have already agreed on their overlaps, and conditionings are attached once on the whole canvas and filtered per tile at the token level rather than cropped by hand. + +The two axes are seamed differently. Neither side of a spatial border holds a known answer, so those overlaps are blended with trapezoidal weights. The temporal tiles are cut instead, on the same keyframe seams the last refine round stitched on: both windows reproduce a shared keyframe there, so the later one drops its run-up rather than averaging it. The keyframes the rounds settled are handed to the epilogue as content — each is decoded on its own, stretched x2 with Lanczos in RGB, encoded again at the output resolution, and pinned fully clean — so nothing along a seam is generated twice. Since the epilogue is a detailing pass, pass `detailing_lora_adapter_name` with it: + +```py +video, audio = pipe( + prompt=prompt, + height=1024, + width=1920, # both divisible by 128 at spatial_upscalings=2 + num_frames=121, + spatial_upscalings=2, + detailing_lora_adapter_name="detailing", + output_type="np", + return_dict=False, +) +``` + +**Decoding with the diffusion decoder.** `LTX2DFRPipeline` decodes with its convolutional `vae`. For maximum detail fidelity, run with `output_type="latent"` and hand the latents to [`LTX2VideoDiffusionDecodePipeline`]. + +```py +from diffusers import LTX2VideoDiffusionDecodePipeline +from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel + +video_latents, audio_latents = pipe(prompt=prompt, output_type="latent", return_dict=False) + +decoder = LTX2VideoDiffusionDecoderModel.from_pretrained( + "Lightricks/LTX-2.5-Diffusers", subfolder="diffusion_decoder", dtype=torch.bfloat16 +) +decode_pipe = LTX2VideoDiffusionDecodePipeline( + diffusion_decoder=decoder, scheduler=pipe.scheduler, vae=pipe.vae +) +decode_pipe.enable_model_cpu_offload() +# `denormalize=False`: the `output_type="latent"` path already applied the latent statistics. +video = decode_pipe(latents=video_latents, denormalize=False, output_type="np", return_dict=False)[0] +``` + ## LTX2Pipeline [[autodoc]] LTX2Pipeline diff --git a/scripts/convert_ltx2_to_diffusers.py b/scripts/convert_ltx2_to_diffusers.py index 57a3e88c7465..33b91790ef1c 100644 --- a/scripts/convert_ltx2_to_diffusers.py +++ b/scripts/convert_ltx2_to_diffusers.py @@ -443,7 +443,7 @@ def get_ltx2_transformer_config(version: str) -> tuple[dict[str, Any], dict[str, "use_prompt_embeddings": False, "perturbed_attn": True, # The only transformer-level deltas from 2.3: the video FFN drops its bias (audio_ff_bias and - # use_prompt_adaln_single keep their True defaults for this checkpoint), and 2.5.1+ carries a + # use_prompt_adaln_single keep their True defaults for this checkpoint), and 2.5 carries a # learned keyframe absolute-position embedding. "ff_bias": False, "use_keyframes_abs_pos_embedding": True, @@ -1222,7 +1222,7 @@ def get_ltx2_spatial_latent_upsampler_config(version: str): "rational_spatial_scale": 2.0, "use_rational_resampler": True, } - elif version == "2.3": + elif version in ("2.3", "2.5"): config = { "in_channels": 128, "mid_channels": 1024, @@ -1238,9 +1238,21 @@ def get_ltx2_spatial_latent_upsampler_config(version: str): return config -def convert_ltx2_spatial_latent_upsampler( - original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype -): +def get_ltx2_temporal_latent_upsampler_config(version: str): + if version != "2.5": + raise ValueError(f"Unsupported version: {version}") + # The temporal x2 upsampler is narrower than its spatial sibling and pixel-shuffles along time only. + return { + "in_channels": 128, + "mid_channels": 512, + "num_blocks_per_stage": 4, + "dims": 3, + "spatial_upsample": False, + "temporal_upsample": True, + } + + +def convert_ltx2_latent_upsampler(original_state_dict: dict[str, Any], config: dict[str, Any], dtype: torch.dtype): with init_empty_weights(): latent_upsampler = LTX2LatentUpsamplerModel(**config) @@ -1379,6 +1391,12 @@ def none_or_str(value: str): "google/gemma-4-E2B-it or google/gemma-4-E4B-it." ), ) + parser.add_argument( + "--temporal_latent_upsampler_filename", + default="ltx-2.5-latent-temporal-upscaler-x2-bf16-1.0.safetensors", + type=none_or_str, + help="Temporal x2 latent upsampler filename (LTX-2.5, used by the DFR pipeline's temporal refine rounds)", + ) parser.add_argument( "--latent_upsampler_filename", default="ltx-2-spatial-upscaler-x2-1.0.safetensors", @@ -1410,6 +1428,11 @@ def none_or_str(value: str): parser.add_argument("--vocoder", action="store_true", help="Whether to convert the vocoder model") parser.add_argument("--text_encoder", action="store_true", help="Whether to conver the text encoder") parser.add_argument("--latent_upsampler", action="store_true", help="Whether to convert the latent upsampler") + parser.add_argument( + "--temporal_latent_upsampler", + action="store_true", + help="Whether to convert the temporal x2 latent upsampler (LTX-2.5)", + ) parser.add_argument( "--full_pipeline", action="store_true", @@ -1608,7 +1631,7 @@ def main(args): repo_id=args.original_state_dict_repo_id, filename=args.latent_upsampler_filename ) latent_upsampler_config = get_ltx2_spatial_latent_upsampler_config(args.version) - latent_upsampler = convert_ltx2_spatial_latent_upsampler( + latent_upsampler = convert_ltx2_latent_upsampler( original_latent_upsampler_ckpt, latent_upsampler_config, dtype=vae_dtype, @@ -1616,6 +1639,17 @@ def main(args): if not args.full_pipeline and not args.upsample_pipeline: latent_upsampler.save_pretrained(os.path.join(args.output_path, "latent_upsampler")) + if args.temporal_latent_upsampler: + original_temporal_upsampler_ckpt = load_hub_or_local_checkpoint( + repo_id=args.original_state_dict_repo_id, filename=args.temporal_latent_upsampler_filename + ) + temporal_latent_upsampler = convert_ltx2_latent_upsampler( + original_temporal_upsampler_ckpt, + get_ltx2_temporal_latent_upsampler_config(args.version), + dtype=vae_dtype, + ) + temporal_latent_upsampler.save_pretrained(os.path.join(args.output_path, "temporal_latent_upsampler")) + if args.full_pipeline: is_distilled_ckpt = "distilled" in args.combined_filename if is_distilled_ckpt: diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b7d79b8ee97d..7dc599b75365 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -724,6 +724,7 @@ "LongCatImageEditPipeline", "LongCatImagePipeline", "LTX2ConditionPipeline", + "LTX2DFRPipeline", "LTX2HDRPipeline", "LTX2ImageToVideoPipeline", "LTX2InContextPipeline", @@ -1570,6 +1571,7 @@ LongCatImageEditPipeline, LongCatImagePipeline, LTX2ConditionPipeline, + LTX2DFRPipeline, LTX2HDRPipeline, LTX2ImageToVideoPipeline, LTX2InContextPipeline, diff --git a/src/diffusers/models/transformers/transformer_ltx2.py b/src/diffusers/models/transformers/transformer_ltx2.py index 27ac7acb89e6..755080198a7f 100644 --- a/src/diffusers/models/transformers/transformer_ltx2.py +++ b/src/diffusers/models/transformers/transformer_ltx2.py @@ -1115,8 +1115,8 @@ class LTX2VideoTransformer3DModel( for a given prompt. use_keyframes_abs_pos_embedding (`bool`, defaults to `False`): Whether to store a learned `(1, inner_dim)` absolute-position embedding for generated-keyframe tokens - (LTX-2.5.1+). When `True`, the weight is kept on the module for load/save; the regular distilled forward - path does not consume it until a dedicated keyframes pipeline wires it in. + (LTX-2.5). When `True`, tokens selected by `video_keyframes_mask` receive this embedding. The argument is + optional; omitting it leaves the distilled forward path unchanged. """ _supports_gradient_checkpointing = True @@ -1388,6 +1388,7 @@ def forward( use_cross_timestep: bool = False, attention_kwargs: dict[str, Any] | None = None, video_self_attention_mask: torch.Tensor | None = None, + video_keyframes_mask: torch.Tensor | None = None, return_dict: bool = True, ) -> torch.Tensor: """ @@ -1458,6 +1459,10 @@ def forward( applied to the video self-attention in each transformer block. Values in `[0, 1]` where `1` means full attention and `0` means masked. Used e.g. by the IC-LoRA pipeline to control attention strength between noisy tokens and appended reference tokens. Audio self-attention is not affected. + video_keyframes_mask (`torch.Tensor`, *optional*): + Optional per-token marker of shape `(batch_size, num_video_tokens, 1)`, non-zero on video tokens whose + latent frame encodes a single pixel frame. Those tokens receive `keyframes_abs_pos_embedding`. Ignored + when the model was built without `use_keyframes_abs_pos_embedding`. return_dict (`bool`, *optional*, defaults to `True`): Whether to return a dict-like structured output of type `AudioVisualModelOutput` or a tuple. @@ -1509,6 +1514,11 @@ def forward( hidden_states = self.proj_in(hidden_states) audio_hidden_states = self.audio_proj_in(audio_hidden_states) + # 2.1. Mark tokens whose latent encodes a single pixel frame (causal first frame, generated keyframe slots). + if self.config.use_keyframes_abs_pos_embedding and video_keyframes_mask is not None: + marker = (video_keyframes_mask > 0).to(dtype=hidden_states.dtype) + hidden_states = hidden_states + marker * self.keyframes_abs_pos_embedding.to(dtype=hidden_states.dtype) + # 3. Prepare timestep embeddings and modulation parameters timestep_cross_attn_gate_scale_factor = ( self.config.cross_attn_timestep_scale_multiplier / self.config.timestep_scale_multiplier diff --git a/src/diffusers/pipelines/__init__.py b/src/diffusers/pipelines/__init__.py index ef1814bbebcb..6a4aa91d60bc 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -340,6 +340,7 @@ _import_structure["ltx2"] = [ "LTX2Pipeline", "LTX2ConditionPipeline", + "LTX2DFRPipeline", "LTX2HDRPipeline", "LTX2InContextPipeline", "LTX2ImageToVideoPipeline", @@ -802,6 +803,7 @@ ) from .ltx2 import ( LTX2ConditionPipeline, + LTX2DFRPipeline, LTX2HDRPipeline, LTX2ImageToVideoPipeline, LTX2InContextPipeline, diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py index fe5b563a3c39..59644fdc2060 100644 --- a/src/diffusers/pipelines/ltx2/__init__.py +++ b/src/diffusers/pipelines/ltx2/__init__.py @@ -28,6 +28,7 @@ _import_structure["latent_upsampler"] = ["LTX2LatentUpsamplerModel"] _import_structure["pipeline_ltx2"] = ["LTX2Pipeline"] _import_structure["pipeline_ltx2_condition"] = ["LTX2ConditionPipeline", "LTX2VideoCondition"] + _import_structure["pipeline_ltx2_dfr"] = ["LTX2DFRPipeline"] _import_structure["pipeline_ltx2_diffusion_decode"] = ["LTX2VideoDiffusionDecodePipeline"] _import_structure["pipeline_ltx2_hdr_lora"] = ["LTX2HDRPipeline", "LTX2HDRReferenceCondition"] _import_structure["pipeline_ltx2_ic_lora"] = ["LTX2InContextPipeline", "LTX2ReferenceCondition"] @@ -49,6 +50,7 @@ from .latent_upsampler import LTX2LatentUpsamplerModel from .pipeline_ltx2 import LTX2Pipeline from .pipeline_ltx2_condition import LTX2ConditionPipeline, LTX2VideoCondition + from .pipeline_ltx2_dfr import LTX2DFRPipeline from .pipeline_ltx2_diffusion_decode import LTX2VideoDiffusionDecodePipeline from .pipeline_ltx2_hdr_lora import LTX2HDRPipeline, LTX2HDRReferenceCondition from .pipeline_ltx2_ic_lora import LTX2InContextPipeline, LTX2ReferenceCondition diff --git a/src/diffusers/pipelines/ltx2/dfr_layout.py b/src/diffusers/pipelines/ltx2/dfr_layout.py new file mode 100644 index 000000000000..d1b0384c7bd3 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/dfr_layout.py @@ -0,0 +1,463 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Canvas layout for the DFR pipeline: keyframe segment grid, temporal tile plan, epilogue tiling, and the token plan a +tiled transformer call walks.""" + +import itertools +from collections.abc import Sequence +from typing import NamedTuple + +import torch + +from ...utils import logging + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + +# Candidate keyframe segment lengths, in *latent* frames -- 24 and 32 pixel frames on the LTX-2.5 VAE. The grid picks +# whichever pads the request least. A segment must be a whole number of latent frames because every keyframe sits on a +# latent border. +SEGMENT_LATENT_CANDIDATES = (3, 4) + +# Latent-grid overlap between the spatial epilogue's height and width tiles. Its tiles agree on their overlaps at +# every denoising step, so the overlap only has to be wide enough to blend the seam away. +EPILOGUE_SPATIAL_OVERLAP = 12 + + +class DimensionInterval(NamedTuple): + """ + One tile's extent along a single axis, in that axis' own units. + + `start` / `end` are half-open. `left_ramp` / `right_ramp` are the lengths of the regions at each end that overlap + the neighbouring tile. How a ramp is resolved is the caller's choice: the temporal tile plan cuts on keyframe seams + and drops the ramp outright (both tiles reproduce a known frame there, so averaging only smears it), while the + spatial epilogue blends its ramps with a trapezoidal mask. + """ + + start: int + end: int + left_ramp: int + right_ramp: int + + +class LTX2DFRTemporalTile(NamedTuple): + """ + One temporal-refine window: its latent interval plus the pixel-frame keyframes it carries. + + `pixel_start` / `pixel_end` are inclusive pixel-frame bounds of `interval`. `anchors` are the seam keyframes inside + the window, carried in from the previous round; on every tile but the first this includes `pixel_start` itself. + `slots` are the mid-segment positions this window invents. + """ + + interval: DimensionInterval + pixel_start: int + pixel_end: int + anchors: tuple[int, ...] + slots: tuple[int, ...] + + +class LTX2DFREpilogueTile(NamedTuple): + """ + One spatial-epilogue tile: its extent on each latent axis plus the window its prediction is weighted by. + + `frames` / `heights` / `widths` are half-open slices into the latent grid. `blend_weight` is the separable `(F, H, + W)` window over that extent; adjacent tiles' windows sum to exactly one, so accumulating `tile * blend_weight` + reconstructs the canvas without a normalization pass. + """ + + frames: slice + heights: slice + widths: slice + blend_weight: torch.Tensor + + +class LTX2DFRTokenPlan(NamedTuple): + """ + One tile's slice of a packed token sequence, ready for a tiled transformer call. + + `keep` are the token indices the tile processes, base-grid tokens first. `weights` carries one blend weight per + kept token. `coords` are those tokens' positions, rebased on the tile's own first base token. + """ + + keep: torch.Tensor + weights: torch.Tensor + coords: torch.Tensor + + +def choose_segment_length(content_frames: int, temporal_compression_ratio: int = 8) -> int: + """ + Pick the keyframe segment length, in pixel frames, that pads `content_frames` least. + + Args: + content_frames (`int`): + `num_frames - 1`, the frame count the segment grid has to cover. + temporal_compression_ratio (`int`, defaults to `8`): + The VAE's temporal compression ratio. Candidates are `SEGMENT_LATENT_CANDIDATES` scaled by it, so the + shipped LTX-2.5 VAE offers 24 and 32 pixel frames. + + Returns: + `int`: the chosen segment length in pixel frames. Ties keep the larger segment. + """ + candidates = [candidate * temporal_compression_ratio for candidate in SEGMENT_LATENT_CANDIDATES] + return max(candidates, key=lambda candidate: (-((candidate - content_frames % candidate) % candidate), candidate)) + + +def resolve_canvas(num_frames: int, temporal_compression_ratio: int = 8) -> tuple[int, int, list[int]]: + """ + Pad `num_frames - 1` up to a multiple of the keyframe segment length. + + Args: + num_frames (`int`): + Requested pixel frame count. Must satisfy `(num_frames - 1) % temporal_compression_ratio == 0` and be at + least `temporal_compression_ratio + 1`. + temporal_compression_ratio (`int`, defaults to `8`): + The VAE's temporal compression ratio. + + Returns: + `tuple[int, int, list[int]]`: the padded frame count, the chosen segment length, and the keyframe positions + `[S, 2S, ..., N' - 1]`. Frame 0 is excluded (under causal encoding its latent already covers a single pixel + frame) and the terminal frame is included. + """ + if (num_frames - 1) % temporal_compression_ratio != 0: + raise ValueError( + f"`num_frames` must satisfy (num_frames - 1) % {temporal_compression_ratio} == 0, got {num_frames}" + ) + content = num_frames - 1 + if content < temporal_compression_ratio: + raise ValueError(f"The DFR canvas needs at least {temporal_compression_ratio + 1} pixel frames") + + segment = choose_segment_length(content, temporal_compression_ratio) + content_padded = content + (segment - content % segment) % segment + positions = [segment * index for index in range(1, content_padded // segment + 1)] + return content_padded + 1, segment, positions + + +def pixel_to_latent_index(pixel_frame: int, temporal_compression_ratio: int = 8) -> int: + """Map a pixel frame sitting on a latent border to its latent index.""" + if pixel_frame < 0: + raise ValueError(f"`pixel_frame` must be >= 0, got {pixel_frame}") + if pixel_frame % temporal_compression_ratio != 0: + raise ValueError(f"Pixel frame {pixel_frame} is not on the x{temporal_compression_ratio} latent border") + return pixel_frame // temporal_compression_ratio + + +def split_canvas_at_seams( + seams: Sequence[int], num_tiles: int, overlap: int, dim_size: int +) -> list[DimensionInterval]: + """ + Split a canvas on keyframe boundary cells, tolerating a remainder segment count. + + Each tile but the first starts `overlap` cells before the boundary it resumes after. That run-up is context only: + it lands in the interval's `left_ramp`, and the temporal rounds drop it, so the earlier tile keeps the boundary + cell and this one contributes strictly after it. + + Args: + seams (`Sequence[int]`): + The `K + 1` segment edges in grid cells, starting at `0` and ending at `dim_size - 1`. + num_tiles (`int`): + Requested number of tiles, clamped to the segment count `K`. + overlap (`int`): + Run-up each non-first tile reaches back, in grid cells. + dim_size (`int`): + Length of the axis being split, in grid cells. + """ + seams = tuple(int(seam) for seam in seams) + if overlap < 0: + raise ValueError(f"`overlap` must be >= 0, got {overlap}") + if len(seams) < 2 or seams[0] != 0: + raise ValueError(f"`seams` must start at 0 and hold at least one segment, got {list(seams)}") + if any(later <= earlier for earlier, later in itertools.pairwise(seams)): + raise ValueError(f"`seams` must be strictly increasing, got {list(seams)}") + if seams[-1] != dim_size - 1: + raise ValueError(f"`seams` must end at the last cell ({dim_size - 1}), got {seams[-1]}") + + if num_tiles < 1: + raise ValueError(f"`num_tiles` must be >= 1, got {num_tiles}") + # Deal the segments out so the leftovers land on the leading tiles. A tile cannot own zero segments, so more tiles + # than segments collapses to one tile per segment. + num_segments = len(seams) - 1 + num_tiles = min(num_tiles, num_segments) + base, remainder = divmod(num_segments, num_tiles) + counts = [base + (1 if index < remainder else 0) for index in range(num_tiles)] + + intervals = [] + cursor = 0 + for tile_index, count in enumerate(counts): + resume = seams[cursor] + 1 + start = 0 if tile_index == 0 else max(0, resume - overlap) + cursor += count + intervals.append( + DimensionInterval( + start=start, + end=seams[cursor] + 1, + left_ramp=0 if tile_index == 0 else resume - start, + right_ramp=0, + ) + ) + return intervals + + +def temporal_tile_plan( + seam_positions: Sequence[int], + num_frames: int, + num_tiles: int, + temporal_compression_ratio: int = 8, +) -> list[LTX2DFRTemporalTile]: + """ + Partition one temporal-refine round's canvas into keyframe-seam tiles. + + The overlap is one canvas segment plus the shared seam cell, so a tile's local latent 0 is the seam it inherits and + its lead-in covers the image latent the keyframe-at-0 lock produces -- neither may be spliced into the mid-canvas + stream, and both fall inside `interval.left_ramp`. + + Args: + seam_positions (`Sequence[int]`): + Keyframe pixel positions on this round's grid, strictly increasing, ending on `num_frames - 1`. + num_frames (`int`): + Pixel frame count of this round's canvas. + num_tiles (`int`): + Requested number of tiles, clamped to the segment count. + temporal_compression_ratio (`int`, defaults to `8`): + The VAE's temporal compression ratio. + """ + seams = [0, *(pixel_to_latent_index(position, temporal_compression_ratio) for position in seam_positions)] + latent_length = (num_frames - 1) // temporal_compression_ratio + 1 + overlap = seams[1] - seams[0] + 1 + tiles = [] + for interval in split_canvas_at_seams(seams, num_tiles, overlap, latent_length): + pixel_start = interval.start * temporal_compression_ratio + pixel_end = (interval.end - 1) * temporal_compression_ratio + anchors = tuple(position for position in seam_positions if pixel_start <= position <= pixel_end) + marks = [pixel_start, *(position for position in seam_positions if pixel_start < position <= pixel_end)] + slots = tuple((left + right) // 2 for left, right in itertools.pairwise(marks)) + tiles.append(LTX2DFRTemporalTile(interval, pixel_start, pixel_end, anchors, slots)) + return tiles + + +def split_by_count(dim_size: int, num_tiles: int, overlap: int) -> list[DimensionInterval]: + """ + Split an axis into `num_tiles` evenly sized tiles sharing `overlap` cells with each neighbour. + + Leading tiles absorb the remainder, so the tiles cover `[0, dim_size)` exactly and every adjacent pair shares + precisely `overlap` cells -- which is what makes `trapezoidal_mask_1d` sum to one across the seam. + """ + if num_tiles == 1: + return [DimensionInterval(start=0, end=dim_size, left_ramp=0, right_ramp=0)] + + total = dim_size + overlap * (num_tiles - 1) + tile_size, remainder = divmod(total, num_tiles) + if tile_size <= overlap: + raise ValueError( + f"Tile size {tile_size} is not larger than the overlap {overlap} for dim_size={dim_size}, " + f"num_tiles={num_tiles}" + ) + stride = tile_size - overlap + + intervals = [] + for index in range(num_tiles): + shift = min(index, remainder) + grow = 1 if index < remainder else 0 + intervals.append( + DimensionInterval( + start=index * stride + shift, + end=index * stride + tile_size + shift + grow, + left_ramp=0 if index == 0 else overlap, + right_ramp=0 if index == num_tiles - 1 else overlap, + ) + ) + return intervals + + +def trapezoidal_mask_1d(length: int, left_ramp: int, right_ramp: int) -> torch.Tensor: + """ + Build a `(length,)` blending weight that fades in over `left_ramp` cells and out over `right_ramp` cells. + + The ramps are the interior of a `linspace`, so two tiles sharing `k` cells contribute `i / (k + 1)` and `(k + 1 - + i) / (k + 1)` there and their weights sum to exactly one. + """ + mask = torch.ones(length) + if left_ramp > 0: + mask[:left_ramp] *= torch.linspace(0.0, 1.0, left_ramp + 2)[1:-1] + if right_ramp > 0: + mask[-right_ramp:] *= torch.linspace(1.0, 0.0, right_ramp + 2)[1:-1] + return mask + + +def rectangular_mask_1d(length: int, left_ramp: int) -> torch.Tensor: + """ + Build a `(length,)` blending weight that drops `left_ramp` cells outright and keeps the rest at full weight. + + This is the seam cut's counterpart to `trapezoidal_mask_1d`. Where a ramp is what two tiles need when neither knows + the truth at the border, a seam is a cell both of them reproduce from the same keyframe, so the later tile discards + its run-up instead of averaging it into the earlier tile's answer. + """ + mask = torch.ones(length) + mask[:left_ramp] = 0.0 + return mask + + +def epilogue_tiles( + latent_shape: tuple[int, int, int], frame_tiles: int, frame_seams: Sequence[int] +) -> list[LTX2DFREpilogueTile]: + """ + Lay the spatial detailing epilogue's `(frames, height, width)` tiling over a latent grid. + + Height and width are always split in two with `EPILOGUE_SPATIAL_OVERLAP` cells of trapezoidal overlap. The temporal + axis is cut on `frame_seams` instead: those cells carry a keyframe the last refine round already settled, so both + neighbours reproduce the same content there and the later tile drops its run-up rather than averaging it into the + earlier tile's answer. The run-up itself is one segment plus the shared seam cell, the same handover the refine + rounds use. Seams that are not interior to the grid leave the axis blended. + + An axis too short to hold its tile count falls back to a single tile, and an overlap it cannot hold is clamped, + since evenly sized tiles need `overlap < tile_size`. + + Returns one `(frame_slice, height_slice, width_slice, blend_weight)` per tile. `blend_weight` is the separable + window over the tile's `(F, H, W)` extent; adjacent windows sum to exactly one, so accumulating `tile * + blend_weight` reconstructs the canvas without a normalization pass. + """ + frame_size = latent_shape[0] + frame_overlap = frame_seams[0] + 1 if frame_tiles > 1 and frame_seams else 0 + + axis_configs = [] + for dim_size, num_tiles, overlap, axis in zip( + latent_shape, + (frame_tiles, 2, 2), + (frame_overlap, EPILOGUE_SPATIAL_OVERLAP, EPILOGUE_SPATIAL_OVERLAP), + ("frames", "height", "width"), + ): + if num_tiles > 1 and dim_size < num_tiles: + logger.warning( + f"Spatial epilogue: latent {axis} is {dim_size} cells, too short for {num_tiles} tiles; running " + f"this axis untiled." + ) + num_tiles, overlap = 1, 0 + elif num_tiles > 1 and overlap > dim_size - num_tiles: + logger.warning( + f"Spatial epilogue: latent {axis} is {dim_size} cells, so an overlap of {overlap} leaves no room " + f"for {num_tiles} tiles; clamping the overlap to {dim_size - num_tiles}." + ) + overlap = dim_size - num_tiles + axis_configs.append((dim_size, num_tiles, overlap)) + + _, frame_num_tiles, frame_overlap = axis_configs[0] + interior_seams = sorted({cell for cell in frame_seams if 0 < cell < frame_size - 1}) + seam_cut = frame_num_tiles > 1 and bool(interior_seams) + if seam_cut: + frame_intervals = split_canvas_at_seams( + [0, *interior_seams, frame_size - 1], frame_num_tiles, frame_overlap, frame_size + ) + else: + if frame_seams and frame_num_tiles > 1: + logger.warning( + f"Spatial epilogue: seams {list(frame_seams)} are not interior to {frame_size} latent frames; " + f"blending the temporal tiles instead of cutting them." + ) + frame_intervals = split_by_count(frame_size, frame_num_tiles, frame_overlap) + + tiles = [] + for frames, heights, widths in itertools.product( + frame_intervals, *(split_by_count(*config) for config in axis_configs[1:]) + ): + frame_mask = ( + rectangular_mask_1d(frames.end - frames.start, frames.left_ramp) + if seam_cut + else trapezoidal_mask_1d(frames.end - frames.start, frames.left_ramp, frames.right_ramp) + ) + height_mask = trapezoidal_mask_1d(heights.end - heights.start, heights.left_ramp, heights.right_ramp) + width_mask = trapezoidal_mask_1d(widths.end - widths.start, widths.left_ramp, widths.right_ramp) + tiles.append( + LTX2DFREpilogueTile( + frames=slice(frames.start, frames.end), + heights=slice(heights.start, heights.end), + widths=slice(widths.start, widths.end), + blend_weight=frame_mask[:, None, None] * height_mask[None, :, None] * width_mask[None, None, :], + ) + ) + return tiles + + +def video_tile_plan( + tiles: list[LTX2DFREpilogueTile], + video_coords: torch.Tensor, + latent_num_frames: int, + latent_height: int, + latent_width: int, +) -> list[LTX2DFRTokenPlan]: + """ + Resolve `tiles` into the per-tile token plan a tiled transformer call needs. + + A tile takes the base-grid tokens inside its `(F, H, W)` window plus every appended conditioning token whose RoPE + interval overlaps that window on all three axes -- token-level filtering, so a keyframe or reference token reaches + exactly the tiles whose picture it describes. + + Positions are rebased on the tile's own first base token, matching what a standalone pass over that crop would have + built. Base tokens carry the tile's separable blend window; a conditioning token kept by `n` tiles carries `1 / n`, + so both groups sum to one across the tiling and the blended prediction needs no normalization pass. + + Returns one dict per tile: `keep` (token indices, base tokens first), `weights` (one per kept token) and `coords` + (the kept tokens' rebased positions). + """ + tokens_per_latent_frame = latent_height * latent_width + num_base_tokens = latent_num_frames * tokens_per_latent_frame + device = video_coords.device + condition_coords = video_coords[:, :, num_base_tokens:, :] + + base_indices, tile_starts, kept_conditions = [], [], [] + for frames, heights, widths, _ in tiles: + frame_range = torch.arange(frames.start, frames.stop, device=device) + height_range = torch.arange(heights.start, heights.stop, device=device) + width_range = torch.arange(widths.start, widths.stop, device=device) + base = ( + frame_range[:, None, None] * tokens_per_latent_frame + + height_range[None, :, None] * latent_width + + width_range[None, None, :] + ).reshape(-1) + base_indices.append(base) + + # The window's extent in RoPE units, read off the tokens it owns rather than recomputed, so it stays in + # step with however `prepare_video_coords` laid the canvas out. + base_coords = video_coords[:, :, base, :] + starts = base_coords[..., 0].amin(dim=2) + ends = base_coords[..., 1].amax(dim=2) + tile_starts.append(starts) + overlaps = (condition_coords[..., 0] < ends[..., None]) & (condition_coords[..., 1] > starts[..., None]) + kept_conditions.append(overlaps.all(dim=1).any(dim=0)) + + condition_keepers = torch.stack(kept_conditions).sum(dim=0) + if not bool((condition_keepers > 0).all()): + unclaimed = (condition_keepers == 0).nonzero(as_tuple=False).squeeze(1).tolist() + raise ValueError( + f"Conditioning tokens {unclaimed} fall outside every tile, so no tile would denoise them. Their RoPE " + f"position lies off the canvas the tiling covers." + ) + + plan = [] + for tile, base, starts, keep_condition in zip(tiles, base_indices, tile_starts, kept_conditions): + condition_indices = num_base_tokens + keep_condition.nonzero(as_tuple=False).squeeze(1) + keep = torch.cat([base, condition_indices]) + weights = torch.cat( + [ + tile.blend_weight.reshape(-1).to(device=device), + 1.0 / condition_keepers[keep_condition].to(device=device, dtype=tile.blend_weight.dtype), + ] + ) + plan.append( + LTX2DFRTokenPlan( + keep=keep, + weights=weights, + coords=video_coords[:, :, keep, :] - starts[..., None, None], + ) + ) + return plan diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py new file mode 100644 index 000000000000..08ef3497ffbf --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py @@ -0,0 +1,2458 @@ +# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import copy +import math +from typing import Any, Callable + +import numpy as np +import PIL.Image +import torch +from transformers import ( + Gemma3ForConditionalGeneration, + Gemma4ForConditionalGeneration, + Gemma4UnifiedForConditionalGeneration, + GemmaTokenizer, + GemmaTokenizerFast, + ProcessorMixin, +) + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...image_processor import PipelineImageInput +from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin +from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.transformers import LTX2VideoTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import is_torch_xla_available, logging, replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..pipeline_utils import DiffusionPipeline +from .connectors import LTX2TextConnectors +from .dfr_layout import ( + epilogue_tiles, + pixel_to_latent_index, + resolve_canvas, + temporal_tile_plan, + video_tile_plan, +) +from .duration_head import LTX2DurationHead +from .latent_upsampler import LTX2LatentUpsamplerModel +from .pipeline_ltx2_condition import LTX2VideoCondition +from .pipeline_output import LTX2PipelineOutput +from .prompt_enhancement import ( + _pad_inputs_for_attention_alignment, + _prepare_enhance_image, + clean_response, +) +from .utils import ( + DISTILLED_SIGMA_VALUES, + GEMMA3_PROMPT_ENHANCEMENT_CONFIG, + GEMMA4_PROMPT_ENHANCEMENT_CONFIG, + LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT, + LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT, + STAGE_2_DISTILLED_SIGMA_VALUES, + TEMPORAL_ROUND_DISTILLED_SIGMA_VALUES, + apply_image_conditioning_crf, + resolve_default_image_crf, +) +from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE + + +if is_torch_xla_available(): + import torch_xla.core.xla_model as xm + + XLA_AVAILABLE = True +else: + XLA_AVAILABLE = False + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +# Keyframes carried between temporal rounds are pinned just short of fully clean so a tile can still settle its seam +# frame. +ANCHOR_KEYFRAME_STRENGTH = 0.95 + +# Ancestral noise fraction used by the temporal refine rounds. Their short schedule densifies detail rather than +# building structure, so a partly stochastic step is what fills in the freshly interpolated frames. +TEMPORAL_ANCESTRAL_ETA = 0.5 + +# RoPE time is `pixel_frame / fps`. The transformer is trained around 24/25/30 and 60 fps, not 48 or 120. A 120 fps +# time base halves every token's temporal span and the model can no longer lay out the VAE's pixel frames inside one +# latent -- it decodes as a motion spike at each latent border followed by a stall. 48 fps stretches the same span the +# other way. Condition at 60 in both cases and treat the decoded frames at the playback rate (120 fps: generate 2x +# frames at 60, mux as 120; 48 fps: generate the 24x2 canvas at 60, mux as 48). Playback fps is used for the returned +# frame count and the audio trim only. +MAX_CONDITIONING_FPS = 60.0 +SNAP_CONDITIONING_FPS_ABOVE = 30.0 + + +def _conditioning_fps(playback_fps: float) -> float: + """Fps the transformer sees. Playback may be 48, 96, 120, ...; those rates only affect muxing.""" + return MAX_CONDITIONING_FPS if playback_fps > SNAP_CONDITIONING_FPS_ABOVE else playback_fps + + +# The epilogue's keyframes arrive as finished frames, rebuilt at the output resolution, so they are pinned fully clean. +EPILOGUE_KEYFRAME_STRENGTH = 1.0 + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTX2DFRPipeline + >>> from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel + >>> from diffusers.utils import encode_video + + >>> # The published `model_index.json` has no upsampler; load it from its subfolder and pass it in. + >>> latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + ... "Lightricks/LTX-2.5-Diffusers", subfolder="latent_upsampler", torch_dtype=torch.bfloat16 + ... ) + >>> pipe = LTX2DFRPipeline.from_pretrained( + ... "Lightricks/LTX-2.5-Diffusers", latent_upsampler=latent_upsampler, torch_dtype=torch.bfloat16 + ... ) + >>> pipe.enable_model_cpu_offload() + + >>> frame_rate = 24.0 + >>> video, audio = pipe( + ... prompt="A tabby cat stretching in a sunlit window, dust motes drifting in the light", + ... height=704, + ... width=1216, + ... num_frames=121, + ... frame_rate=frame_rate, + ... output_type="np", + ... return_dict=False, + ... ) + + >>> encode_video( + ... video[0], + ... fps=frame_rate, + ... audio=audio[0].float().cpu(), + ... audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + ... output_path="dfr_output.mp4", + ... ) + ``` +""" + + +# Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img.retrieve_latents +def retrieve_latents( + encoder_output: torch.Tensor, generator: torch.Generator | None = None, sample_mode: str = "sample" +): + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +def _audio_window_for_tile( + audio_latents: torch.Tensor, + pixel_start: int, + tile_frames: int, + playback_fps: float, + source_seconds: float, + conditioning_fps: float, + audio_latents_per_second: float, +) -> torch.Tensor: + """ + Cut the frozen stage-1 audio to one temporal tile's window and resample it to that tile's token count. + + The window is wall clock: `pixel_start / playback_fps` through `(pixel_start + tile_frames) / playback_fps`, as a + fraction of `source_seconds`. Taking a fraction of the *canvas* instead would drift, because a refine round maps `N + -> 2 (N - 1) + 1` while the frame rate doubles, so each round's canvas is a hair shorter than twice the last one + and the tail tiles would pull audio from past their own playback. `conditioning_fps` only sizes the output token + count, matching what the video side asks the transformer for. + + Returns the packed `(batch_size, tile_audio_frames, channels * mel_bins)` window and its frame count. + """ + source_frames = audio_latents.shape[1] + tile_audio_frames = round(tile_frames / conditioning_fps * audio_latents_per_second) + start = pixel_start / playback_fps / source_seconds * source_frames + span = tile_frames / playback_fps / source_seconds * source_frames + positions = start + (span / tile_audio_frames) * torch.arange( + tile_audio_frames, device=audio_latents.device, dtype=torch.float32 + ) + positions = positions.clamp(0, source_frames - 1) + low = positions.floor().long() + high = (low + 1).clamp(max=source_frames - 1) + weight = (positions - low).to(audio_latents.dtype).view(1, -1, 1) + return audio_latents[:, low] * (1 - weight) + audio_latents[:, high] * weight + + +def ancestral_euler_step( + sample: torch.Tensor, + denoised: torch.Tensor, + sigma: torch.Tensor, + sigma_next: torch.Tensor, + eta: float, + noise: torch.Tensor, +) -> torch.Tensor: + """ + One ancestral (SDE) Euler step in the rectified-flow parameterization (`alpha = 1 - sigma`). + + The step advances deterministically to an intermediate noise level `sigma_down <= sigma_next` and then renoises + back up to `sigma_next`, rescaling the signal component by `alpha_next / alpha_down` so the transition stays + variance-preserving. `eta` interpolates between a plain Euler step (`eta=0`, `sigma_down == sigma_next`, no noise + added) and a fully ancestral one (`eta=1`). + + This is not expressible through [`FlowMatchEulerDiscreteScheduler`]: its `stochastic_sampling` option renoises all + the way from `x0` to `sigma_next` (the `eta=1` amount) and omits the variance-preserving rescale, so it gives a + different trajectory for every `eta` in `(0, 1)`. + """ + if sigma_next == 0: + return denoised.to(sample.dtype) + + sample = sample.to(torch.float32) + denoised = denoised.to(torch.float32) + + downstep_ratio = 1.0 + (sigma_next / sigma - 1.0) * eta + sigma_down = sigma_next * downstep_ratio + sigma_down_ratio = sigma_down / sigma + prev_sample = sigma_down_ratio * sample + (1.0 - sigma_down_ratio) * denoised + + if eta > 0: + alpha_next = 1.0 - sigma_next + alpha_down = 1.0 - sigma_down + renoise_coeff = (sigma_next**2 - sigma_down**2 * alpha_next**2 / alpha_down**2).clamp(min=0) ** 0.5 + prev_sample = (alpha_next / alpha_down) * prev_sample + noise.to(torch.float32) * renoise_coeff + return prev_sample + + +class LTX2DFRPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): + r""" + Pipeline for Diffusion Fidelity Rendering (DFR) with LTX-2.5: generated keyframe slots, a spatial detailing pass, + and optional temporal 2x/4x refinement. + + Stage 1 generates video *and* extra single-pixel-frame keyframe slots at a fraction of the requested resolution, + placing the slots on a segment grid aligned to the VAE's temporal border. Slots relax the effective temporal + compression at those positions, so the surrounding video can be conditioned on genuinely new frames rather than + interpolated ones. The half-resolution result is kept as an IC-LoRA reference while both the video and the slot + keyframes are upsampled in latent space. Stage 2 re-denoises at twice stage 1's resolution, seeded from the + upsampled latents and re-attaching the upsampled slots. + + `spatial_upscalings` chooses how many x2 spatial stages separate the base canvas from the returned frames. `1` is + the two-stage recipe above, with stage 2 already at the output resolution. `2` starts stage 1 a further factor of + two down and, *after* the temporal rounds, adds a third detailing pass at the output resolution -- tiled two ways + spatially and `2 ** temporal_upscalings` ways temporally, because a full-resolution pass over the whole refined + canvas does not fit in one go. Spatial tiles are blended over their overlaps; temporal tiles are cut on the last + refine round's keyframe seams, exactly where that round stitched. + + Audio comes from **stage 1**. Stage 2 still runs a joint audio pass so the video branch has cross-modal attention; + temporal refine tiles freeze that stage-1 audio (sigma 0, shared across tiles) so newly densified frames can follow + the same speech without each tile re-denoising a different audio realization. The shipped waveform is still stage + 1's. + + `temporal_upscalings` (0-2) adds rounds that each double the frame rate: the canvas is upsampled temporally, split + into `2 ** round` tiles that meet at shared keyframes, given fresh mid-segment slots, and densified with ancestral + Euler. Each tile cross-attends to the slice of the frozen stage-1 audio covering its own playback window. Whatever + padding the canvas needs internally, the caller always gets `(num_frames - 1) * 2 ** temporal_upscalings + 1` + frames back. + + Requires a transformer whose config sets `use_keyframes_abs_pos_embedding` (LTX-2.5 and later): each slot costs a + full latent frame of tokens to buy one pixel frame, so a checkpoint without the learned marker would spend that + budget on tokens it cannot interpret. + + Reference: https://github.com/Lightricks/LTX-2 + + Args: + scheduler ([`FlowMatchEulerDiscreteScheduler`]): + A scheduler to be used in combination with `transformer` to denoise the encoded video latents. + vae ([`AutoencoderKLLTX2Video`]): + Variational Auto-Encoder (VAE) Model to encode and decode videos to and from latent representations. + audio_vae ([`AutoencoderKLLTX2Audio`]): + Audio VAE to encode and decode audio spectrograms. + text_encoder ([`Gemma3ForConditionalGeneration`] or [`Gemma4UnifiedForConditionalGeneration`]): + Text encoder model. + tokenizer (`GemmaTokenizer` or `GemmaTokenizerFast`): + Tokenizer for the text encoder. + connectors ([`LTX2TextConnectors`]): + Text connector stack used to adapt text encoder hidden states for the video and audio branches. + transformer ([`LTX2VideoTransformer3DModel`]): + Conditional Transformer architecture to denoise the encoded video latents. + vocoder ([`LTX2Vocoder`] or [`LTX2VocoderWithBWE`]): + Vocoder to convert mel spectrograms to audio waveforms. + latent_upsampler ([`LTX2LatentUpsamplerModel`]): + Spatial x2 latent upsampler applied between stage 1 and stage 2. + temporal_latent_upsampler ([`LTX2LatentUpsamplerModel`], *optional*): + Temporal x2 latent upsampler. Required when `temporal_upscalings > 0`. + processor (`ProcessorMixin`, *optional*): + Processor used for prompt enhancement chat templating. + prompt_enhancer ([`Gemma4ForConditionalGeneration`], *optional*): + Dedicated prompt enhancement model (LTX-2.5). + duration_head ([`LTX2DurationHead`], *optional*): + Predicts `num_frames` from the prompt embeddings when `num_frames` is not supplied. + """ + + model_cpu_offload_seq = ( + "prompt_enhancer->text_encoder->connectors->duration_head->transformer->latent_upsampler->" + "temporal_latent_upsampler->vae->audio_vae->vocoder" + ) + _optional_components = ["temporal_latent_upsampler", "processor", "prompt_enhancer", "duration_head"] + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: FlowMatchEulerDiscreteScheduler, + vae: AutoencoderKLLTX2Video, + audio_vae: AutoencoderKLLTX2Audio, + text_encoder: Gemma3ForConditionalGeneration | Gemma4UnifiedForConditionalGeneration, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + connectors: LTX2TextConnectors, + transformer: LTX2VideoTransformer3DModel, + vocoder: LTX2Vocoder | LTX2VocoderWithBWE, + latent_upsampler: LTX2LatentUpsamplerModel, + temporal_latent_upsampler: LTX2LatentUpsamplerModel | None = None, + processor: ProcessorMixin | None = None, + prompt_enhancer: Gemma4ForConditionalGeneration | None = None, + duration_head: LTX2DurationHead | None = None, + ): + super().__init__() + + self.register_modules( + vae=vae, + audio_vae=audio_vae, + text_encoder=text_encoder, + tokenizer=tokenizer, + connectors=connectors, + transformer=transformer, + vocoder=vocoder, + scheduler=scheduler, + latent_upsampler=latent_upsampler, + temporal_latent_upsampler=temporal_latent_upsampler, + processor=processor, + prompt_enhancer=prompt_enhancer, + duration_head=duration_head, + ) + + self.vae_spatial_compression_ratio = ( + self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32 + ) + self.vae_temporal_compression_ratio = ( + self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8 + ) + self.audio_vae_mel_compression_ratio = ( + self.audio_vae.mel_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.audio_vae_temporal_compression_ratio = ( + self.audio_vae.temporal_compression_ratio if getattr(self, "audio_vae", None) is not None else 4 + ) + self.transformer_spatial_patch_size = ( + self.transformer.config.patch_size if getattr(self, "transformer", None) is not None else 1 + ) + self.transformer_temporal_patch_size = ( + self.transformer.config.patch_size_t if getattr(self, "transformer", None) is not None else 1 + ) + + self.audio_sampling_rate = ( + self.audio_vae.config.sample_rate if getattr(self, "audio_vae", None) is not None else 16000 + ) + self.audio_hop_length = ( + self.audio_vae.config.mel_hop_length if getattr(self, "audio_vae", None) is not None else 160 + ) + self.audio_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64 + self.audio_latent_channels = ( + self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8 + ) + + self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio, resample="bilinear") + self.tokenizer_max_length = ( + self.tokenizer.model_max_length if getattr(self, "tokenizer", None) is not None else 1024 + ) + tokenizer_padding_side = "left" + if getattr(self, "tokenizer", None) is not None: + tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left") + self.tokenizer_padding_side = tokenizer_padding_side + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.check_inputs + def check_inputs( + self, + prompt, + height, + width, + callback_on_step_end_tensor_inputs=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + spatio_temporal_guidance_blocks=None, + stg_scale=None, + audio_stg_scale=None, + system_prompt=None, + enable_prompt_enhancement=None, + num_frames=None, + min_seconds=1.0, + max_seconds=20.0, + image=None, + image_crf=None, + latents=None, + audio_latents=None, + ): + if height % 32 != 0 or width % 32 != 0: + raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.") + + if callback_on_step_end_tensor_inputs is not None and not all( + k in self._callback_tensor_inputs for k in callback_on_step_end_tensor_inputs + ): + raise ValueError( + f"`callback_on_step_end_tensor_inputs` has to be in {self._callback_tensor_inputs}, but found {[k for k in callback_on_step_end_tensor_inputs if k not in self._callback_tensor_inputs]}" + ) + + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. Please make sure to" + " only forward one of the two." + ) + elif prompt is None and prompt_embeds is None: + raise ValueError( + "Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined." + ) + elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): + raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") + + if prompt_embeds is not None and prompt_attention_mask is None: + raise ValueError("Must provide `prompt_attention_mask` when specifying `prompt_embeds`.") + + if negative_prompt_embeds is not None and negative_prompt_attention_mask is None: + raise ValueError("Must provide `negative_prompt_attention_mask` when specifying `negative_prompt_embeds`.") + + if prompt_embeds is not None and negative_prompt_embeds is not None: + if prompt_embeds.shape != negative_prompt_embeds.shape: + raise ValueError( + "`prompt_embeds` and `negative_prompt_embeds` must have the same shape when passed directly, but" + f" got: `prompt_embeds` {prompt_embeds.shape} != `negative_prompt_embeds`" + f" {negative_prompt_embeds.shape}." + ) + if prompt_attention_mask.shape != negative_prompt_attention_mask.shape: + raise ValueError( + "`prompt_attention_mask` and `negative_prompt_attention_mask` must have the same shape when passed directly, but" + f" got: `prompt_attention_mask` {prompt_attention_mask.shape} != `negative_prompt_attention_mask`" + f" {negative_prompt_attention_mask.shape}." + ) + + if latents is not None and latents.ndim != 5: + raise ValueError( + f"Only unpacked (5D) video latents of shape `[batch_size, latent_channels, latent_frames," + f" latent_height, latent_width] are supported, but got {latents.ndim} dims. If you have packed (3D)" + f" latents, please unpack them (e.g. using the `_unpack_latents` method)." + ) + if audio_latents is not None and audio_latents.ndim != 4: + raise ValueError( + f"Only unpacked (4D) audio latents of shape `[batch_size, num_channels, audio_length, mel_bins] are" + f" supported, but got {audio_latents.ndim} dims. If you have packed (3D) latents, please unpack them" + f" (e.g. using the `_unpack_audio_latents` method)." + ) + + if ((stg_scale > 0.0) or (audio_stg_scale > 0.0)) and not spatio_temporal_guidance_blocks: + raise ValueError( + "Spatio-Temporal Guidance (STG) is specified but no STG blocks are supplied. Please supply a list of" + "block indices at which to apply STG in `spatio_temporal_guidance_blocks`" + ) + + if ( + enable_prompt_enhancement + and prompt is not None + and system_prompt is None + and getattr(self, "prompt_enhancer", None) is None + ): + raise ValueError( + "`system_prompt` must be supplied to enable prompt enhancement when no dedicated " + "`prompt_enhancer` component is configured (LTX-2.0/2.3)." + ) + + if min_seconds >= max_seconds: + raise ValueError( + f"`min_seconds` ({min_seconds}) must be less than `max_seconds` ({max_seconds})." + " A collapsed range leaves no room for a prediction, and cannot generally be satisfied by a frame" + " count on the VAE's temporal grid." + ) + + # Auto-duration path: `num_frames` omitted on a pipeline that has a `duration_head`. + if num_frames is None and getattr(self, "duration_head", None) is not None: + num_prompts = len(prompt) if isinstance(prompt, list) else 1 if prompt is not None else len(prompt_embeds) + if num_prompts > 1: + raise ValueError( + f"`num_frames` was omitted so the duration head would auto-predict, but {num_prompts} prompts were" + " supplied. The duration head predicts one duration, and prompts with different natural lengths" + " cannot share a single frame count. Call the pipeline once per prompt, or pass `num_frames` as an" + " integer." + ) + + if latents is None and image is not None: + crf = image_crf if image_crf is not None else resolve_default_image_crf(self.text_encoder) + if crf != 0 and not isinstance(image, PIL.Image.Image): + raise ValueError( + f"`image_crf` re-compression requires a `PIL.Image.Image` input, got {type(image)}. " + "Pass a PIL image, or set `image_crf=0` to skip re-compression." + ) + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.enhance_prompt + def enhance_prompt( + self, + prompt: str, + system_prompt: str, + max_new_tokens: int | None = None, + seed: int = 10, + generator: torch.Generator | None = None, + generation_kwargs: dict[str, Any] | None = None, + device: str | torch.device | None = None, + image: PipelineImageInput | None = None, + ): + """ + Enhances the supplied `prompt` by generating a new prompt using the prompt enhancer (a Gemma + conditional-generation model) from it and a system prompt. When `image` is supplied, the enhancer is also + conditioned on that reference frame (I2V / keyframe-style enhancement). Uses the dedicated `prompt_enhancer` + component if one is configured (e.g. LTX-2.5, whose text encoder isn't trained for enhancement), otherwise + falls back to the main `text_encoder` (LTX-2.0/2.3, which double as their own enhancer). + + Message templates, decoding kwargs, response cleaning, and image long-side prep match `ltx-core` / + `ltx-pipelines` (`enhance_t2v` / `enhance_i2v` / `generate_enhanced_prompt`). + """ + device = device or self._execution_device + using_dedicated_enhancer = getattr(self, "prompt_enhancer", None) is not None + enhancer = self.prompt_enhancer if using_dedicated_enhancer else self.text_encoder + config = GEMMA4_PROMPT_ENHANCEMENT_CONFIG if using_dedicated_enhancer else GEMMA3_PROMPT_ENHANCEMENT_CONFIG + + generation_kwargs = ( + dict(generation_kwargs) if generation_kwargs is not None else dict(config.generation_kwargs) + ) + if max_new_tokens is None: + max_new_tokens = config.max_new_tokens + + # Templates match ltx-core `LTXGemmaTextEncoder.enhance_t2v` / `enhance_i2v` for both Gemma 3 and 4. + if image is None: + messages = [ + {"role": "system", "content": system_prompt}, + {"role": "user", "content": f"user prompt: {prompt}"}, + ] + enhance_image = None + else: + enhance_image = _prepare_enhance_image(image) + messages = [ + {"role": "system", "content": system_prompt}, + { + "role": "user", + "content": [ + {"type": "image"}, + {"type": "text", "text": f"User Raw Input Prompt: {prompt}."}, + ], + }, + ] + + template = self.processor.tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True) + model_inputs = self.processor(text=template, images=enhance_image, return_tensors="pt").to(device) + pad_token_id = ( + self.processor.tokenizer.pad_token_id if self.processor.tokenizer.pad_token_id is not None else 0 + ) + model_inputs = _pad_inputs_for_attention_alignment(model_inputs, pad_token_id=pad_token_id) + enhancer.to(device) + + # `transformers.GenerationMixin.generate` does not support using a `torch.Generator` to control randomness, + # so manually apply a seed for reproducible generation. + if generator is not None: + seed = generator.initial_seed() if not isinstance(generator, list) else generator[0].initial_seed() + torch.manual_seed(seed) + generated_sequences = enhancer.generate( + **model_inputs, + max_new_tokens=max_new_tokens, + **generation_kwargs, + ) # tensor of shape [batch_size, seq_len] + + generated_ids = [seq[len(model_inputs.input_ids[i]) :] for i, seq in enumerate(generated_sequences)] + enhanced_prompt = self.processor.tokenizer.batch_decode(generated_ids, skip_special_tokens=True) + return [clean_response(text) for text in enhanced_prompt] + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._get_gemma_prompt_embeds + def _get_gemma_prompt_embeds( + self, + prompt: str | list[str], + num_videos_per_prompt: int = 1, + max_sequence_length: int = 1024, + scale_factor: int = 8, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + device: (`str` or `torch.device`): + torch device to place the resulting embeddings on + dtype: (`torch.dtype`): + torch dtype to cast the prompt embeds to + max_sequence_length (`int`, defaults to 1024): Maximum sequence length to use for the prompt. + """ + device = device or self._execution_device + dtype = dtype or self.text_encoder.dtype + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) + + if getattr(self, "tokenizer", None) is not None: + # Gemma expects left padding for chat-style prompts + self.tokenizer.padding_side = "left" + if self.tokenizer.pad_token is None: + self.tokenizer.pad_token = self.tokenizer.eos_token + + prompt = [p.strip() for p in prompt] + text_inputs = self.tokenizer( + prompt, + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids + prompt_attention_mask = text_inputs.attention_mask + text_input_ids = text_input_ids.to(device) + prompt_attention_mask = prompt_attention_mask.to(device) + + text_encoder_outputs = self.text_encoder( + input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True + ) + text_encoder_hidden_states = text_encoder_outputs.hidden_states + text_encoder_hidden_states = torch.stack(text_encoder_hidden_states, dim=-1) + prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype) # Pack to 3D + + # duplicate text embeddings for each generation per prompt, using mps friendly method + _, seq_len, _ = prompt_embeds.shape + prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) + + prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) + prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) + + return prompt_embeds, prompt_attention_mask + + def encode_prompt( + self, + prompt: str | list[str], + num_videos_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + max_sequence_length: int = 1024, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ): + r""" + Encodes the prompt into text encoder hidden states. + + DFR runs the distilled sigma schedule, which is trained to be used without classifier-free guidance, so there + is no negative branch here. + + Args: + prompt (`str` or `list[str]`, *optional*): + prompt to be encoded + num_videos_per_prompt (`int`, *optional*, defaults to 1): + Number of videos that should be generated per prompt. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not + provided, text embeddings will be generated from `prompt` input argument. + prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for `prompt_embeds`. + device: (`torch.device`, *optional*): + torch device + dtype: (`torch.dtype`, *optional*): + torch dtype + """ + device = device or self._execution_device + + if prompt_embeds is None: + prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + max_sequence_length=max_sequence_length, + device=device, + dtype=dtype, + ) + + return prompt_embeds, prompt_attention_mask + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_latents + def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: + # Unpacked latents of shape are [B, C, F, H, W] are patched into tokens of shape [B, C, F // p_t, p_t, H // p, p, W // p, p]. + # The patch dimensions are then permuted and collapsed into the channel dimension of shape: + # [B, F // p_t * H // p * W // p, C * p_t * p * p] (an ndim=3 tensor). + # dim=0 is the batch size, dim=1 is the effective video sequence length, dim=2 is the effective number of input features + batch_size, num_channels, num_frames, height, width = latents.shape + post_patch_num_frames = num_frames // patch_size_t + post_patch_height = height // patch_size + post_patch_width = width // patch_size + latents = latents.reshape( + batch_size, + -1, + post_patch_num_frames, + patch_size_t, + post_patch_height, + patch_size, + post_patch_width, + patch_size, + ) + latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents + def _unpack_latents( + latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1 + ) -> torch.Tensor: + # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) + # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of + # what happens in the `_pack_latents` method. + batch_size = latents.size(0) + latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size) + latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2_image2video.LTX2ImageToVideoPipeline._normalize_latents + def _normalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Normalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = (latents - latents_mean) * scaling_factor / latents_std + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents + def _denormalize_latents( + latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 + ) -> torch.Tensor: + # Denormalize latents across the channel dimension [B, C, F, H, W] + latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype) + latents = latents * latents_std / scaling_factor + latents_mean + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._normalize_audio_latents + def _normalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents - latents_mean) / latents_std + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_audio_latents + def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): + latents_mean = latents_mean.to(latents.device, latents.dtype) + latents_std = latents_std.to(latents.device, latents.dtype) + return (latents * latents_std) + latents_mean + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_audio_latents + def _pack_audio_latents( + latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None + ) -> torch.Tensor: + # Audio latents shape: [B, C, L, M], where L is the latent audio length and M is the number of mel bins + if patch_size is not None and patch_size_t is not None: + # Packs the latents into a patch sequence of shape [B, L // p_t * M // p, C * p_t * p] (a ndim=3 tnesor). + # dim=1 is the effective audio sequence length and dim=2 is the effective audio input feature size. + batch_size, num_channels, latent_length, latent_mel_bins = latents.shape + post_patch_latent_length = latent_length / patch_size_t + post_patch_mel_bins = latent_mel_bins / patch_size + latents = latents.reshape( + batch_size, -1, post_patch_latent_length, patch_size_t, post_patch_mel_bins, patch_size + ) + latents = latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + else: + # Packs the latents into a patch sequence of shape [B, L, C * M]. This implicitly assumes a (mel) + # patch_size of M (all mel bins constitutes a single patch) and a patch_size_t of 1. + latents = latents.transpose(1, 2).flatten(2, 3) # [B, C, L, M] --> [B, L, C * M] + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_audio_latents + def _unpack_audio_latents( + latents: torch.Tensor, + latent_length: int, + num_mel_bins: int, + patch_size: int | None = None, + patch_size_t: int | None = None, + ) -> torch.Tensor: + # Unpacks an audio patch sequence of shape [B, S, D] into a latent spectrogram tensor of shape [B, C, L, M], + # where L is the latent audio length and M is the number of mel bins. + if patch_size is not None and patch_size_t is not None: + batch_size = latents.size(0) + latents = latents.reshape(batch_size, latent_length, num_mel_bins, -1, patch_size_t, patch_size) + latents = latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + else: + # Assume [B, S, D] = [B, L, C * M], which implies that patch_size = M and patch_size_t = 1. + latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) + return latents + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2_condition.LTX2ConditionPipeline.trim_conditioning_sequence + def trim_conditioning_sequence(self, start_frame: int, sequence_num_frames: int, target_num_frames: int) -> int: + """ + Trim a conditioning sequence to the allowed number of frames. + + Args: + start_frame (int): The target frame number of the first frame in the sequence. + sequence_num_frames (int): The number of frames in the sequence. + target_num_frames (int): The target number of frames in the generated video. + Returns: + int: updated sequence length + """ + scale_factor = self.vae_temporal_compression_ratio + num_frames = min(sequence_num_frames, target_num_frames - start_frame) + # Trim down to a multiple of temporal_scale_factor frames plus 1 + num_frames = (num_frames - 1) // scale_factor * scale_factor + 1 + return num_frames + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2_condition.LTX2ConditionPipeline.preprocess_conditions + def preprocess_conditions( + self, + conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, + height: int = 512, + width: int = 768, + num_frames: int = 121, + device: torch.device | None = None, + ) -> tuple[list[torch.Tensor], list[float], list[int], list[int]]: + """ + Preprocesses the condition images/videos to torch tensors. + + Args: + conditions (`LTX2VideoCondition` or `List[LTX2VideoCondition]`, *optional*, defaults to `None`): + A list of image/video condition instances. + height (`int`, *optional*, defaults to `512`): + The desired height in pixels. + width (`int`, *optional*, defaults to `768`): + The desired width in pixels. + num_frames (`int`, *optional*, defaults to `121`): + The desired number of frames in the generated video. + device (`torch.device`, *optional*, defaults to `None`): + The device on which to put the preprocessed image/video tensors. + + Returns: + `Tuple[List[torch.Tensor], List[float], List[int], List[int]]`: + Returns a 4-tuple of lists of length `len(conditions)` as follows: + 1. The first list is a list of preprocessed video tensors of shape [batch_size=1, num_channels, + num_frames, height, width]. + 2. The second list is a list of conditioning strengths. + 3. The third list is a list of latent-space indices for each condition. + 4. The fourth list is a list of (trimmed) pixel-space frame counts per condition. This is needed + for keyframe coord semantics (single-pixel-frame keyframes have a clamped temporal extent). + """ + conditioning_frames, conditioning_strengths, conditioning_indices, conditioning_pixel_frames = [], [], [], [] + + if conditions is None: + conditions = [] + if isinstance(conditions, LTX2VideoCondition): + conditions = [conditions] + + frame_scale_factor = self.vae_temporal_compression_ratio + latent_num_frames = (num_frames - 1) // frame_scale_factor + 1 + for i, condition in enumerate(conditions): + # Create a channels-last video-like array of shape (F, H, W, C) in preparation for resizing. + if isinstance(condition.frames, PIL.Image.Image): + arr = np.array(condition.frames.convert("RGB"))[None] # (1, H, W, 3) + elif isinstance(condition.frames, list) and all(isinstance(f, PIL.Image.Image) for f in condition.frames): + arr = np.stack([np.array(f.convert("RGB")) for f in condition.frames]) # (F, H, W, 3) + elif isinstance(condition.frames, np.ndarray): + arr = condition.frames if condition.frames.ndim == 4 else condition.frames[None] + elif isinstance(condition.frames, torch.Tensor): + t = condition.frames if condition.frames.ndim == 4 else condition.frames.unsqueeze(0) + # Reference layout for video tensors is (F, C, H, W); convert to (F, H, W, C) for the + # resize logic, which expects channels-last. + arr = t.detach().cpu().permute(0, 2, 3, 1).numpy() + else: + raise TypeError(f"Unsupported `frames` type for condition {i}: {type(condition.frames)}") + + # Single-frame image keyframes are H.264 re-compressed at the model CRF (ltx-pipelines + # `ImageConditioner.resolve_crf` + `media_io.preprocess`). Multi-frame video conditions are not. + if arr.shape[0] == 1: + crf = condition.crf if condition.crf is not None else resolve_default_image_crf(self.text_encoder) + if crf != 0 and arr.dtype != np.uint8: + raise ValueError( + f"Image conditioning CRF expects a uint8 RGB frame, got dtype={arr.dtype}. " + "Pass a PIL image / uint8 array, or set `crf=0` on the condition to skip re-compression." + ) + arr = apply_image_conditioning_crf(arr[0], crf)[None] + + src_h, src_w = arr.shape[1], arr.shape[2] + num_cond_frames = arr.shape[0] + # Convert the NumPy array to a channels-first tensor of shape (1, C, F, H, W) + pixels = torch.from_numpy(np.ascontiguousarray(arr)).to(torch.float32) + pixels = pixels.permute(3, 0, 1, 2).unsqueeze(0).to(device) # (1, C, F, H, W) + + # Resize so the longer side fills the target, then center-crop to exact (height, width). + scale = max(height / src_h, width / src_w) + new_h = math.ceil(src_h * scale) + new_w = math.ceil(src_w * scale) + # Flatten (B, C, F, H, W) → (B*F, C, H, W) for the per-frame interpolation + pixels = pixels.permute(0, 2, 1, 3, 4).reshape(num_cond_frames, 3, src_h, src_w) + # NOTE: we avoid using VideoProcessor.preprocess_video here because it uses PIL.Image.resize under the + # hood, which will apply an anti-aliasing pre-filter when downsampling. The original LTX-2.X code simply + # uses F.interpolate, which is reproduced here. + pixels = torch.nn.functional.interpolate(pixels, size=(new_h, new_w), mode="bilinear", align_corners=False) + top = (new_h - height) // 2 + left = (new_w - width) // 2 + pixels = pixels[:, :, top : top + height, left : left + width] + pixels = pixels.reshape(1, num_cond_frames, 3, height, width).permute(0, 2, 1, 3, 4) + + # Map [0, 255] → [-1, 1] (VAE input convention). + condition_pixels = pixels / 127.5 - 1.0 + + # Interpret the index as a latent index, following the original LTX-2 code. + latent_start_idx = condition.index + # Support negative latent indices (e.g. -1 for the last latent index) + if latent_start_idx < 0: + # latent_start_idx will be positive because latent_num_frames is positive + latent_start_idx = latent_start_idx % latent_num_frames + if latent_start_idx >= latent_num_frames: + logger.warning( + f"The starting latent index {latent_start_idx} of condition {i} is too big for the specified number" + f" of latent frames {latent_num_frames}. This condition will be skipped." + ) + continue + + cond_num_frames = condition_pixels.size(2) + start_idx = max((latent_start_idx - 1) * frame_scale_factor + 1, 0) + truncated_cond_frames = self.trim_conditioning_sequence(start_idx, cond_num_frames, num_frames) + condition_pixels = condition_pixels[:, :, :truncated_cond_frames] + + conditioning_frames.append(condition_pixels.to(dtype=self.vae.dtype, device=device)) + conditioning_strengths.append(condition.strength) + conditioning_indices.append(latent_start_idx) + conditioning_pixel_frames.append(truncated_cond_frames) + + return conditioning_frames, conditioning_strengths, conditioning_indices, conditioning_pixel_frames + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2_condition.LTX2ConditionPipeline._prepare_keyframe_coords + def _prepare_keyframe_coords( + self, + keyframe_latent_num_frames: int, + keyframe_latent_height: int, + keyframe_latent_width: int, + pixel_frame_idx: int, + num_pixel_frames: int, + fps: float, + device: torch.device, + ) -> torch.Tensor: + """ + Compute positional coordinates for a keyframe condition being appended as extra tokens. + + Mirrors `VideoConditionByKeyframeIndex.apply_to` in the reference implementation: + - Latent coords scaled to pixel space *without* the causal fix (since non-zero-index keyframes don't need the + first-frame causal adjustment). + - Temporal axis offset by `pixel_frame_idx` (the pixel-space index at which the keyframe appears). + - For single-pixel-frame keyframes, the per-patch temporal extent is clamped to `[idx, idx + 1)` so the + keyframe occupies a single pixel timestep rather than the VAE-scaled range. + - Temporal coords divided by `fps` to produce seconds. + """ + patch_size = self.transformer_spatial_patch_size + patch_size_t = self.transformer_temporal_patch_size + scale_factors = ( + self.vae_temporal_compression_ratio, + self.vae_spatial_compression_ratio, + self.vae_spatial_compression_ratio, + ) + + grid_f = torch.arange( + start=0, end=keyframe_latent_num_frames, step=patch_size_t, dtype=torch.float32, device=device + ) + grid_h = torch.arange(start=0, end=keyframe_latent_height, step=patch_size, dtype=torch.float32, device=device) + grid_w = torch.arange(start=0, end=keyframe_latent_width, step=patch_size, dtype=torch.float32, device=device) + grid = torch.meshgrid(grid_f, grid_h, grid_w, indexing="ij") + grid = torch.stack(grid, dim=0) + + patch_size_delta = torch.tensor((patch_size_t, patch_size, patch_size), dtype=grid.dtype, device=device) + patch_ends = grid + patch_size_delta.view(3, 1, 1, 1) + + latent_coords = torch.stack([grid, patch_ends], dim=-1) # [3, N_F, N_H, N_W, 2] + latent_coords = latent_coords.flatten(1, 3) # [3, num_patches, 2] + latent_coords = latent_coords.unsqueeze(0) # [1, 3, num_patches, 2] + + scale_tensor = torch.tensor(scale_factors, device=device, dtype=latent_coords.dtype) + broadcast_shape = [1] * latent_coords.ndim + broadcast_shape[1] = -1 + pixel_coords = latent_coords * scale_tensor.view(*broadcast_shape) + + # No causal fix: keyframe coords place the keyframe at `pixel_frame_idx` without the first-frame adjustment. + pixel_coords[:, 0, :, :] = pixel_coords[:, 0, :, :] + pixel_frame_idx + + if num_pixel_frames == 1: + # Single-pixel-frame keyframe: clamp temporal extent to [idx, idx + 1). + pixel_coords[:, 0, :, 1:] = pixel_coords[:, 0, :, :1] + 1 + + pixel_coords[:, 0, :, :] = pixel_coords[:, 0, :, :] / fps + + return pixel_coords + + def prepare_latents( + self, + conditions: list[LTX2VideoCondition] | None = None, + condition_latents: list[tuple[int, torch.Tensor, float, int]] | None = None, + keyframe_latents: list[tuple[int, torch.Tensor, float]] | None = None, + slot_frame_indices: list[int] | None = None, + slot_initial_latents: torch.Tensor | None = None, + reference_latents: torch.Tensor | None = None, + reference_downscale_factor: int = 1, + batch_size: int = 1, + num_channels_latents: int = 128, + height: int = 512, + width: int = 768, + num_frames: int = 121, + frame_rate: float = 24.0, + noise_scale: float = 1.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, slice | None]: + """ + Prepare the noisy packed video latents for one DFR denoising pass. + + The packed sequence is laid out as `[base | keyframes | slots | reference]`: + + - Base tokens cover the target latent grid, seeded from `latents` when supplied. + - Frame conditions with `index == 0` set the clean target at the first-frame positions; those with `index > 0` + and every entry of `keyframe_latents` are appended as extra keyframe tokens with a per-token conditioning + mask equal to their strength. + - `slot_frame_indices` appends one latent frame's worth of *generated* keyframe tokens per position, with + conditioning mask `0` (fully denoised) and a RoPE temporal extent of exactly one pixel frame. These are the + keyframe slots that give DFR its extra frames; `slot_initial_latents` seeds their content. + - `reference_latents` appends the stage-1 half-resolution latent as a fully clean IC-LoRA reference, with + spatial coordinates scaled by `reference_downscale_factor` so it maps into the target coordinate space. + + Appended conditioning tokens carry their content in `clean_latents` and a zero placeholder in `latents`, while + keyframe slots carry their seed in `latents` and zeros in `clean_latents` -- the returned `latents` are the + noised mix of the two (see the noising step at the end of this method). + + Args: + conditions (`list[LTX2VideoCondition]`, *optional*): + Frame-level image / video conditions, positioned by latent index. + condition_latents (`list[tuple[int, torch.Tensor, float, int]]`, *optional*): + Already-encoded stand-in for `conditions`, as `(pixel_frame_index, latent, strength, + num_pixel_frames)`. Pixel rather than latent index, because a temporal refine round scales a + condition's position by `2 ** round` and the result does not generally land on a latent boundary -- + only an appended keyframe token can sit there, and it is placed by pixel. `pixel_frame_index == 0` + still means "replace the first frame". + keyframe_latents (`list[tuple[int, torch.Tensor, float]]`, *optional*): + Already-encoded keyframe guidance as `(pixel_frame_index, latent, strength)`, where `latent` has shape + `(batch_size, num_channels_latents, 1, latent_height, latent_width)`. Used by the temporal refine + rounds to pin the seam keyframes carried in from the previous round. + slot_frame_indices (`list[int]`, *optional*): + Pixel-frame positions of the generated keyframe slots. + slot_initial_latents (`torch.Tensor`, *optional*): + `(batch_size, num_channels_latents, len(slot_frame_indices), latent_height, latent_width)` content + written into the slot tokens before noising. + reference_latents (`torch.Tensor`, *optional*): + Normalized `(batch_size, num_channels_latents, F, H, W)` IC-LoRA reference latent. + reference_downscale_factor (`int`, defaults to `1`): + Ratio between the target and the reference resolution. + latents (`torch.Tensor`, *optional*): + Normalized `(batch_size, num_channels_latents, F, H, W)` initial content for the base tokens. + noise_scale (`float`, defaults to `1.0`): + Noise level the unconditioned tokens are initialized at, i.e. the schedule's first sigma. + + Returns: + `tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, slice | None]`: + `(latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, slot_token_slice)`. + `slot_token_slice` indexes the generated keyframe slot tokens in the packed sequence, or `None` when no + slots were requested. + """ + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + patch_size = self.transformer_spatial_patch_size + patch_size_t = self.transformer_temporal_patch_size + + # `randn_tensor` draws per batch element from a list of generators, but the VAE's latent distribution is + # sampled once for a batch-1 condition tensor, so that call takes a single generator. + encode_generator = generator[0] if isinstance(generator, list) else generator + + if latents is None: + # NOTE: zeros rather than a Gaussian sample, because the per-token noise level is only known once the + # conditioning mask below is complete. + latents = torch.zeros( + (batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width), + device=device, + dtype=dtype, + ) + latents = self._pack_latents(latents.to(device=device, dtype=dtype), patch_size, patch_size_t) + conditioning_mask = latents.new_zeros((*latents.shape[:2], 1)) + clean_latents = torch.zeros_like(latents) + + # Frame conditions: encode each one, then either overwrite the first-frame tokens or queue it as a keyframe. + if condition_latents is None: + condition_latents = self.encode_conditions( + conditions, height, width, num_frames, device=device, dtype=dtype, generator=encode_generator + ) + appended_keyframes: list[tuple[torch.Tensor, torch.Tensor, float]] = [] + for pixel_frame_idx, encoded, strength, num_pixel_frames in condition_latents: + encoded = encoded.to(device=device, dtype=dtype) + # Conditions are preprocessed as batch-1 tensors and shared across the batch. + condition_tokens = self._pack_latents(encoded, patch_size, patch_size_t).expand(batch_size, -1, -1) + + if pixel_frame_idx == 0: + # Overwrite the clean target and mask only; the noisy sequence keeps the base-grid seed. + num_condition_tokens = condition_tokens.shape[1] + conditioning_mask[:, :num_condition_tokens] = strength + clean_latents[:, :num_condition_tokens] = condition_tokens + continue + + coords = self._prepare_keyframe_coords( + keyframe_latent_num_frames=encoded.shape[2], + keyframe_latent_height=encoded.shape[3], + keyframe_latent_width=encoded.shape[4], + pixel_frame_idx=pixel_frame_idx, + num_pixel_frames=num_pixel_frames, + fps=frame_rate, + device=device, + ) + appended_keyframes.append((condition_tokens, coords, strength)) + + # Pre-encoded keyframe guidance carried in from a previous temporal round. + for pixel_frame_index, keyframe_latent, strength in keyframe_latents or []: + keyframe_latent = keyframe_latent.to(device=device, dtype=dtype) + appended_keyframes.append( + ( + self._pack_latents(keyframe_latent, patch_size, patch_size_t), + self._prepare_keyframe_coords( + keyframe_latent_num_frames=keyframe_latent.shape[2], + keyframe_latent_height=keyframe_latent.shape[3], + keyframe_latent_width=keyframe_latent.shape[4], + pixel_frame_idx=pixel_frame_index, + num_pixel_frames=1, + fps=frame_rate, + device=device, + ), + strength, + ) + ) + + appended_coords = [] + # Guidance keyframes carry given content and are not marked. The learned embedding is for generated + # single-pixel-frame latents (causal first frame + slots). + keyframes_mask = torch.zeros_like(conditioning_mask) + # Causal encoding gives the first latent frame a temporal stride of 1, so it is marked like a slot. + tokens_per_latent_frame = latent_height * latent_width + keyframes_mask[:, :tokens_per_latent_frame] = 1.0 + for tokens, coords, strength in appended_keyframes: + latents = torch.cat([latents, torch.zeros_like(tokens)], dim=1) + clean_latents = torch.cat([clean_latents, tokens], dim=1) + conditioning_mask = torch.cat( + [conditioning_mask, conditioning_mask.new_full((batch_size, tokens.shape[1], 1), float(strength))], + dim=1, + ) + keyframes_mask = torch.cat([keyframes_mask, keyframes_mask.new_zeros((batch_size, tokens.shape[1], 1))], 1) + appended_coords.append(coords) + + # Generated keyframe slots: empty, fully-denoised single-pixel-frame token blocks the model fills in. + slot_token_slice = None + if slot_frame_indices: + if slot_initial_latents is None: + num_slot_tokens = tokens_per_latent_frame * len(slot_frame_indices) + slot_tokens = latents.new_zeros((batch_size, num_slot_tokens, latents.shape[2])) + else: + slot_initial_latents = slot_initial_latents.to(device=device, dtype=dtype) + slot_tokens = torch.cat( + [ + self._pack_latents(slot_initial_latents[:, :, index : index + 1], patch_size, patch_size_t) + for index in range(slot_initial_latents.shape[2]) + ], + dim=1, + ) + slot_token_slice = slice(latents.shape[1], latents.shape[1] + slot_tokens.shape[1]) + + latents = torch.cat([latents, slot_tokens], dim=1) + clean_latents = torch.cat([clean_latents, torch.zeros_like(slot_tokens)], dim=1) + conditioning_mask = torch.cat( + [conditioning_mask, conditioning_mask.new_zeros((batch_size, slot_tokens.shape[1], 1))], dim=1 + ) + keyframes_mask = torch.cat( + [keyframes_mask, keyframes_mask.new_ones((batch_size, slot_tokens.shape[1], 1))], dim=1 + ) + appended_coords.extend( + self._prepare_keyframe_coords( + keyframe_latent_num_frames=1, + keyframe_latent_height=latent_height, + keyframe_latent_width=latent_width, + pixel_frame_idx=position, + num_pixel_frames=1, + fps=frame_rate, + device=device, + ) + for position in slot_frame_indices + ) + + # IC-LoRA reference: the stage-1 half-resolution latent, held fully clean. + if reference_latents is not None: + reference_latents = reference_latents.to(device=device, dtype=dtype) + reference_tokens = self._pack_latents(reference_latents, patch_size, patch_size_t) + reference_coords = self.transformer.rope.prepare_video_coords( + batch_size=1, + num_frames=reference_latents.shape[2], + height=reference_latents.shape[3], + width=reference_latents.shape[4], + device=device, + fps=frame_rate, + ) + reference_coords[:, 1:, :, :] = reference_coords[:, 1:, :, :] * reference_downscale_factor + + latents = torch.cat([latents, torch.zeros_like(reference_tokens)], dim=1) + clean_latents = torch.cat([clean_latents, reference_tokens], dim=1) + conditioning_mask = torch.cat( + [conditioning_mask, conditioning_mask.new_ones((batch_size, reference_tokens.shape[1], 1))], dim=1 + ) + keyframes_mask = torch.cat( + [keyframes_mask, keyframes_mask.new_zeros((batch_size, reference_tokens.shape[1], 1))], dim=1 + ) + appended_coords.append(reference_coords) + + video_coords = self.transformer.rope.prepare_video_coords( + batch_size, latent_num_frames, latent_height, latent_width, device, fps=frame_rate + ) + if appended_coords: + appended = torch.cat(appended_coords, dim=2).expand(batch_size, -1, -1, -1) + video_coords = torch.cat([video_coords, appended], dim=2) + + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + latents = noise * noise_scale + latents * (1 - noise_scale) + latents = clean_latents * conditioning_mask + latents * (1 - conditioning_mask) + + return latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, slot_token_slice + + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2_condition.LTX2ConditionPipeline.prepare_audio_latents + def prepare_audio_latents( + self, + batch_size: int = 1, + num_channels_latents: int = 8, + audio_latent_length: int = 1, # 1 is just a dummy value + num_mel_bins: int = 64, + noise_scale: float = 0.0, + dtype: torch.dtype | None = None, + device: torch.device | None = None, + generator: torch.Generator | None = None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + if latents is not None: + # latents expected to be unpacked (4D) with shape [B, C, L, M] + latents = self._pack_audio_latents(latents) + latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) + latents = self._create_noised_state(latents, noise_scale, generator) + return latents.to(device=device, dtype=dtype) + + latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio + + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective batch" + f" size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + + # Sample in packed shape (B, L, C * M), following the original LTX-2.X code + packed_shape = (batch_size, audio_latent_length, num_channels_latents * latent_mel_bins) + latents = randn_tensor(packed_shape, generator=generator, device=device, dtype=dtype) + return latents + + @staticmethod + # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._create_noised_state + def _create_noised_state( + latents: torch.Tensor, noise_scale: float | torch.Tensor, generator: torch.Generator | None = None + ): + noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) + noised_latents = noise_scale * noise + (1 - noise_scale) * latents + return noised_latents + + def _unpack_video_latents(self, tokens: torch.Tensor, num_frames: int, height: int, width: int) -> torch.Tensor: + """Unpack a `(batch_size, tokens, channels)` block onto this transformer's patch grid.""" + return self._unpack_latents( + tokens, + num_frames, + height, + width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, + ) + + def upsample_latents(self, latents: torch.Tensor, upsampler: LTX2LatentUpsamplerModel) -> torch.Tensor: + """Run `upsampler` on normalized latents, round-tripping through raw VAE latent space as it expects.""" + latents = self._denormalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + latents = upsampler(latents.to(upsampler.dtype)) + return self._normalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + + def encode_conditions( + self, + conditions: list[LTX2VideoCondition] | None, + height: int, + width: int, + num_frames: int, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + generator: torch.Generator | None = None, + ) -> list[tuple[int, torch.Tensor, float, int]]: + """ + Preprocess and VAE-encode frame conditions, positioned by pixel frame. + + Returns `(pixel_frame_index, latent, strength, num_pixel_frames)` per condition, ready for + [`~LTX2DFRPipeline.prepare_latents`]'s `condition_latents`. Encoding is kept separate from placement because + the temporal refine rounds scale a condition's position by `2 ** round` and re-base it per tile, and should not + re-encode the same still once per tile to do so. + + The returned index is on `num_frames`' own pixel grid; carrying it onto a refined canvas is the caller's job. + """ + condition_frames, condition_strengths, condition_indices, condition_pixel_frames = self.preprocess_conditions( + conditions, height, width, num_frames, device=device + ) + encoded = [] + for pixels, strength, latent_index, num_pixel_frames in zip( + condition_frames, condition_strengths, condition_indices, condition_pixel_frames + ): + latent = self._normalize_latents( + retrieve_latents(self.vae.encode(pixels), generator=generator, sample_mode="argmax"), + self.vae.latents_mean, + self.vae.latents_std, + ).to(device=device, dtype=dtype) + pixel_index = 0 if latent_index == 0 else (latent_index - 1) * self.vae_temporal_compression_ratio + 1 + encoded.append((pixel_index, latent, strength, num_pixel_frames)) + return encoded + + def _rebuild_epilogue_keyframes( + self, + keyframe_latents: torch.Tensor, + decode_timestep: float, + decode_noise_scale: float, + seed: int, + device: torch.device, + dtype: torch.dtype, + ) -> torch.Tensor: + """ + Re-encode the carry keyframes at twice their resolution by way of RGB. + + These are frames the refine rounds already settled, so the epilogue is *given* them rather than asked to + generate them. Decoding to pixels, stretching x2 with Lanczos and encoding again preserves the frame while + landing it on the output grid, which is what lets the epilogue pin it fully clean. + + Each plane is decoded as its own one-frame clip. The VAE is causal, so a stacked decode would let neighbouring + planes bleed into each other -- they are independent stills, not a sequence. + + Args: + keyframe_latents (`torch.Tensor`): + `(batch_size, C, K, H, W)` normalized carry keyframe latents. + seed (`int`): + Base seed; plane `i` decodes under `seed + 4000 + i`, so a plane's pixels do not depend on how many + planes were decoded before it. + + Returns: + `torch.Tensor`: `(batch_size, C, K, 2H, 2W)` normalized latents. + """ + if keyframe_latents.ndim != 5: + raise ValueError(f"Expected carry keyframes (B, C, K, H, W), got {tuple(keyframe_latents.shape)}") + + encoded = [] + for index in range(keyframe_latents.shape[2]): + plane = self._denormalize_latents( + keyframe_latents[:, :, index : index + 1], + self.vae.latents_mean, + self.vae.latents_std, + self.vae.config.scaling_factor, + ) + timestep = None + if self.vae.config.timestep_conditioning: + noise = randn_tensor( + plane.shape, + generator=torch.Generator(device=device).manual_seed(seed + 4000 + index), + device=device, + dtype=plane.dtype, + ) + plane = (1 - decode_noise_scale) * plane + decode_noise_scale * noise + timestep = torch.full((plane.shape[0],), decode_timestep, device=device, dtype=plane.dtype) + frames = self.vae.decode(plane.to(self.vae.dtype), timestep, return_dict=False)[0] + + # PIL is the reference's Lanczos, and it wants one 8-bit channels-last image at a time. + stretched_batch = [] + for frame in frames: + rgb = ((frame[:, 0].float().clamp(-1, 1) + 1.0) * 127.5).round().to(torch.uint8) + image = PIL.Image.fromarray(rgb.permute(1, 2, 0).cpu().numpy(), mode="RGB") + image = image.resize((image.width * 2, image.height * 2), resample=PIL.Image.Resampling.LANCZOS) + stretched = torch.from_numpy(np.asarray(image, dtype=np.float32) / 127.5 - 1.0) + stretched_batch.append(stretched.permute(2, 0, 1).unsqueeze(1)) + stretched = torch.stack(stretched_batch).to(device=device, dtype=self.vae.dtype) + + encoded.append( + self._normalize_latents( + retrieve_latents(self.vae.encode(stretched), sample_mode="argmax"), + self.vae.latents_mean, + self.vae.latents_std, + self.vae.config.scaling_factor, + ) + ) + return torch.cat(encoded, dim=2).to(device=device, dtype=dtype) + + def _activate_adapters(self, adapter_names: list[str]) -> None: + """ + Make exactly `adapter_names` active on the transformer, switching LoRA off entirely for an empty list. + + Activation only. Adapter weights belong to the caller, who sets them through + [`~loaders.LTX2LoraLoaderMixin.set_adapters`], so switching between the base and detailing sets must not + rewrite them. + """ + if adapter_names: + self.transformer.enable_adapters() + self.transformer.set_adapter(adapter_names) + else: + self.transformer.disable_adapters() + + def denoise( + self, + latents: torch.Tensor, + conditioning_mask: torch.Tensor, + clean_latents: torch.Tensor, + video_coords: torch.Tensor, + keyframes_mask: torch.Tensor, + prompt_embeds: torch.Tensor, + audio_prompt_embeds: torch.Tensor, + prompt_attention_mask: torch.Tensor, + sigmas: list[float], + frame_rate: float, + audio_latents: torch.Tensor, + freeze_audio: bool = False, + ancestral_eta: float = 0.0, + video_tile_plan: list[dict[str, torch.Tensor]] | None = None, + generator: torch.Generator | None = None, + ancestral_generator: torch.Generator | None = None, + use_cross_timestep: bool = True, + attention_kwargs: dict[str, Any] | None = None, + progress_bar=None, + step_offset: int = 0, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Run one DFR denoising pass over `sigmas` and return `(latents, audio_latents)`, both still packed. + + The distilled schedule is used without classifier-free guidance, so this is a single transformer call per step. + Every pass runs both streams, because the video branch needs the cross-modal attention even where the audio it + produces is thrown away. `freeze_audio=True` keeps the audio stream at sigma 0 (no Euler step) so video can + still cross-attend to it — the temporal refine tiles and the epilogue use this to follow stage-1 speech without + each tile re-denoising a different audio realization. + + `ancestral_eta > 0` switches from the scheduler's Euler step to [`ancestral_euler_step`]. + + Args: + sigmas (`list[float]`): + Noise schedule for this pass, without the terminal `0.0` (the scheduler appends it). + ancestral_generator (`torch.Generator`, *optional*): + Generator for the ancestral renoise draws, kept separate from `generator` so those draws do not consume + the state the next tile's initial noising reads. Required when `ancestral_eta > 0`. + freeze_audio (`bool`, *optional*, defaults to `False`): + Hold `audio_latents` clean (timestep/sigma 0, no audio Euler step) while still running audio-to-video + cross-attention. Ignored when `audio_latents` is `None`. + video_tile_plan (`list[dict[str, torch.Tensor]]`, *optional*): + Per-tile token plan from [`~LTX2DFRPipeline._video_tile_plan`]. When given, each step runs the + transformer once per tile and blends the predictions, so the sampler still steps a single full canvas + and the tiles agree on their overlaps at every step. This is what puts a resolution too large for one + forward pass within reach. + step_offset (`int`): + Index of this pass's first step within the pipeline's whole schedule, used for `callback_on_step_end` + and the shared progress bar. + """ + device = latents.device + # The audio stream's own length; passing it separately only invites the two disagreeing. + audio_num_frames = audio_latents.shape[1] + + self.scheduler.set_timesteps(sigmas=sigmas, device=device) + timesteps = self.scheduler.timesteps + # The video and audio streams step the same schedule, but `step` tracks its own index, so audio needs its own + # scheduler instance. + audio_scheduler = copy.deepcopy(self.scheduler) + + audio_coords = self.transformer.audio_rope.prepare_audio_coords( + audio_latents.shape[0], audio_num_frames, device + ) + + for index, t in enumerate(timesteps): + if self.interrupt: + continue + + self._current_timestep = t + timestep_scalar = t.expand(latents.shape[0]) + # Conditioned tokens see a proportionally lower noise level, which is what holds them near their clean + # content while the rest of the sequence denoises. + video_timestep = timestep_scalar.unsqueeze(-1) * (1 - conditioning_mask.squeeze(-1)) + if freeze_audio: + audio_timestep = torch.zeros(audio_latents.shape[0], device=device, dtype=t.dtype) + else: + audio_timestep = audio_scheduler.timesteps[index].expand(audio_latents.shape[0]) + + # Everything a tile shares with the full canvas: audio, text, and the schedule's scalar sigma. + transformer_kwargs = { + "encoder_hidden_states": prompt_embeds, + "audio_encoder_hidden_states": audio_prompt_embeds, + "audio_timestep": audio_timestep, + "sigma": timestep_scalar, + "audio_sigma": audio_timestep, + "encoder_attention_mask": prompt_attention_mask, + "audio_encoder_attention_mask": prompt_attention_mask, + "fps": frame_rate, + "audio_num_frames": audio_num_frames, + "audio_coords": audio_coords, + "use_cross_timestep": use_cross_timestep, + "attention_kwargs": attention_kwargs, + "return_dict": False, + } + + if video_tile_plan is None: + noise_pred_video, noise_pred_audio = self.transformer( + hidden_states=latents.to(prompt_embeds.dtype), + audio_hidden_states=audio_latents.to(prompt_embeds.dtype), + timestep=video_timestep, + video_keyframes_mask=keyframes_mask, + video_coords=video_coords, + **transformer_kwargs, + ) + else: + # Blending the velocity is the same as blending x0: `x0 = x - v * sigma` is affine in `v`, every tile + # reads the same `latents` and the same per-token sigma, and the weights sum to one. + noise_pred_video = torch.zeros_like(latents, dtype=torch.float32) + noise_pred_audio = None + for tile in video_tile_plan: + keep = tile.keep + tile_pred, tile_audio_pred = self.transformer( + hidden_states=latents[:, keep].to(prompt_embeds.dtype), + audio_hidden_states=audio_latents.to(prompt_embeds.dtype), + timestep=video_timestep[:, keep], + video_keyframes_mask=keyframes_mask[:, keep], + video_coords=tile.coords, + **transformer_kwargs, + ) + noise_pred_video.index_add_( + 1, keep, tile_pred.float() * tile.weights.to(torch.float32).view(1, -1, 1) + ) + if tile_audio_pred is not None: + # Every tile saw the whole audio under a different video context; average them. + contribution = tile_audio_pred.float() / len(video_tile_plan) + noise_pred_audio = ( + contribution if noise_pred_audio is None else noise_pred_audio + contribution + ) + noise_pred_video = noise_pred_video.to(latents.dtype) + + # Conditioning is applied in x0 space. Convert velocity -> x0 with each token's own noise level: a token + # held at strength `s` sits at `(1 - s) * sigma`, so the scalar schedule sigma would mis-scale it. + # Deliberately not the scheduler's `per_token_timesteps` path: that steps each token from its own sigma to + # the nearest schedule sigma below it, whereas a conditioned token has to stay pinned while the rest of the + # sequence advances on the shared schedule. + sigma = self.scheduler.sigmas[index] + per_token_sigma = (video_timestep / self.scheduler.config.num_train_timesteps).unsqueeze(-1) + denoised = latents.float() - noise_pred_video.float() * per_token_sigma + denoised = denoised * (1 - conditioning_mask) + clean_latents.float() * conditioning_mask + + if ancestral_eta > 0: + noise = randn_tensor(latents.shape, generator=ancestral_generator, device=device, dtype=torch.float32) + stepped = ancestral_euler_step( + latents, denoised, sigma, self.scheduler.sigmas[index + 1], ancestral_eta, noise + ) + # Ancestral Euler noises every token, so re-apply the blend or strength-0.95 seam anchors erode. + latents = (stepped * (1 - conditioning_mask) + clean_latents.float() * conditioning_mask).to( + latents.dtype + ) + else: + latents = self.scheduler.step((latents.float() - denoised) / sigma, t, latents, return_dict=False)[0] + + if not freeze_audio: + audio_sigma = audio_scheduler.sigmas[index] + audio_denoised = audio_latents.float() - noise_pred_audio.float() * audio_sigma + audio_latents = audio_scheduler.step( + (audio_latents.float() - audio_denoised) / audio_sigma, t, audio_latents, return_dict=False + )[0] + + if callback_on_step_end is not None: + callback_kwargs = {} + for k in callback_on_step_end_tensor_inputs or []: + callback_kwargs[k] = locals()[k] + callback_outputs = callback_on_step_end(self, step_offset + index, t, callback_kwargs) + + latents = callback_outputs.pop("latents", latents) + prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) + + if progress_bar is not None: + progress_bar.update() + + if XLA_AVAILABLE: + xm.mark_step() + + return latents, audio_latents + + @property + def num_timesteps(self): + return self._num_timesteps + + @property + def current_timestep(self): + return self._current_timestep + + @property + def attention_kwargs(self): + return self._attention_kwargs + + @property + def interrupt(self): + return self._interrupt + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] = None, + conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, + height: int = 704, + width: int = 1216, + num_frames: int | None = None, + frame_rate: float = 24.0, + min_seconds: float = 1.0, + max_seconds: float = 20.0, + temporal_upscalings: int = 0, + spatial_upscalings: int = 1, + detailing_lora_adapter_name: str | None = None, + detailing_reference_downscale_factor: int = 2, + stage_1_sigmas: list[float] = DISTILLED_SIGMA_VALUES, + stage_2_sigmas: list[float] = STAGE_2_DISTILLED_SIGMA_VALUES, + temporal_round_sigmas: list[float] = TEMPORAL_ROUND_DISTILLED_SIGMA_VALUES, + num_videos_per_prompt: int | None = 1, + generator: torch.Generator | list[torch.Generator] | None = None, + prompt_embeds: torch.Tensor | None = None, + prompt_attention_mask: torch.Tensor | None = None, + decode_timestep: float | list[float] = 0.0, + decode_noise_scale: float | list[float] | None = None, + use_cross_timestep: bool = True, + system_prompt: str | None = None, + enable_prompt_enhancement: bool = False, + prompt_max_new_tokens: int | None = None, + prompt_enhancement_kwargs: dict[str, Any] | None = None, + prompt_enhancement_seed: int = 10, + output_type: str = "pil", + return_dict: bool = True, + attention_kwargs: dict[str, Any] | None = None, + callback_on_step_end: Callable[[int, int], None] | None = None, + callback_on_step_end_tensor_inputs: list[str] = ["latents"], + max_sequence_length: int = 1024, + ): + r""" + Function invoked when calling the pipeline for generation. + + Args: + prompt (`str` or `List[str]`, *optional*): + The prompt or prompts to guide the video generation. If not defined, one has to pass `prompt_embeds`. + conditions (`LTX2VideoCondition` or `List[LTX2VideoCondition]`, *optional*): + Frame-level image or video conditions. `index` is a *latent* index, so a condition cannot sit on an + interior keyframe seam (pixel 24, 32, ...); only on the latent boundary at or below it. It is read on + the canvas `num_frames` asks for, and the moment it names is carried onto each refine round's longer + canvas, so a condition stays at the same point in the clip however many rounds run. + height (`int`, *optional*, defaults to `704`): + The height in pixels of the *returned* video. Every stage below it runs at a halved resolution, so this + must be divisible by `2 ** spatial_upscalings` times the VAE's spatial compression ratio. + width (`int`, *optional*, defaults to `1216`): + The width in pixels of the returned video, subject to the same divisibility rule as `height`. + num_frames (`int`, *optional*): + The number of video frames to generate, before temporal refinement. If not supplied, the duration is + predicted from the prompt by the `duration_head`. Must satisfy `(num_frames - 1) % 8 == 0`; DFR pads + the canvas internally to a whole number of keyframe segments and trims back to `(num_frames - 1) * 2 ** + temporal_upscalings + 1` before decoding. + frame_rate (`float`, *optional*, defaults to `24.0`): + The frames per second (FPS) the base canvas is generated at. Each temporal refine round doubles the + frame rate of the returned video. + min_seconds (`float`, *optional*, defaults to `1.0`): + Lower bound on the auto-predicted duration when `num_frames` is omitted. + max_seconds (`float`, *optional*, defaults to `20.0`): + Upper bound on the auto-predicted duration when `num_frames` is omitted. + temporal_upscalings (`int`, *optional*, defaults to `0`): + Number of temporal x2 refine rounds (`0` -> base frame rate, `1` -> 2x with 2 tiles, `2` -> 4x with 4 + tiles). Anything above `0` requires the `temporal_latent_upsampler` component. + spatial_upscalings (`int`, *optional*, defaults to `1`): + Number of spatial x2 stages between the base canvas and `height` x `width`. `1` runs the two-stage + recipe: a base pass at half resolution, then the detailing pass at full resolution. `2` starts a + quarter resolution and adds a third, tiled full-resolution detailing pass *after* the temporal rounds, + which needs `height` and `width` divisible by four times the VAE's spatial compression ratio. That + epilogue is a detailing pass, so pass `detailing_lora_adapter_name` too unless you have a reason not + to. + detailing_lora_adapter_name (`str`, *optional*): + Adapter name of the 2x spatial detailing IC-LoRA + ([`Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler`](https://huggingface.co/Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler)), + loaded beforehand with [`~loaders.LTX2LoraLoaderMixin.load_lora_weights`]. When set, stage 2 runs with + this adapter active *and* attends to the stage-1 half-resolution latent as an in-context reference; + stage 1 and the temporal rounds run with it deactivated, and whatever adapters were active on entry are + restored before returning. The shipped adapter is calibrated for strength `0.5`, so set that weight + when you load it. Leave unset to run without the detailing pass. + detailing_reference_downscale_factor (`int`, *optional*, defaults to `2`): + Ratio between the target and the stage-1 reference resolution, used to scale the reference tokens' + spatial coordinates into the target's coordinate space. Defaults to `2`, matching the shipped detailing + LoRA. Ignored unless `detailing_lora_adapter_name` is set. + stage_1_sigmas (`list[float]`, *optional*): + Noise schedule for the half-resolution keyframe-slot stage. + stage_2_sigmas (`list[float]`, *optional*): + Noise schedule for the full-resolution detailing stage. Its first value is also the level stage 1's + upsampled latents are re-noised to. + temporal_round_sigmas (`list[float]`, *optional*): + Noise schedule for each temporal refine round's tiles. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of videos to generate per prompt. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Random generator(s) for reproducibility. + prompt_embeds (`torch.Tensor`, *optional*): + Pre-generated text embeddings. + prompt_attention_mask (`torch.Tensor`, *optional*): + Pre-generated attention mask for text embeddings. + decode_timestep (`float`, defaults to `0.0`): + The timestep at which generated video is decoded. + decode_noise_scale (`float`, defaults to `None`): + Noise scale at decode time. + use_cross_timestep (`bool`, *optional*, defaults to `True`): + Whether to use cross-modality sigma for cross attention modulation. `True` for LTX-2.3+. + system_prompt (`str`, *optional*): + Optional system prompt for prompt enhancement. See `enable_prompt_enhancement`. + enable_prompt_enhancement (`bool`, *optional*, defaults to `False`): + Whether to run prompt enhancement. + prompt_max_new_tokens (`int`, *optional*): + The maximum number of new tokens to generate when performing prompt enhancement. + prompt_enhancement_kwargs (`dict[str, Any]`, *optional*): + Keyword arguments for the prompt enhancer's `.generate` call. + prompt_enhancement_seed (`int`, *optional*, defaults to `10`): + Random seed for any random operations during prompt enhancement. + output_type (`str`, *optional*, defaults to `"pil"`): + Output format. Choose `"pil"`, `"np"`, `"pt"` or `"latent"`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`LTX2PipelineOutput`] or a plain tuple. + attention_kwargs (`dict`, *optional*): + Additional kwargs passed to the attention processor. + callback_on_step_end (`Callable`, *optional*): + A function called at the end of each denoising step, across every stage and temporal tile. + callback_on_step_end_tensor_inputs (`List`, *optional*, defaults to `["latents"]`): + Tensor inputs for the callback function. + max_sequence_length (`int`, *optional*, defaults to `1024`): + Maximum sequence length for the text prompt. + + Examples: + + Returns: + [`LTX2PipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`LTX2PipelineOutput`] is returned, otherwise a `tuple` of `(video, audio)` + is returned. + """ + if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): + callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs + + # 1. Check inputs. The resolution rule comes first: it is stricter than the generic one every LTX-2 pipeline + # applies, so checking it up front reports the divisor a DFR caller actually has to satisfy. A 4K run at + # `spatial_upscalings=2` needs multiples of 128, which is why UHD is 3840x2176 here and not 3840x2160. + if spatial_upscalings not in (1, 2): + raise ValueError(f"`spatial_upscalings` must be 1 or 2, got {spatial_upscalings}") + # Every stage below the output runs at a halved resolution, so the smallest one must still land on the VAE's + # spatial grid. + spatial_divisor = 2**spatial_upscalings + stage_1_divisor = spatial_divisor * self.vae_spatial_compression_ratio + if height % stage_1_divisor != 0 or width % stage_1_divisor != 0: + raise ValueError( + f"`height` and `width` have to be divisible by {stage_1_divisor} for a DFR recipe with " + f"`spatial_upscalings={spatial_upscalings}` on a VAE that compresses space by " + f"{self.vae_spatial_compression_ratio}, but are {height} and {width}. The nearest valid size at or " + f"below is {height // stage_1_divisor * stage_1_divisor}x{width // stage_1_divisor * stage_1_divisor}." + ) + self.check_inputs( + prompt=prompt, + height=height, + width=width, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + system_prompt=system_prompt, + enable_prompt_enhancement=enable_prompt_enhancement, + num_frames=num_frames, + min_seconds=min_seconds, + max_seconds=max_seconds, + # The distilled schedule runs without classifier-free or spatio-temporal guidance. + stg_scale=0.0, + audio_stg_scale=0.0, + ) + if temporal_upscalings not in (0, 1, 2): + raise ValueError(f"`temporal_upscalings` must be 0, 1 or 2, got {temporal_upscalings}") + if getattr(self, "latent_upsampler", None) is None: + raise ValueError( + "DFR requires the pipeline's `latent_upsampler` component. Load it from the `latent_upsampler` " + "subfolder and pass it to `from_pretrained`." + ) + if temporal_upscalings > 0 and getattr(self, "temporal_latent_upsampler", None) is None: + raise ValueError( + "`temporal_upscalings > 0` requires the pipeline's optional `temporal_latent_upsampler` component." + ) + if not self.transformer.config.use_keyframes_abs_pos_embedding: + raise ValueError( + "DFR generates keyframe slots, which requires a transformer whose config sets " + "`use_keyframes_abs_pos_embedding` (LTX-2.5 and later). Each slot costs a full latent frame of tokens, " + "so a checkpoint without the learned marker would spend that budget on tokens it cannot interpret." + ) + + # The detailing IC-LoRA belongs to stage 2 alone: stage 1 and the temporal refine rounds run on the base + # adapter set. Snapshot what was active on entry so the transformer is handed back as it was found. + detailing = detailing_lora_adapter_name is not None + original_adapters: list[str] = [] + base_adapters: list[str] = [] + detailing_adapters: list[str] = [] + if detailing: + loaded_adapters = list(getattr(self.transformer, "peft_config", None) or []) + if detailing_lora_adapter_name not in loaded_adapters: + raise ValueError( + f"`detailing_lora_adapter_name={detailing_lora_adapter_name!r}` is not loaded on the transformer. " + f"Load the detailing IC-LoRA with `load_lora_weights(..., adapter_name=...)` first; currently " + f"loaded adapters are {loaded_adapters}." + ) + original_adapters = list(self.transformer.active_adapters()) + base_adapters = [name for name in original_adapters if name != detailing_lora_adapter_name] + detailing_adapters = [*base_adapters, detailing_lora_adapter_name] + + try: + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + # 2. Define call parameters + if prompt is not None and isinstance(prompt, str): + batch_size = 1 + elif prompt is not None and isinstance(prompt, list): + batch_size = len(prompt) + else: + batch_size = prompt_embeds.shape[0] + batch_size *= num_videos_per_prompt + + if conditions is not None and not isinstance(conditions, list): + conditions = [conditions] + + # Temporal tiles seed ancestral noise as `seed + 1000 * round + tile`. A list of generators uses the first. + seed_source = generator[0] if isinstance(generator, list) else generator + ancestral_seed_base = seed_source.initial_seed() if seed_source is not None else 0 + + device = self._execution_device + + # 3. Prepare text embeddings + if enable_prompt_enhancement and prompt is not None: + enhancement_image = None + for condition in conditions or []: + frames = condition.frames + if isinstance(frames, PIL.Image.Image): + enhancement_image = frames + break + if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], PIL.Image.Image): + enhancement_image = frames[0] + break + if system_prompt is None: + system_prompt = ( + LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT + if enhancement_image is not None + else LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT + ) + prompt = self.enhance_prompt( + prompt=prompt, + system_prompt=system_prompt, + max_new_tokens=prompt_max_new_tokens, + seed=prompt_enhancement_seed, + generator=generator, + generation_kwargs=prompt_enhancement_kwargs, + device=device, + image=enhancement_image, + ) + + prompt_embeds, prompt_attention_mask = self.encode_prompt( + prompt=prompt, + num_videos_per_prompt=num_videos_per_prompt, + prompt_embeds=prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + max_sequence_length=max_sequence_length, + device=device, + ) + video_prompt_embeds, audio_prompt_embeds, connector_attention_mask = self.connectors( + prompt_embeds, prompt_attention_mask, padding_side=self.tokenizer_padding_side + ) + + if num_frames is None: + if getattr(self, "duration_head", None) is None: + raise ValueError( + "`num_frames` must be supplied when the pipeline has no `duration_head` component to predict it." + ) + # `num_videos_per_prompt` duplicates sit after the first row. + num_frames = self.duration_head.predict_num_frames( + video_prompt_embeds[:1], + audio_prompt_embeds[:1], + frame_rate=frame_rate, + temporal_compression_ratio=self.vae_temporal_compression_ratio, + min_seconds=min_seconds, + max_seconds=max_seconds, + ) + + # 4. Resolve the keyframe segment grid. The canvas pads the request up to a whole number of segments; the + # excess tail is trimmed before decoding. + requested_frames = num_frames + canvas_frames, _, slot_frame_indices = resolve_canvas(num_frames, self.vae_temporal_compression_ratio) + + num_channels_latents = self.transformer.config.in_channels + latent_mel_bins = self.audio_mel_bins // self.audio_vae_mel_compression_ratio + audio_latents_per_second = ( + self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) + ) + audio_num_frames = round(canvas_frames / frame_rate * audio_latents_per_second) + # The shipped audio covers the stage-1 canvas; the refine rounds slice it by wall clock, and both the + # canvas and the frame rate change under them, so its duration is snapshotted here. + stage_1_seconds = canvas_frames / frame_rate + # Stages 1 and 2 lay out RoPE time at the snapped rate too: `frame_rate` above 30 is a rate the + # transformer never saw. Audio stays sized off the playback rate. + base_conditioning_fps = _conditioning_fps(frame_rate) + + # Plan the temporal rounds up front. Each round's tiling is fully determined by the previous round's keyframe + # positions, so resolving it here gives an exact step total for the progress bar and `callback_on_step_end`. + round_plan: list[tuple[int, list[int], list]] = [] + plan_canvas, plan_positions = canvas_frames, list(slot_frame_indices) + for round_index in range(1, temporal_upscalings + 1): + plan_canvas = 2 * (plan_canvas - 1) + 1 + seam_positions = [2 * position for position in plan_positions] + tiles = temporal_tile_plan( + seam_positions, plan_canvas, 2**round_index, self.vae_temporal_compression_ratio + ) + round_plan.append((plan_canvas, seam_positions, tiles)) + slots = {position for tile in tiles for position in tile.slots} + plan_positions = sorted(set(seam_positions) | slots) + + # The epilogue tiles the finished canvas at the output resolution, so its layout is known here too. Its + # temporal cuts land on the last refine round's window seams, which are the keyframes the canvas it + # inherits was actually stitched on. + epilogue_layout: list[tuple[slice, slice, slice, torch.Tensor]] = [] + if spatial_upscalings == 2: + epilogue_layout = epilogue_tiles( + latent_shape=( + (plan_canvas - 1) // self.vae_temporal_compression_ratio + 1, + height // self.vae_spatial_compression_ratio, + width // self.vae_spatial_compression_ratio, + ), + frame_tiles=2**temporal_upscalings, + frame_seams=[ + pixel_to_latent_index(position, self.vae_temporal_compression_ratio) + for position in (round_plan[-1][1] if round_plan else []) + ], + ) + + self._num_timesteps = ( + len(stage_1_sigmas) + + len(stage_2_sigmas) + + len(temporal_round_sigmas) * sum(len(tiles) for _, _, tiles in round_plan) + # The epilogue is a single loop over the whole canvas; its tiling lives inside the transformer call. + + (len(stage_2_sigmas) if epilogue_layout else 0) + ) + progress_bar = self.progress_bar(total=self._num_timesteps) + step_offset = 0 + + # 5. Stage 1: half-resolution video plus generated keyframe slots. + if detailing: + self._activate_adapters(base_adapters) + stage_1_height = height // spatial_divisor + stage_1_width = width // spatial_divisor + latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, slot_token_slice = ( + self.prepare_latents( + conditions=conditions, + slot_frame_indices=slot_frame_indices, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=stage_1_height, + width=stage_1_width, + num_frames=canvas_frames, + frame_rate=base_conditioning_fps, + noise_scale=stage_1_sigmas[0], + dtype=torch.float32, + device=device, + generator=generator, + ) + ) + audio_latents = self.prepare_audio_latents( + batch_size=batch_size, + num_channels_latents=self.audio_latent_channels, + audio_latent_length=audio_num_frames, + num_mel_bins=self.audio_mel_bins, + dtype=torch.float32, + device=device, + generator=generator, + ) + + latent_num_frames = (canvas_frames - 1) // self.vae_temporal_compression_ratio + 1 + stage_1_latent_height = stage_1_height // self.vae_spatial_compression_ratio + stage_1_latent_width = stage_1_width // self.vae_spatial_compression_ratio + latents, audio_latents = self.denoise( + latents=latents, + conditioning_mask=conditioning_mask, + clean_latents=clean_latents, + video_coords=video_coords, + keyframes_mask=keyframes_mask, + prompt_embeds=video_prompt_embeds, + audio_prompt_embeds=audio_prompt_embeds, + prompt_attention_mask=connector_attention_mask, + sigmas=stage_1_sigmas, + frame_rate=base_conditioning_fps, + audio_latents=audio_latents, + generator=generator, + use_cross_timestep=use_cross_timestep, + attention_kwargs=attention_kwargs, + progress_bar=progress_bar, + step_offset=step_offset, + callback_on_step_end=callback_on_step_end, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + step_offset += len(stage_1_sigmas) + + # The shipped audio is stage 1's: later stages still run an audio pass so video can cross-attend, but they + # re-noise it and the result is discarded. + stage_1_audio_latents = audio_latents + half_resolution_latents = self._unpack_video_latents( + latents[:, : latent_num_frames * stage_1_latent_height * stage_1_latent_width], + latent_num_frames, + stage_1_latent_height, + stage_1_latent_width, + ) + slot_keyframes = self._unpack_video_latents( + latents[:, slot_token_slice], + len(slot_frame_indices), + stage_1_latent_height, + stage_1_latent_width, + ) + + # 6. Upsample the stage-1 video and its keyframe slots into the full-resolution latent grid. + upsampled_latents = self.upsample_latents(half_resolution_latents, self.latent_upsampler) + upsampled_slot_keyframes = self.upsample_latents(slot_keyframes, self.latent_upsampler) + + # 7. Stage 2: re-denoise at full resolution, seeded from the upsampled latents and slots, with the detailing + # IC-LoRA and its stage-1 in-context reference. + if detailing: + self._activate_adapters(detailing_adapters) + stage_2_height = height // (spatial_divisor // 2) + stage_2_width = width // (spatial_divisor // 2) + latent_height = stage_2_height // self.vae_spatial_compression_ratio + latent_width = stage_2_width // self.vae_spatial_compression_ratio + latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, slot_token_slice = ( + self.prepare_latents( + conditions=conditions, + slot_frame_indices=slot_frame_indices, + slot_initial_latents=upsampled_slot_keyframes, + reference_latents=half_resolution_latents if detailing else None, + reference_downscale_factor=detailing_reference_downscale_factor, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=stage_2_height, + width=stage_2_width, + num_frames=canvas_frames, + frame_rate=base_conditioning_fps, + noise_scale=stage_2_sigmas[0], + dtype=torch.float32, + device=device, + generator=generator, + latents=upsampled_latents, + ) + ) + latents, _ = self.denoise( + latents=latents, + conditioning_mask=conditioning_mask, + clean_latents=clean_latents, + video_coords=video_coords, + keyframes_mask=keyframes_mask, + prompt_embeds=video_prompt_embeds, + audio_prompt_embeds=audio_prompt_embeds, + prompt_attention_mask=connector_attention_mask, + sigmas=stage_2_sigmas, + frame_rate=base_conditioning_fps, + audio_latents=self._create_noised_state(stage_1_audio_latents, stage_2_sigmas[0], generator), + generator=generator, + use_cross_timestep=use_cross_timestep, + attention_kwargs=attention_kwargs, + progress_bar=progress_bar, + step_offset=step_offset, + callback_on_step_end=callback_on_step_end, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + step_offset += len(stage_2_sigmas) + if detailing: + self._activate_adapters(base_adapters) + + tokens_per_latent_frame = latent_height * latent_width + video_latents = self._unpack_video_latents( + latents[:, : latent_num_frames * tokens_per_latent_frame], + latent_num_frames, + latent_height, + latent_width, + ) + carry_keyframes = self._unpack_video_latents( + latents[:, slot_token_slice], + len(slot_frame_indices), + latent_height, + latent_width, + ) + + carry_positions = list(slot_frame_indices) + + # 8. Temporal refine rounds. + temporal_ratio = self.vae_temporal_compression_ratio + for round_index, (canvas_frames, seam_positions, tiles) in enumerate(round_plan, start=1): + video_latents = self.upsample_latents(video_latents, self.temporal_latent_upsampler) + frame_rate = 2 * frame_rate + anchor_keyframes = carry_keyframes + conditioning_fps = _conditioning_fps(frame_rate) + # `condition.index` is an original-canvas latent index; this round's pixel grid is scaled by `2**round`. + pixel_scale = 2**round_index + round_conditions = self.encode_conditions( + conditions, + stage_2_height, + stage_2_width, + requested_frames, + device=device, + dtype=torch.float32, + ) + + tile_latents = [] + slot_positions: list[int] = [] + slot_latents: list[torch.Tensor] = [] + seam_to_index = {seam: index for index, seam in enumerate(seam_positions)} + for tile_index, tile in enumerate(tiles): + tile_frames = (tile.interval.end - tile.interval.start - 1) * temporal_ratio + 1 + tile_video_latents = video_latents[:, :, tile.interval.start : tile.interval.end] + + # A condition's moment moves with the canvas: after this round it sits at `pixel_scale` times its + # original pixel position. Only the windows that actually cover it re-attach it, re-based on their + # own first frame. + tile_conditions = [ + (pixel * pixel_scale - tile.pixel_start, latent, strength, num_pixel_frames) + for pixel, latent, strength, num_pixel_frames in round_conditions + if tile.pixel_start <= pixel * pixel_scale <= tile.pixel_end + ] + + # Every seam inside the window is a hard keyframe, including the one at local frame 0. They are + # the keyframes carried in from the previous round, pinned just short of fully clean. + tile_keyframe_latents = [ + ( + position - tile.pixel_start, + anchor_keyframes[:, :, seam_to_index[position] : seam_to_index[position] + 1], + ANCHOR_KEYFRAME_STRENGTH, + ) + for position in tile.anchors + ] + + # Seed each invented slot from the nearest latent frame of the temporally upsampled tile. + tile_slot_positions = [position - tile.pixel_start for position in tile.slots] + seed_indices = [ + min(round(position / temporal_ratio), tile_video_latents.shape[2] - 1) + for position in tile_slot_positions + ] + tile_slot_initials = torch.cat( + [tile_video_latents[:, :, index : index + 1] for index in seed_indices], dim=2 + ) + + ( + tile_packed_latents, + tile_conditioning_mask, + tile_clean_latents, + tile_video_coords, + tile_keyframes_mask, + tile_slot_slice, + ) = self.prepare_latents( + condition_latents=tile_conditions, + keyframe_latents=tile_keyframe_latents, + slot_frame_indices=tile_slot_positions, + slot_initial_latents=tile_slot_initials, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=stage_2_height, + width=stage_2_width, + num_frames=tile_frames, + frame_rate=conditioning_fps, + noise_scale=temporal_round_sigmas[0], + dtype=torch.float32, + device=device, + generator=generator, + latents=tile_video_latents, + ) + # Per-tile ancestral seed, kept off `generator` so these draws do not consume the next tile's noising. + ancestral_generator = torch.Generator(device=device).manual_seed( + ancestral_seed_base + 1000 * round_index + tile_index + ) + # Frozen stage-1 audio, cut to this window's slice of the *playback* clock and resampled to the + # token count this tile's frame count and conditioning fps ask for. Re-noising audio per tile (or + # isolating it) made the face jump at the seam: each side followed a different realization. + tile_audio_latents = _audio_window_for_tile( + stage_1_audio_latents, + pixel_start=tile.pixel_start, + tile_frames=tile_frames, + playback_fps=frame_rate, + source_seconds=stage_1_seconds, + conditioning_fps=conditioning_fps, + audio_latents_per_second=audio_latents_per_second, + ) + tile_packed_latents, _ = self.denoise( + latents=tile_packed_latents, + conditioning_mask=tile_conditioning_mask, + clean_latents=tile_clean_latents, + video_coords=tile_video_coords, + keyframes_mask=tile_keyframes_mask, + prompt_embeds=video_prompt_embeds, + audio_prompt_embeds=audio_prompt_embeds, + prompt_attention_mask=connector_attention_mask, + sigmas=temporal_round_sigmas, + frame_rate=conditioning_fps, + audio_latents=tile_audio_latents, + freeze_audio=True, + ancestral_eta=TEMPORAL_ANCESTRAL_ETA, + generator=generator, + ancestral_generator=ancestral_generator, + use_cross_timestep=use_cross_timestep, + attention_kwargs=attention_kwargs, + progress_bar=progress_bar, + step_offset=step_offset, + callback_on_step_end=callback_on_step_end, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + step_offset += len(temporal_round_sigmas) + + unpacked = self._unpack_video_latents( + tile_packed_latents[:, : tile_video_latents.shape[2] * tokens_per_latent_frame], + tile_video_latents.shape[2], + latent_height, + latent_width, + ) + tile_latents.append(unpacked) + if tile_slot_slice is not None: + slot_positions.extend(tile.slots) + slot_latents.append( + self._unpack_video_latents( + tile_packed_latents[:, tile_slot_slice], + len(tile_slot_positions), + latent_height, + latent_width, + ) + ) + + # Each tile contributes strictly after the seam the previous one kept, so the ramp is dropped + # rather than blended: both sides reproduce a known keyframe there and averaging only smears it. + video_latents = torch.cat( + [latent[:, :, tile.interval.left_ramp :] for latent, tile in zip(tile_latents, tiles)], dim=2 + ) + expected_latent_frames = (canvas_frames - 1) // temporal_ratio + 1 + if video_latents.shape[2] != expected_latent_frames: + raise RuntimeError( + f"Stitched round {round_index} has T={video_latents.shape[2]} latent frames, expected " + f"{expected_latent_frames}" + ) + + # Next round's anchor bag: the keyframes carried into this round plus the slots it invented. Two tiles + # invent the slot that falls in the later one's dropped lead-in, and the stitch keeps the earlier + # tile's frames there, so the earlier tile's copy is the one the canvas actually holds. A slot + # overwrites an anchor at the same position. + carry: dict[int, torch.Tensor] = { + position: anchor_keyframes[:, :, index : index + 1] + for index, position in enumerate(seam_positions) + } + all_slot_latents = torch.cat(slot_latents, dim=2) if slot_latents else None + first_slot_index: dict[int, int] = {} + for index, position in enumerate(slot_positions): + first_slot_index.setdefault(position, index) + for position, index in first_slot_index.items(): + carry[position] = all_slot_latents[:, :, index : index + 1] + carry_positions = sorted(carry) + carry_keyframes = torch.cat([carry[position] for position in carry_positions], dim=2) + + # 9. Spatial detailing epilogue. Stage 2 ran at half the output resolution, so the finished canvas is + # upsampled once more and re-detailed at full resolution. One denoising loop covers the whole canvas and + # the transformer call inside it is tiled, because a full-resolution forward pass over the refined canvas + # does not fit. Every Euler step therefore steps a canvas whose tiles have already agreed on their + # overlaps. Spatial tiles blend, since neither side of a height or width border holds a known frame; + # temporal tiles are cut on the last refine round's seams instead. + if spatial_upscalings == 2: + guide_latents = video_latents + video_latents = self.upsample_latents(guide_latents, self.latent_upsampler) + conditioning_fps = _conditioning_fps(frame_rate) + # Only the video latent is spatially upsampled. The keyframes go the long way round -- decode, Lanczos + # x2, encode -- which hands the epilogue finished frames to pin rather than positions to fill in. + # The planes decode one at a time, so the per-batch-element decode settings collapse to a scalar. + plane_timestep = decode_timestep[0] if isinstance(decode_timestep, list) else decode_timestep + plane_noise_scale = plane_timestep if decode_noise_scale is None else decode_noise_scale + if isinstance(plane_noise_scale, list): + plane_noise_scale = plane_noise_scale[0] + epilogue_keyframes = self._rebuild_epilogue_keyframes( + carry_keyframes, + decode_timestep=plane_timestep, + decode_noise_scale=plane_noise_scale, + seed=ancestral_seed_base, + device=device, + dtype=torch.float32, + ) + + # The epilogue runs on the fully refined canvas, so a condition's moment is carried onto that grid the + # same way the rounds carry it. There is one window, so nothing is filtered or re-based. + epilogue_conditions = [ + (pixel * 2**temporal_upscalings, latent, strength, num_pixel_frames) + for pixel, latent, strength, num_pixel_frames in self.encode_conditions( + conditions, height, width, requested_frames, device=device, dtype=torch.float32 + ) + ] + + if detailing: + self._activate_adapters(detailing_adapters) + ( + epilogue_latents, + epilogue_conditioning_mask, + epilogue_clean_latents, + epilogue_video_coords, + epilogue_keyframes_mask, + _, + ) = self.prepare_latents( + condition_latents=epilogue_conditions, + keyframe_latents=[ + (position, epilogue_keyframes[:, :, index : index + 1], EPILOGUE_KEYFRAME_STRENGTH) + for index, position in enumerate(carry_positions) + ], + reference_latents=guide_latents if detailing else None, + reference_downscale_factor=detailing_reference_downscale_factor, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=height, + width=width, + num_frames=canvas_frames, + frame_rate=conditioning_fps, + noise_scale=stage_2_sigmas[0], + dtype=torch.float32, + device=device, + # One draw covers the whole canvas, so the noise is independent of how the tiles are laid out. + generator=torch.Generator(device=device).manual_seed(ancestral_seed_base + 2000), + latents=video_latents, + ) + epilogue_latent_frames = (canvas_frames - 1) // temporal_ratio + 1 + epilogue_latent_height = height // self.vae_spatial_compression_ratio + epilogue_latent_width = width // self.vae_spatial_compression_ratio + # Frozen stage-1 audio over the whole canvas, so video-audio cross attention still sees it. The + # returned audio is discarded; the shipped waveform is stage 1's. + epilogue_audio_latents = _audio_window_for_tile( + stage_1_audio_latents, + pixel_start=0, + tile_frames=canvas_frames, + playback_fps=frame_rate, + source_seconds=stage_1_seconds, + conditioning_fps=conditioning_fps, + audio_latents_per_second=audio_latents_per_second, + ) + epilogue_latents, _ = self.denoise( + latents=epilogue_latents, + conditioning_mask=epilogue_conditioning_mask, + clean_latents=epilogue_clean_latents, + video_coords=epilogue_video_coords, + keyframes_mask=epilogue_keyframes_mask, + prompt_embeds=video_prompt_embeds, + audio_prompt_embeds=audio_prompt_embeds, + prompt_attention_mask=connector_attention_mask, + sigmas=stage_2_sigmas, + frame_rate=conditioning_fps, + audio_latents=epilogue_audio_latents, + freeze_audio=True, + video_tile_plan=video_tile_plan( + epilogue_layout, + epilogue_video_coords, + epilogue_latent_frames, + epilogue_latent_height, + epilogue_latent_width, + ), + generator=generator, + use_cross_timestep=use_cross_timestep, + attention_kwargs=attention_kwargs, + progress_bar=progress_bar, + step_offset=step_offset, + callback_on_step_end=callback_on_step_end, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + step_offset += len(stage_2_sigmas) + video_latents = self._unpack_video_latents( + epilogue_latents[:, : epilogue_latent_frames * epilogue_latent_height * epilogue_latent_width], + epilogue_latent_frames, + epilogue_latent_height, + epilogue_latent_width, + ) + if detailing: + self._activate_adapters(base_adapters) + + progress_bar.close() + + # 10. Trim the padded tail. `requested_frames - 1` is a multiple of the VAE's temporal ratio, so each round's + # `N -> 2(N - 1) + 1` mapping keeps the trim on a latent boundary. + num_frames = (requested_frames - 1) * 2**temporal_upscalings + 1 + if num_frames > canvas_frames: + raise ValueError(f"Target {num_frames} frames exceeds the generated canvas of {canvas_frames}") + video_latents = video_latents[:, :, : (num_frames - 1) // self.vae_temporal_compression_ratio + 1] + + # 11. Decode. Audio is stage 1's, cut to the video's duration so a muxed container does not outlast the + # picture. + audio_latents = self._denormalize_audio_latents( + stage_1_audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std + ) + audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins) + + if output_type == "latent": + video = self._denormalize_latents( + video_latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + audio = audio_latents + else: + video_latents = video_latents.to(prompt_embeds.dtype) + if not self.vae.config.timestep_conditioning: + timestep = None + else: + noise = randn_tensor( + video_latents.shape, generator=generator, device=device, dtype=video_latents.dtype + ) + if not isinstance(decode_timestep, list): + decode_timestep = [decode_timestep] * batch_size + if decode_noise_scale is None: + decode_noise_scale = decode_timestep + elif not isinstance(decode_noise_scale, list): + decode_noise_scale = [decode_noise_scale] * batch_size + + timestep = torch.tensor(decode_timestep, device=device, dtype=video_latents.dtype) + decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=video_latents.dtype)[ + :, None, None, None, None + ] + video_latents = (1 - decode_noise_scale) * video_latents + decode_noise_scale * noise + + video_latents = self._denormalize_latents( + video_latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + video = self.vae.decode(video_latents.to(self.vae.dtype), timestep, return_dict=False)[0] + video = self.video_processor.postprocess_video(video, output_type=output_type) + + audio = self.vocoder( + self.audio_vae.decode(audio_latents.to(self.audio_vae.dtype), return_dict=False)[0] + ) + audio_samples = min( + audio.shape[-1], round(num_frames / frame_rate * self.vocoder.config.output_sampling_rate) + ) + audio = audio[..., :audio_samples] + + # Offload all models + self.maybe_free_model_hooks() + + if not return_dict: + return (video, audio) + + return LTX2PipelineOutput(frames=video, audio=audio) + finally: + if detailing: + self._activate_adapters(original_adapters) diff --git a/src/diffusers/pipelines/ltx2/utils.py b/src/diffusers/pipelines/ltx2/utils.py index 1ba0cfd6eae2..c46f82782af0 100644 --- a/src/diffusers/pipelines/ltx2/utils.py +++ b/src/diffusers/pipelines/ltx2/utils.py @@ -34,6 +34,10 @@ # Reduced schedule for super-resolution stage 2 (subset of distilled values) STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875] +# Schedule for the DFR pipeline's temporal refine rounds: the distilled schedule from its fifth sigma on. The rounds +# densify an already-structured canvas, so they skip the near-1.0 head of the schedule. +TEMPORAL_ROUND_DISTILLED_SIGMA_VALUES = DISTILLED_SIGMA_VALUES[4:] + # Default negative prompt from # https://github.com/Lightricks/LTX-2/blob/ae855f8538843825f9015a419cf4ba5edaf5eec2/packages/ltx-pipelines/src/ltx_pipelines/utils/constants.py#L131-L143 diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 376596d632ea..be4beee46905 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3122,6 +3122,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class LTX2DFRPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class LTX2HDRPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/models/transformers/test_models_transformer_ltx2.py b/tests/models/transformers/test_models_transformer_ltx2.py index ec0cc2d4a0d1..0131d372341b 100644 --- a/tests/models/transformers/test_models_transformer_ltx2.py +++ b/tests/models/transformers/test_models_transformer_ltx2.py @@ -116,6 +116,48 @@ def get_dummy_inputs(self) -> dict[str, torch.Tensor]: class TestLTX2Transformer(LTX2TransformerTesterConfig, ModelTesterMixin): """Core model tests for LTX2 Video Transformer.""" + def test_keyframes_abs_pos_embedding_marks_only_masked_tokens(self): + init_dict = self.get_init_dict() + init_dict["use_keyframes_abs_pos_embedding"] = True + torch.manual_seed(0) + model = self.model_class(**init_dict).to(torch_device).eval() + # The parameter is zero-initialized, so an untrained checkpoint is an exact no-op. + torch.nn.init.normal_(model.keyframes_abs_pos_embedding, std=0.1) + + inputs = self.get_dummy_inputs() + num_tokens = inputs["hidden_states"].shape[1] + keyframes_mask = torch.zeros( + (inputs["hidden_states"].shape[0], num_tokens, 1), device=torch_device, dtype=torch.float32 + ) + + with torch.no_grad(): + unmarked = model(**inputs, return_dict=False)[0] + all_zero_mask = model(**inputs, video_keyframes_mask=keyframes_mask, return_dict=False)[0] + keyframes_mask[:, : num_tokens // 2] = 1.0 + half_marked = model(**inputs, video_keyframes_mask=keyframes_mask, return_dict=False)[0] + + # An all-zero mask marks nothing, so it must match omitting the mask. + assert torch.allclose(unmarked, all_zero_mask, atol=1e-5) + assert not torch.allclose(unmarked, half_marked, atol=1e-5) + + def test_keyframes_mask_is_ignored_without_the_embedding(self): + torch.manual_seed(0) + model = self.model_class(**self.get_init_dict()).to(torch_device).eval() + assert not hasattr(model, "keyframes_abs_pos_embedding") + + inputs = self.get_dummy_inputs() + keyframes_mask = torch.ones( + (inputs["hidden_states"].shape[0], inputs["hidden_states"].shape[1], 1), + device=torch_device, + dtype=torch.float32, + ) + + with torch.no_grad(): + without_mask = model(**inputs, return_dict=False)[0] + with_mask = model(**inputs, video_keyframes_mask=keyframes_mask, return_dict=False)[0] + + assert torch.allclose(without_mask, with_mask, atol=1e-5) + class TestLTX2TransformerMemory(LTX2TransformerTesterConfig, MemoryTesterMixin): """Memory optimization tests for LTX2 Video Transformer.""" diff --git a/tests/pipelines/ltx2/test_ltx2_dfr_layout.py b/tests/pipelines/ltx2/test_ltx2_dfr_layout.py new file mode 100644 index 000000000000..d622da252089 --- /dev/null +++ b/tests/pipelines/ltx2/test_ltx2_dfr_layout.py @@ -0,0 +1,257 @@ +# Copyright 2026 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import itertools + +import pytest +import torch + +from diffusers.pipelines.ltx2.dfr_layout import ( + epilogue_tiles, + pixel_to_latent_index, + rectangular_mask_1d, + resolve_canvas, + split_by_count, + split_canvas_at_seams, + temporal_tile_plan, + trapezoidal_mask_1d, +) + + +class TestResolveCanvas: + def test_exact_multiple_of_a_segment_is_not_padded(self): + canvas_frames, segment, positions = resolve_canvas(97) + + assert (canvas_frames, segment) == (97, 32) + assert positions == [32, 64, 96] + + def test_smaller_segment_wins_when_it_pads_strictly_less(self): + canvas_frames, segment, positions = resolve_canvas(121) + + assert (canvas_frames, segment) == (121, 24) + assert positions == [24, 48, 72, 96, 120] + + def test_tail_is_padded_up_to_a_whole_segment(self): + canvas_frames, segment, positions = resolve_canvas(41) + + assert canvas_frames > 41 + assert (canvas_frames - 1) % segment == 0 + assert positions[-1] == canvas_frames - 1 + + def test_segment_grid_scales_with_the_vae_temporal_ratio(self): + canvas_frames, segment, positions = resolve_canvas(9, temporal_compression_ratio=2) + + assert (canvas_frames, segment) == (9, 8) + assert positions == [8] + + def test_a_request_off_the_latent_grid_is_rejected(self): + with pytest.raises(ValueError, match="num_frames"): + resolve_canvas(100) + + +class TestPixelToLatentIndex: + def test_a_border_frame_maps_to_its_latent(self): + assert pixel_to_latent_index(64) == 8 + + def test_a_frame_off_the_border_is_rejected(self): + with pytest.raises(ValueError, match="latent border"): + pixel_to_latent_index(65) + + +class TestSegmentDealing: + def test_leftover_segments_go_to_the_leading_tiles(self): + # 5 segments over 3 tiles deals 2/2/1, so the tiles end on the 2nd, 4th and 5th seam. + intervals = split_canvas_at_seams([0, 1, 2, 3, 4, 5], num_tiles=3, overlap=0, dim_size=6) + assert [interval.end - 1 for interval in intervals] == [2, 4, 5] + + def test_tile_count_is_clamped_to_the_segment_count(self): + # 2 segments cannot fill 4 tiles; each tile owns one. + intervals = split_canvas_at_seams([0, 1, 2], num_tiles=4, overlap=0, dim_size=3) + assert [(interval.start, interval.end) for interval in intervals] == [(0, 2), (2, 3)] + + +class TestTemporalTilePlan: + def test_a_single_tile_covers_the_whole_canvas(self): + (tile,) = temporal_tile_plan([32, 64], 65, 1) + + assert (tile.pixel_start, tile.pixel_end) == (0, 64) + assert (tile.interval.start, tile.interval.end) == (0, 9) + assert tile.interval.left_ramp == 0 + assert tile.anchors == (32, 64) + assert tile.slots == (16, 48) + + def test_tiles_are_gapless_after_dropping_their_ramps(self): + tiles = temporal_tile_plan([32, 64, 96, 128], 129, 4) + + covered = [] + for tile in tiles: + covered.extend(range(tile.interval.start + tile.interval.left_ramp, tile.interval.end)) + assert covered == list(range(0, 17)) + + def test_non_first_tiles_reach_back_one_segment_for_the_shared_seam(self): + first, second = temporal_tile_plan([32, 64, 96, 128], 129, 2) + + assert second.pixel_start == 32 + assert first.pixel_end == 64 + assert second.pixel_end == 128 + # The ramp swallows the lead-in *and* the seam latent, so the previous tile keeps the shared keyframe. + assert second.interval.start + second.interval.left_ramp == first.interval.end + + def test_the_first_tile_window_start_contributes_no_anchor(self): + tiles = temporal_tile_plan([32, 64, 96], 97, 2) + + assert 0 not in tiles[0].anchors + assert tiles[1].pixel_start in tiles[1].anchors + + def test_owned_segment_runs_are_balanced_largest_first(self): + tiles = temporal_tile_plan([32, 64, 96], 97, 2) + + # 3 segments / 2 tiles: first owns 2, second owns 1. Lead-in re-invents slot 48; earlier tile wins. + assert tiles[0].slots == (16, 48) + assert tiles[1].slots == (48, 80) + + def test_every_seam_in_a_window_is_an_anchor(self): + tiles = temporal_tile_plan([32, 64, 96, 128], 129, 2) + + assert tiles[0].anchors == (32, 64) + assert tiles[1].anchors == (32, 64, 96, 128) + + def test_a_ramp_drop_stitch_reproduces_the_canvas(self): + tiles = temporal_tile_plan([32, 64, 96, 128], 129, 2) + tile_latents = [ + torch.arange(tile.interval.start, tile.interval.end, dtype=torch.float32).reshape(1, 1, -1, 1, 1) + for tile in tiles + ] + + stitched = torch.cat([latent[:, :, tile.interval.left_ramp :] for latent, tile in zip(tile_latents, tiles)], 2) + + assert stitched.shape == (1, 1, 17, 1, 1) + assert torch.equal(stitched.flatten(), torch.arange(17, dtype=torch.float32)) + + +class TestSplitCanvasAtSeams: + def test_seams_must_end_on_the_last_cell(self): + with pytest.raises(ValueError, match="last cell"): + split_canvas_at_seams([0, 4, 8], num_tiles=2, overlap=5, dim_size=13) + + def test_seams_must_be_increasing(self): + with pytest.raises(ValueError, match="strictly increasing"): + split_canvas_at_seams([0, 8, 4, 12], num_tiles=2, overlap=5, dim_size=13) + + +class TestSplitByCount: + @pytest.mark.parametrize( + ("dim_size", "num_tiles", "overlap"), + list(itertools.product((13, 16, 17, 21, 34), (1, 2, 3, 4), (0, 2, 6))), + ) + def test_tiles_cover_the_axis_and_share_exactly_the_overlap(self, dim_size, num_tiles, overlap): + if num_tiles > 1 and overlap > dim_size - num_tiles: + pytest.skip("layout the caller is required to clamp first") + intervals = split_by_count(dim_size, num_tiles, overlap) + + assert len(intervals) == num_tiles + assert intervals[0].start == 0 + assert intervals[-1].end == dim_size + for earlier, later in itertools.pairwise(intervals): + assert earlier.end - later.start == overlap + assert earlier.right_ramp == overlap + assert later.left_ramp == overlap + + def test_trapezoidal_weights_over_a_split_sum_to_one_everywhere(self): + dim_size, num_tiles, overlap = 34, 3, 6 + summed = torch.zeros(dim_size) + for interval in split_by_count(dim_size, num_tiles, overlap): + mask = trapezoidal_mask_1d(interval.end - interval.start, interval.left_ramp, interval.right_ramp) + summed[interval.start : interval.end] += mask + + assert torch.allclose(summed, torch.ones(dim_size), atol=1e-6) + + def test_a_single_tile_is_the_whole_axis_unramped(self): + (interval,) = split_by_count(9, 1, 6) + + assert (interval.start, interval.end) == (0, 9) + assert (interval.left_ramp, interval.right_ramp) == (0, 0) + + +class TestRectangularMask: + def test_the_lead_in_is_dropped_and_the_rest_kept_whole(self): + assert torch.equal(rectangular_mask_1d(5, 2), torch.tensor([0.0, 0.0, 1.0, 1.0, 1.0])) + # The opening tile has no lead-in and so keeps every cell. + assert torch.equal(rectangular_mask_1d(4, 0), torch.ones(4)) + + def test_weights_over_a_seam_cut_sum_to_one_everywhere(self): + seams, dim_size = [0, 4, 8, 12, 16], 17 + summed = torch.zeros(dim_size) + for interval in split_canvas_at_seams(seams, 2, overlap=5, dim_size=dim_size): + summed[interval.start : interval.end] += rectangular_mask_1d( + interval.end - interval.start, interval.left_ramp + ) + + assert torch.equal(summed, torch.ones(dim_size)) + + +class TestTrapezoidalMask: + def test_ramps_are_the_interior_of_a_linspace(self): + assert torch.allclose(trapezoidal_mask_1d(5, 3, 0), torch.tensor([0.25, 0.5, 0.75, 1.0, 1.0])) + assert torch.allclose(trapezoidal_mask_1d(5, 0, 3), torch.tensor([1.0, 1.0, 0.75, 0.5, 0.25])) + # A tile with neither neighbour keeps every cell. + assert torch.equal(trapezoidal_mask_1d(4, 0, 0), torch.ones(4)) + + +class TestEpilogueTiles: + def test_epilogue_tiles_cover_the_canvas_with_unit_weight(self): + tiles = epilogue_tiles(latent_shape=(9, 16, 16), frame_tiles=2, frame_seams=[3, 6]) + + assert len(tiles) == 2 * 2 * 2 + summed = torch.zeros(9, 16, 16) + for frames, heights, widths, weight in tiles: + summed[frames, heights, widths] += weight + assert torch.allclose(summed, torch.ones_like(summed), atol=1e-6) + + def test_epilogue_temporal_tiles_are_cut_on_the_last_round_seams(self): + # Ten segments over 61 latent frames dealt to four tiles: the leading two take three each. + tiles = epilogue_tiles(latent_shape=(61, 8, 8), frame_tiles=4, frame_seams=[6 * i for i in range(1, 11)]) + + windows = sorted({(frames.start, frames.stop) for frames, _, _, _ in tiles}) + assert windows == [(0, 19), (12, 37), (30, 49), (42, 61)] + # The run-up is dropped outright rather than blended, so every non-first window opens on exact zeros and the + # window before it keeps the seam cell. + head = next(weight for frames, _, _, weight in tiles if frames.start == 12)[:8, 0, 0] + assert torch.equal(head, torch.tensor([0.0] * 7 + [1.0])) + + def test_epilogue_temporal_tiles_blend_when_the_seams_are_not_interior(self): + # A canvas with no refine rounds behind it has only its own end cell as a "seam", which cuts nothing. + tiles = epilogue_tiles(latent_shape=(9, 8, 8), frame_tiles=2, frame_seams=[8]) + + # Nothing is cut, so the later window ramps in instead of opening on zeros. + head = next(weight for frames, _, _, weight in tiles if frames.start > 0)[0, 0, 0] + assert 0.0 < head < 1.0 + summed = torch.zeros(9, 8, 8) + for frames, heights, widths, weight in tiles: + summed[frames, heights, widths] += weight + assert torch.allclose(summed, torch.ones_like(summed), atol=1e-6) + + def test_epilogue_tiling_falls_back_on_an_axis_too_short_to_split(self): + tiles = epilogue_tiles(latent_shape=(3, 16, 16), frame_tiles=4, frame_seams=[1, 2]) + + # 3 cells cannot hold 4 tiles, so the temporal axis runs whole and only the spatial split survives. + assert len(tiles) == 1 * 2 * 2 + assert {(frames.start, frames.stop) for frames, _, _, _ in tiles} == {(0, 3)} + + def test_epilogue_tiling_clamps_an_overlap_the_axis_cannot_hold(self): + tiles = epilogue_tiles(latent_shape=(9, 3, 16), frame_tiles=1, frame_seams=[]) + + # Height is 3 cells, so the requested 12-cell overlap is clamped to 1 and the two tiles still tile it. + height_slices = sorted({(heights.start, heights.stop) for _, heights, _, _ in tiles}) + assert height_slices == [(0, 2), (1, 3)] diff --git a/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py new file mode 100644 index 000000000000..baea5b1d58dc --- /dev/null +++ b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py @@ -0,0 +1,772 @@ +# Copyright 2026 The HuggingFace Team. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import numpy as np +import PIL.Image +import pytest +import torch +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from diffusers import ( + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + LTX2DFRPipeline, + LTX2VideoTransformer3DModel, +) +from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel, LTX2TextConnectors +from diffusers.pipelines.ltx2.dfr_layout import LTX2DFREpilogueTile, epilogue_tiles, video_tile_plan +from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition +from diffusers.pipelines.ltx2.pipeline_ltx2_dfr import ( + ANCHOR_KEYFRAME_STRENGTH, + EPILOGUE_KEYFRAME_STRENGTH, + MAX_CONDITIONING_FPS, + _audio_window_for_tile, +) +from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder +from diffusers.utils.torch_utils import randn_tensor + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin + + +enable_full_determinism() + + +class LTX2DFRPipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = LTX2DFRPipeline + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "num_frames", "frame_rate", "prompt_embeds", "temporal_upscalings"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 32, 32) + # DFR stages use fixed distilled schedules (`stage_*_sigmas`); `latents` cannot seed stage 1, which always + # generates at half resolution. Videos are counted with `num_videos_per_prompt`. + optional_input_params = BasePipelineTesterConfig.optional_input_params - { + "num_inference_steps", + "num_images_per_prompt", + "latents", + } + + base_text_encoder_ckpt_id = "hf-internal-testing/tiny-gemma3" + + def get_dummy_components(self): + tokenizer = AutoTokenizer.from_pretrained(self.base_text_encoder_ckpt_id) + text_encoder = Gemma3ForConditionalGeneration.from_pretrained(self.base_text_encoder_ckpt_id) + + torch.manual_seed(0) + transformer = LTX2VideoTransformer3DModel( + in_channels=4, + out_channels=4, + patch_size=1, + patch_size_t=1, + num_attention_heads=2, + attention_head_dim=8, + cross_attention_dim=16, + audio_in_channels=4, + audio_out_channels=4, + audio_num_attention_heads=2, + audio_attention_head_dim=4, + audio_cross_attention_dim=8, + num_layers=2, + qk_norm="rms_norm_across_heads", + caption_channels=text_encoder.config.text_config.hidden_size, + rope_double_precision=False, + rope_type="split", + vae_scale_factors=(2, 2, 2), + use_keyframes_abs_pos_embedding=True, + ) + # Zero-init would make a broken `video_keyframes_mask` an exact no-op. + torch.nn.init.normal_(transformer.keyframes_abs_pos_embedding, std=0.1) + + torch.manual_seed(0) + connectors = LTX2TextConnectors( + caption_channels=text_encoder.config.text_config.hidden_size, + text_proj_in_factor=text_encoder.config.text_config.num_hidden_layers + 1, + video_connector_num_attention_heads=4, + video_connector_attention_head_dim=8, + video_connector_num_layers=1, + video_connector_num_learnable_registers=None, + audio_connector_num_attention_heads=4, + audio_connector_attention_head_dim=8, + audio_connector_num_layers=1, + audio_connector_num_learnable_registers=None, + connector_rope_base_seq_len=32, + rope_theta=10000.0, + rope_double_precision=False, + causal_temporal_positioning=False, + rope_type="split", + ) + + torch.manual_seed(0) + vae = AutoencoderKLLTX2Video( + in_channels=3, + out_channels=3, + latent_channels=4, + block_out_channels=(8,), + decoder_block_out_channels=(8,), + layers_per_block=(1,), + decoder_layers_per_block=(1, 1), + spatio_temporal_scaling=(True,), + decoder_spatio_temporal_scaling=(True,), + decoder_inject_noise=(False, False), + downsample_type=("spatial",), + upsample_residual=(False,), + upsample_factor=(1,), + timestep_conditioning=False, + patch_size=1, + patch_size_t=1, + encoder_causal=True, + decoder_causal=False, + ) + vae.use_framewise_encoding = False + vae.use_framewise_decoding = False + + torch.manual_seed(0) + audio_vae = AutoencoderKLLTX2Audio( + base_channels=4, + output_channels=2, + ch_mult=(1,), + num_res_blocks=1, + attn_resolutions=None, + in_channels=2, + resolution=32, + latent_channels=2, + norm_type="pixel", + causality_axis="height", + dropout=0.0, + mid_block_add_attention=False, + sample_rate=16000, + mel_hop_length=160, + is_causal=True, + mel_bins=8, + ) + + torch.manual_seed(0) + vocoder = LTX2Vocoder( + in_channels=audio_vae.config.output_channels * audio_vae.config.mel_bins, + hidden_channels=32, + out_channels=2, + upsample_kernel_sizes=[4, 4], + upsample_factors=[2, 2], + resnet_kernel_sizes=[3], + resnet_dilations=[[1, 3, 5]], + leaky_relu_negative_slope=0.1, + output_sampling_rate=16000, + ) + + # `mid_channels` feeds a GroupNorm with 32 groups, so it cannot be shrunk below 32. + torch.manual_seed(0) + latent_upsampler = LTX2LatentUpsamplerModel( + in_channels=4, + mid_channels=32, + num_blocks_per_stage=1, + dims=3, + spatial_upsample=True, + temporal_upsample=False, + use_rational_resampler=False, + ) + + torch.manual_seed(0) + temporal_latent_upsampler = LTX2LatentUpsamplerModel( + in_channels=4, + mid_channels=32, + num_blocks_per_stage=1, + dims=3, + spatial_upsample=False, + temporal_upsample=True, + ) + + scheduler = FlowMatchEulerDiscreteScheduler() + + return { + "transformer": transformer, + "vae": vae, + "audio_vae": audio_vae, + "scheduler": scheduler, + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "connectors": connectors, + "vocoder": vocoder, + "latent_upsampler": latent_upsampler, + "temporal_latent_upsampler": temporal_latent_upsampler, + "processor": None, + "prompt_enhancer": None, + "duration_head": None, + } + + def get_dummy_inputs(self): + # Dummy VAE temporal ratio is 2, so the segment grid is 6 / 8 pixel frames (not 24 / 32). + return { + "prompt": "a robot dancing", + "generator": self.get_generator(0), + "height": 32, + "width": 32, + "num_frames": 9, + "frame_rate": 25.0, + "stage_1_sigmas": [1.0, 0.5], + "stage_2_sigmas": [0.75, 0.25], + "temporal_round_sigmas": [0.75, 0.25], + "use_cross_timestep": False, + "max_sequence_length": 16, + "output_type": "pt", + } + + +class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin): + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3): + # Mixin default copies `num_inference_steps` into batched inputs (DFR has none). Tolerance is 1e-3 because + # batched vs single matmuls on accelerator reduce in a different order (max diff ~6e-4 on CUDA). + super().test_inference_batch_single_identical( + batch_size=batch_size, + expected_max_diff=expected_max_diff, + additional_params_copy_to_batched_inputs=[], + ) + + def test_padded_canvas_is_trimmed_back_to_the_request(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + inputs = self.get_dummy_inputs() + # 10 content frames: segment 6 pads by 2, so the canvas is 13 and must be trimmed to 11. + inputs["num_frames"] = 11 + output = pipe(**inputs) + + assert output.frames.shape[1] == 11 + + def test_temporal_upsample_round_doubles_the_frame_count(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + inputs = self.get_dummy_inputs() + inputs["temporal_upscalings"] = 1 + output = pipe(**inputs) + + assert output.frames.shape[1] == (self.output_shape[0] - 1) * 2 + 1 + # One 8-frame segment -> one tile: 2 + 2 stage steps plus 2 tile steps. `callback_on_step_end` indexes this. + assert pipe.num_timesteps == 2 + 2 + 2 + + def test_temporal_round_tiles_get_distinct_ancestral_noise(self): + # Tiles are positionally identical; each must draw from `seed + 1000 * round + tile`, not a shared stream. + import diffusers.pipelines.ltx2.pipeline_ltx2_dfr as dfr_module + + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + captured = [] + original_step = dfr_module.ancestral_euler_step + + def capture(sample, denoised, sigma, sigma_next, eta, noise): + captured.append(noise.clone()) + return original_step(sample, denoised, sigma, sigma_next, eta, noise) + + inputs = self.get_dummy_inputs() + inputs["num_frames"] = 17 + inputs["temporal_upscalings"] = 1 + dfr_module.ancestral_euler_step = capture + try: + pipe(**inputs) + finally: + dfr_module.ancestral_euler_step = original_step + + assert len(captured) >= 2, f"expected ancestral draws from at least 2 tiles, got {len(captured)}" + + first_tile = captured[0] + expected = randn_tensor( + first_tile.shape, + generator=torch.Generator(device=torch_device).manual_seed(0 + 1000 * 1 + 0), + device=torch.device(torch_device), + dtype=first_tile.dtype, + ) + assert torch.equal(first_tile, expected) + + @pytest.mark.parametrize( + ("num_frames", "temporal_upscalings", "output_frames", "num_timesteps"), + [ + # No rounds: the epilogue is one 2-step loop, whatever its tiling, because the tiles live inside the + # transformer call rather than around the schedule. + (9, 0, 9, 2 + 2 + 2), + # One round leaves a 33-frame canvas the epilogue then cuts 2 ways in time as well; still one loop. + (17, 1, 33, 2 + 2 + 2 * 2 + 2), + ], + ) + def test_the_epilogue_runs_on_the_canvas_the_earlier_stages_left_behind( + self, num_frames, temporal_upscalings, output_frames, num_timesteps + ): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + inputs = self.get_dummy_inputs() + inputs["num_frames"] = num_frames + inputs["temporal_upscalings"] = temporal_upscalings + inputs["spatial_upscalings"] = 2 + output = pipe(**inputs) + + # `height` / `width` are the output either way; only the stages below them move. + assert output.frames.shape[1:] == (output_frames, *self.output_shape[1:]) + assert pipe.num_timesteps == num_timesteps + + @pytest.mark.parametrize(("temporal_upscalings", "expected_pixel"), [(0, 7), (1, 14), (2, 28)]) + def test_a_condition_keeps_its_moment_through_the_refine_rounds(self, temporal_upscalings, expected_pixel): + # `index` is a latent index on the canvas the caller asked for. Each round doubles the canvas and the frame + # rate, so a fixed moment moves to twice its pixel position, and the epilogue runs on the refined canvas. The + # scaled position does not generally land on a latent boundary, so it has to travel as a pixel index -- the + # dummy VAE ratio is 2, latent 4 is pixel 7, and one round puts it at 14, which is not on the latent grid. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + coords = [] + original = pipe.prepare_latents + + def capture(**kwargs): + result = original(**kwargs) + coords.append((kwargs["frame_rate"], result[3], kwargs.get("num_frames"))) + return result + + image = PIL.Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)) + inputs = self.get_dummy_inputs() + inputs["num_frames"] = 33 + inputs["temporal_upscalings"] = temporal_upscalings + inputs["spatial_upscalings"] = 2 + inputs["conditions"] = [LTX2VideoCondition(frames=image, index=4, strength=1.0)] + pipe.prepare_latents = capture + try: + pipe(**inputs) + finally: + del pipe.prepare_latents + + # The epilogue is the last pass and covers the whole canvas, so its condition sits where the caller asked. + # Read the position back off the RoPE coordinates the pass will actually attend with, in seconds / fps. + frame_rate, video_coords, num_frames = coords[-1] + latent_frames = (num_frames - 1) // pipe.vae_temporal_compression_ratio + 1 + tokens_per_frame = (32 // pipe.vae_spatial_compression_ratio) ** 2 + condition_token = latent_frames * tokens_per_frame + start = video_coords[0, 0, condition_token, 0].item() * frame_rate + assert round(start) == expected_pixel, f"condition landed at pixel {start}, expected {expected_pixel}" + # A single-pixel-frame keyframe spans exactly one frame. + end = video_coords[0, 0, condition_token, 1].item() * frame_rate + assert round(end - start) == 1 + + def test_the_epilogue_resolution_rule_is_reported_before_the_generic_one(self): + # `spatial_upscalings=2` needs multiples of `2 ** upscalings * vae_spatial_compression_ratio` -- 128 on the + # shipped VAE, which is why UHD is 3840x2176 rather than 3840x2160. A caller who misses it must be told that + # divisor, not the looser one every LTX-2 pipeline applies. + pipe = self.get_pipeline(**self.get_dummy_components()) + divisor = 4 * pipe.vae_spatial_compression_ratio + + inputs = self.get_dummy_inputs() + inputs["spatial_upscalings"] = 2 + inputs["height"] = inputs["width"] = 32 + divisor // 2 + with pytest.raises(ValueError, match=f"divisible by {divisor}"): + pipe(**inputs) + + def test_the_epilogue_keeps_every_batch_element_distinct(self): + # The carry keyframes are rebuilt through RGB one plane at a time, and every batch element has its own picture + # to stretch. Collapsing the batch there would hand one element's frames to all of them. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + inputs = self.get_dummy_inputs() + inputs["spatial_upscalings"] = 2 + inputs["num_videos_per_prompt"] = 2 + output = pipe(**inputs) + + assert output.frames.shape[0] == 2 + + def _epilogue_inputs(self, **overrides): + inputs = self.get_dummy_inputs() + inputs["num_frames"] = 33 + inputs["temporal_upscalings"] = 1 + inputs["spatial_upscalings"] = 2 + inputs.update(overrides) + return inputs + + def _epilogue_plan(self, pipe, keyframe_positions=(8, 16, 24, 32, 40, 48, 56, 64)): + """The epilogue's real tiling and token plan for a 33-frame request with one refine round.""" + # One round turns the 33-frame canvas into 65 frames; at the dummy ratio of 2 that is a 33-frame latent grid, + # cut on the round's seams and split two ways on each spatial axis. + latent_frames = 33 + latent_height = latent_width = 32 // pipe.vae_spatial_compression_ratio + seams = [position // pipe.vae_temporal_compression_ratio for position in keyframe_positions] + tiles = epilogue_tiles( + latent_shape=(latent_frames, latent_height, latent_width), frame_tiles=2, frame_seams=seams + ) + keyframe = torch.randn(1, 4, 1, latent_height, latent_width, device=torch_device) + _, _, _, video_coords, _, _ = pipe.prepare_latents( + keyframe_latents=[(position, keyframe, EPILOGUE_KEYFRAME_STRENGTH) for position in keyframe_positions], + num_channels_latents=4, + height=32, + width=32, + num_frames=(latent_frames - 1) * pipe.vae_temporal_compression_ratio + 1, + dtype=torch.float32, + device=torch_device, + generator=self.get_generator(0), + ) + plan = video_tile_plan(tiles, video_coords, latent_frames, latent_height, latent_width) + return tiles, video_coords, plan, len(keyframe_positions), latent_frames * latent_height * latent_width + + def test_a_tiled_epilogue_pass_routes_every_token_with_unit_total_weight(self): + # Conditionings are attached once on the whole canvas and filtered per tile at the token level, so the routing + # is what decides which window sees which keyframe. Every token -- base grid and appended alike -- must come + # out with a total weight of exactly one, or the blended canvas is dimmed, doubled, or missing tokens. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + _, video_coords, plan, _, _ = self._epilogue_plan(pipe) + + totals = torch.zeros(video_coords.shape[2], device=video_coords.device) + for tile in plan: + totals.index_add_(0, tile.keep, tile.weights.to(totals.dtype)) + assert torch.allclose(totals, torch.ones_like(totals), atol=1e-6) + + def test_a_keyframe_two_epilogue_windows_share_is_a_single_token(self): + # A later window's dropped run-up is content the window before it owns. There is one token per keyframe and + # every window covering it reads that same token, so the two cannot disagree about a frame they both hold. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + tiles, _, plan, num_keyframes, first_keyframe_token = self._epilogue_plan(pipe) + assert len({(frames.start, frames.stop) for frames, _, _, _ in tiles}) > 1, "expected 2 time windows" + + tokens_per_keyframe = (32 // pipe.vae_spatial_compression_ratio) ** 2 + shared = 0 + for index in range(num_keyframes): + token = first_keyframe_token + index * tokens_per_keyframe + windows = { + (frames.start, frames.stop) + for (frames, _, _, _), tile in zip(tiles, plan) + if bool((tile.keep == token).any()) + } + assert windows, f"keyframe token {token} reaches no window" + shared += len(windows) > 1 + assert shared, "no keyframe token is shared across windows" + + def test_the_epilogue_is_given_its_keyframes_rather_than_regenerating_them(self): + # The carry keyframes are frames the refine rounds already settled. Re-attaching them as generated slots asked + # the pass to invent content it already had; they arrive rebuilt at the output resolution and pinned clean. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + calls = [] + original = pipe.prepare_latents + + def capture(**kwargs): + result = original(**kwargs) + calls.append((kwargs, result)) + return result + + pipe.prepare_latents = capture + try: + pipe(**self._epilogue_inputs()) + finally: + del pipe.prepare_latents + + # The epilogue is the last pass, so its call is the one left standing. Read the answer off the state it built + # rather than the arguments it was handed. + _, (_, conditioning_mask, _, _, keyframes_mask, slot_token_slice) = calls[-1] + assert slot_token_slice is None, "the epilogue must not allocate generated slot tokens" + + # Everything appended past the base grid is the carry keyframes, at the output resolution. + tokens_per_frame = (32 // pipe.vae_spatial_compression_ratio) ** 2 + base_tokens = (65 - 1) // pipe.vae_temporal_compression_ratio + 1 + appended = conditioning_mask[:, base_tokens * tokens_per_frame :] + assert appended.shape[1] == 8 * tokens_per_frame, "expected the 8 carry keyframes of a one-round canvas" + # Pinned fully clean... + assert torch.allclose(appended, torch.full_like(appended, EPILOGUE_KEYFRAME_STRENGTH)) + # ...and unmarked, which is what distinguishes given content from a generated slot. + assert torch.count_nonzero(keyframes_mask[:, base_tokens * tokens_per_frame :]) == 0 + + @pytest.mark.parametrize(("pixel_start", "first_latent"), [(0, 0), (48, 20)]) + def test_tile_audio_is_the_stage_1_window_on_the_playback_clock(self, pixel_start, first_latent): + # 40 stage-1 audio frames spanning 2 s. A refine round doubles the frame rate, so a 48-frame tile at 48 fps + # is 1 s of sound however long the canvas got -- the tile starting at 48 hears the second second, and only it. + source = torch.arange(40, dtype=torch.float32).reshape(1, 40, 1) + window = _audio_window_for_tile( + source, + pixel_start=pixel_start, + tile_frames=48, + playback_fps=48.0, + source_seconds=2.0, + conditioning_fps=48.0, + audio_latents_per_second=20.0, + ) + + assert window.shape == (1, 20, 1) + expected = torch.arange(first_latent, first_latent + 20, dtype=torch.float32) + assert torch.allclose(window.flatten(), expected) + + def test_a_single_tile_plan_reproduces_the_untiled_pass(self): + # `_video_tile_plan` is the whole correctness surface of the tiled path: which tokens a tile sees, how its + # positions are rebased, and what weight its prediction carries. Degenerate it to one tile covering the whole + # canvas and the two paths must agree exactly -- anything else is a bug in the plan rather than in the tiling. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe._interrupt = False + pipe._current_timestep = None + pipe._attention_kwargs = None + + keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) + reference = torch.randn(1, 4, 2, 8, 8, device=torch_device) + prepared = pipe.prepare_latents( + keyframe_latents=[(8, keyframe, 0.95)], + slot_frame_indices=[4], + reference_latents=reference, + reference_downscale_factor=2, + num_channels_latents=4, + height=32, + width=32, + num_frames=9, + noise_scale=0.9, + dtype=torch.float32, + device=torch_device, + generator=self.get_generator(0), + ) + latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, _ = prepared + + text_embeds, text_mask = pipe.encode_prompt( + prompt="a robot dancing", max_sequence_length=16, device=torch_device + ) + video_embeds, audio_embeds, connector_mask = pipe.connectors( + text_embeds, text_mask, padding_side=pipe.tokenizer_padding_side + ) + # 9 pixel frames at the dummy VAE's temporal ratio of 2 is a 5-frame latent grid. + whole_canvas = [ + LTX2DFREpilogueTile( + frames=slice(0, 5), + heights=slice(0, 16), + widths=slice(0, 16), + blend_weight=torch.ones(5, 16, 16), + ) + ] + plan = video_tile_plan(whole_canvas, video_coords, 5, 16, 16) + assert len(plan) == 1 + assert plan[0].keep.numel() == latents.shape[1], "one tile must keep every token" + assert torch.equal(plan[0].coords, video_coords) + + audio_latents = pipe.prepare_audio_latents( + num_channels_latents=pipe.audio_latent_channels, + audio_latent_length=8, + num_mel_bins=pipe.audio_mel_bins, + dtype=torch.float32, + device=torch_device, + generator=self.get_generator(0), + ) + + def run(video_tile_plan): + out, _ = pipe.denoise( + latents=latents, + conditioning_mask=conditioning_mask, + clean_latents=clean_latents, + video_coords=video_coords, + keyframes_mask=keyframes_mask, + prompt_embeds=video_embeds, + audio_prompt_embeds=audio_embeds, + prompt_attention_mask=connector_mask, + sigmas=[0.9, 0.7], + frame_rate=25.0, + audio_latents=audio_latents, + video_tile_plan=video_tile_plan, + generator=self.get_generator(0), + ) + return out + + assert torch.allclose(run(plan), run(None), atol=1e-5) + + def test_ancestral_step_does_not_erode_conditioning(self): + # Ancestral Euler noises every token; the strength-0.95 blend has to be re-applied or tile seams drift. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe._interrupt = False + pipe._current_timestep = None + pipe._attention_kwargs = None + + keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) + strength = 0.95 + latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, _ = pipe.prepare_latents( + keyframe_latents=[(8, keyframe, strength)], + num_channels_latents=4, + height=32, + width=32, + num_frames=9, + noise_scale=0.975, + dtype=torch.float32, + device=torch_device, + generator=self.get_generator(0), + ) + block = slice(latents.shape[1] - pipe._pack_latents(keyframe).shape[1], latents.shape[1]) + packed_keyframe = pipe._pack_latents(keyframe) + + text_embeds, text_mask = pipe.encode_prompt( + prompt="a robot dancing", max_sequence_length=16, device=torch_device + ) + video_embeds, audio_embeds, connector_mask = pipe.connectors( + text_embeds, text_mask, padding_side=pipe.tokenizer_padding_side + ) + out, _ = pipe.denoise( + latents=latents, + conditioning_mask=conditioning_mask, + clean_latents=clean_latents, + video_coords=video_coords, + keyframes_mask=keyframes_mask, + prompt_embeds=video_embeds, + audio_prompt_embeds=audio_embeds, + prompt_attention_mask=connector_mask, + sigmas=[0.975, 0.9, 0.7], + frame_rate=25.0, + audio_latents=pipe.prepare_audio_latents( + num_channels_latents=pipe.audio_latent_channels, + audio_latent_length=8, + num_mel_bins=pipe.audio_mel_bins, + dtype=torch.float32, + device=torch_device, + generator=self.get_generator(0), + ), + ancestral_eta=0.5, + generator=self.get_generator(0), + ancestral_generator=torch.Generator(device=torch_device).manual_seed(1), + ) + + anchor_out = out[:, block].float() + cos = torch.nn.functional.cosine_similarity(anchor_out.flatten(), packed_keyframe.float().flatten(), dim=0) + assert cos > 0.9, f"anchor tokens drifted from their conditioned content (cos={cos:.3f})" + + def test_partially_conditioned_keyframe_starts_from_its_clean_content(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) + latents, conditioning_mask, clean_latents, _, _, _ = pipe.prepare_latents( + keyframe_latents=[(8, keyframe, 0.95)], + num_channels_latents=4, + height=32, + width=32, + num_frames=9, + noise_scale=0.0, + dtype=torch.float32, + device=torch_device, + ) + + packed_keyframe = pipe._pack_latents(keyframe) + block = slice(latents.shape[1] - packed_keyframe.shape[1], latents.shape[1]) + assert torch.allclose(conditioning_mask[:, block], torch.full_like(conditioning_mask[:, block], 0.95)) + assert torch.allclose(clean_latents[:, block], packed_keyframe) + assert torch.allclose(latents[:, block], packed_keyframe * 0.95, atol=1e-6) + + def test_keyframe_marker_reaches_the_transformer(self): + components = self.get_dummy_components() + pipe = self.get_pipeline(**components).to(torch_device) + + marked = pipe(**self.get_dummy_inputs()).frames + + with torch.no_grad(): + components["transformer"].keyframes_abs_pos_embedding.zero_() + unmarked = pipe(**self.get_dummy_inputs()).frames + + assert not torch.allclose(marked, unmarked) + + def test_last_frame_condition_stays_at_the_end_across_a_temporal_round(self): + # 17 frames -> 2 segments, so the temporal round runs 2 tiles. `index=-1` must land on the last original + # latent (scaled into the 2x canvas), not wrap to the first tile's last frame or to frame 0. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + frame = np.full((32, 32, 3), 200, dtype=np.uint8) + + inputs = self.get_dummy_inputs() + inputs["num_frames"] = 17 + inputs["temporal_upscalings"] = 1 + last = dict(inputs) + last["conditions"] = LTX2VideoCondition(frames=frame, index=-1, strength=1.0, crf=0) + first = dict(inputs) + first["conditions"] = LTX2VideoCondition(frames=frame, index=0, strength=1.0, crf=0) + + last_out = pipe(**last).frames + first_out = pipe(**first).frames + + assert last_out.shape[1] == (17 - 1) * 2 + 1 + assert not torch.allclose(last_out, first_out) + + def test_stages_1_and_2_condition_at_the_snapped_fps(self): + # RoPE time is `pixel_frame / fps`, and the transformer never saw 48. Every stage must lay its tokens out at + # the snapped rate; only the returned frame count and the audio trim use the playback rate. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + base_coords = [] + original = pipe.prepare_latents + + def capture(**kwargs): + result = original(**kwargs) + base_coords.append((result[3], kwargs["num_frames"], kwargs["height"], kwargs["width"])) + return result + + inputs = self.get_dummy_inputs() + inputs["frame_rate"] = 48.0 + pipe.prepare_latents = capture + try: + pipe(**inputs) + finally: + del pipe.prepare_latents + + assert len(base_coords) == 2, "expected stage 1 and stage 2" + for coords, num_frames, height, width in base_coords: + # Compare against what the real RoPE produces at 60, rather than trusting the argument that was plumbed in. + expected = pipe.transformer.rope.prepare_video_coords( + batch_size=1, + num_frames=(num_frames - 1) // pipe.vae_temporal_compression_ratio + 1, + height=height // pipe.vae_spatial_compression_ratio, + width=width // pipe.vae_spatial_compression_ratio, + device=coords.device, + fps=MAX_CONDITIONING_FPS, + ) + assert torch.allclose(coords[:1, :, : expected.shape[2]], expected) + + def test_a_carried_slot_is_the_copy_the_stitched_canvas_kept(self): + # Two tiles invent the slot that falls in the later one's dropped lead-in. The stitch keeps the earlier tile's + # frames there, so the earlier tile's slot is the one the next round must anchor on; the later tile's describes + # frames the canvas does not hold. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + + slot_slices, anchors, denoised = [], [], [] + original_prepare, original_denoise = pipe.prepare_latents, pipe.denoise + + def capture_prepare(**kwargs): + result = original_prepare(**kwargs) + slot_slices.append(result[-1]) + anchors.append(kwargs.get("keyframe_latents")) + return result + + def capture_denoise(**kwargs): + result = original_denoise(**kwargs) + denoised.append(result[0]) + return result + + inputs = self.get_dummy_inputs() + # 17 frames -> segment 8, so round 1 runs 2 tiles whose slot at pixel 8 falls in tile 1's dropped lead-in, + # and round 2 tile 0 anchors on it at pixel 16. + inputs["num_frames"] = 17 + inputs["temporal_upscalings"] = 2 + pipe.prepare_latents, pipe.denoise = capture_prepare, capture_denoise + try: + pipe(**inputs) + finally: + del pipe.prepare_latents, pipe.denoise + + # Passes are stage 1, stage 2, then round 1's 2 tiles, then round 2's 4 tiles. + assert len(denoised) == 2 + 2 + 4 + first_tile_slot = denoised[2][:, slot_slices[2]] + second_tile_slot = denoised[3][:, slot_slices[3]][:, : first_tile_slot.shape[1]] + assert not torch.allclose(first_tile_slot, second_tile_slot), "the two tiles must invent different content" + + round_2_anchors = anchors[4] + assert [position for position, _, _ in round_2_anchors] == [16] + _, anchor_latent, strength = round_2_anchors[0] + assert strength == ANCHOR_KEYFRAME_STRENGTH + assert torch.equal(pipe._pack_latents(anchor_latent), first_tile_slot) + + +class TestLTX2DFRPipelineMemory(LTX2DFRPipelineTesterConfig, MemoryTesterMixin): + @pytest.mark.skip( + "Pre-existing for the whole LTX-2 family, not DFR-specific: the shared harness group-offloads only " + "`text_encoder` / `transformer` and moves `vae`, leaving the LTX-2-specific `connectors` on the CPU while it " + "receives accelerator tensors from the offloaded text encoder. Verified to fail identically on the stock " + "`LTX2Pipeline`. `test_pipeline_level_group_offloading_inference`, which offloads every component, passes." + ) + def test_group_offloading_inference(self): + pass