diff --git a/docs/source/en/api/pipelines/ltx2.md b/docs/source/en/api/pipelines/ltx2.md index f45a21014a4a..5a04574df465 100644 --- a/docs/source/en/api/pipelines/ltx2.md +++ b/docs/source/en/api/pipelines/ltx2.md @@ -1030,6 +1030,237 @@ 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. Each `__call__` is **one denoise pass** at `height` × `width`: it generates video plus extra single-pixel-frame **keyframe slots**, or re-denoises supplied latents seeded from those slots. Callers compose stages the same way as other LTX two-stage pipelines — this pipeline, [`LTX2LatentUpsamplePipeline`], this pipeline again, then [`LTX2DFRTemporalRefinePipeline`] for each temporal round. + +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; `output_type="latent"` returns that padded grid so a slot on the pad is not dropped. Trim with `trim_canvas` before VAE decode. + +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 it 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. + +`return_dict=False` stays `(frames, audio)` so the diffusion-decoder path does not break. Composition uses `return_dict=True` for `keyframes` and `keyframe_positions`. + +The full recipe below is the one worth starting from: 1088×1920 image-to-video, one x2 temporal refine round, and the x2 spatial detailing IC-LoRA on stage 2. + +```py +import torch +from diffusers import ( + LTX2DFRPipeline, + LTX2DFRTemporalRefinePipeline, + LTX2LatentUpsamplePipeline, + LTXEulerAncestralRFScheduler, +) +from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2.dfr_core import trim_canvas +from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition +from diffusers.pipelines.ltx2.utils import STAGE_2_DISTILLED_SIGMA_VALUES +from diffusers.utils import encode_video, load_image + +pipe = LTX2DFRPipeline.from_pretrained("Lightricks/LTX-2.5-Diffusers", torch_dtype=torch.bfloat16) +latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + "Lightricks/LTX-2.5-Diffusers", subfolder="latent_upsampler", torch_dtype=torch.bfloat16 +) +# The x2 temporal upsampler is not in the published `model_index.json` — convert it with +# `--temporal_latent_upsampler`. +temporal_latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + "path/to/converted/temporal_latent_upsampler", torch_dtype=torch.bfloat16 +) +upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=latent_upsampler) +temporal_pipe = LTX2DFRTemporalRefinePipeline( + scheduler=LTXEulerAncestralRFScheduler(eta=0.5), + vae=pipe.vae, + audio_vae=pipe.audio_vae, + text_encoder=pipe.text_encoder, + tokenizer=pipe.tokenizer, + connectors=pipe.connectors, + transformer=pipe.transformer, + vocoder=pipe.vocoder, + temporal_latent_upsampler=temporal_latent_upsampler, +) + +# All three pipelines share the same components, so place them together. Do not call +# `enable_model_cpu_offload()` on one of them: its hooks would own modules the other two also call, +# and `temporal_latent_upsampler` — held only by `temporal_pipe` — would never reach the device. +pipe.to("cuda") +upsample_pipe.to("cuda") +temporal_pipe.to("cuda") + +image = load_image("https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/cat.png") +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, width = 1088, 1920 +frame_rate = 24.0 +generator = torch.Generator(device="cuda").manual_seed(0) + +num_frames = 121 +out = pipe( + prompt=prompt, + conditions=conditions, + height=height // 2, + width=width // 2, + num_frames=num_frames, + frame_rate=frame_rate, + generator=generator, + output_type="latent", +) +up_video = upsample_pipe(latents=out.frames, output_type="latent", return_dict=False)[0] +up_keyframes = upsample_pipe(latents=out.keyframes, output_type="latent", return_dict=False)[0] + +# Load after stage 1 so the adapter is never disabled. `set_adapters` does not re-enable a +# transformer that already had `disable_adapters()` called on it. +pipe.load_lora_weights("Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", adapter_name="detailing") +pipe.set_adapters(["detailing"], adapter_weights=[0.5]) +out2 = pipe( + prompt=prompt, + conditions=conditions, + latents=up_video, + audio_latents=out.audio, + keyframes_latents=up_keyframes, + keyframe_positions=out.keyframe_positions, + reference_latents=out.frames, + height=height, + width=width, + num_frames=num_frames, + frame_rate=frame_rate, + noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], + sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, + generator=generator, + output_type="latent", +) +pipe.transformer.disable_adapters() + +# `num_frames` here is the *padded* canvas the latents actually cover, which `resolve_canvas` may have +# grown past the 121 that were asked for. `condition_num_frames` stays the original request so a +# negative `condition.index` does not wrap onto the pad. +ratio = pipe.vae.temporal_compression_ratio +canvas_frames = (out2.frames.shape[2] - 1) * ratio + 1 +out3 = temporal_pipe( + latents=out2.frames, + keyframes_latents=out2.keyframes, + keyframe_positions=out2.keyframe_positions, + audio_latents=out.audio, + prompt=prompt, + conditions=conditions, + height=height, + width=width, + num_frames=canvas_frames, + frame_rate=frame_rate, + source_seconds=canvas_frames / frame_rate, + condition_num_frames=num_frames, + generator=generator, + output_type="latent", +) + +# Keep the padded canvas until decode so a slot on the pad is not dropped. `trim_canvas` counts *pixel* +# frames, and the round mapped `N -> 2 (N - 1) + 1`. +playback_fps = frame_rate * 2 +requested_frames = (num_frames - 1) * 2 + 1 +video_latents = trim_canvas(out3.frames, requested_frames, ratio) +timestep = None +if pipe.vae.config.timestep_conditioning: + timestep = torch.zeros(video_latents.shape[0], device=video_latents.device, dtype=pipe.vae.dtype) +video = pipe.vae.decode(video_latents.to(pipe.vae.dtype), timestep, return_dict=False)[0] +video = pipe.video_processor.postprocess_video(video, output_type="np") + +# Audio is stage 1's. Cut it to the video's duration so a muxed container does not outlast the picture. +audio = pipe.vocoder(pipe.audio_vae.decode(out.audio.to(pipe.audio_vae.dtype), return_dict=False)[0]) +audio_samples = round(requested_frames / playback_fps * pipe.vocoder.config.output_sampling_rate) +audio = audio[..., : min(audio.shape[-1], audio_samples)] + +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 **this pass**, not the final output. Stage 1 runs at half the 1080p canvas (544×960); stage 2 at 1088×1920. Each must be divisible by the VAE's spatial compression ratio (32 on LTX-2.5 for a single pass; 64 when stage 1 is half of 1080p). This is why 1080p is **1920×1088** and 4K is **3840×2176**. Each pass runs a fixed distilled schedule (`sigmas`), so there is no `num_inference_steps`; the distilled schedules are trained without guidance, so there is no `negative_prompt` or `guidance_scale` either. The shipped audio is stage 1's — later passes still run an audio stream 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 **after stage 1** and activate it for stage 2 only (`set_adapters(["detailing"], adapter_weights=[0.5])`). Stage 2 then attends to the stage-1 half-resolution latent as `reference_latents`. Stage 1 and the temporal rounds run with the adapter off. If you load the LoRA before stage 1, `transformer.disable_adapters()` turns it off, and `set_adapters` does **not** turn it back on — call `transformer.enable_adapters()` before stage 2, or load the weights after stage 1 as in the recipe. `reference_downscale_factor` (default `2`) scales the reference tokens' spatial coordinates into the target's coordinate space. + +**Temporal refinement.** [`LTX2DFRTemporalRefinePipeline`] is one round: temporally upsample, tile on keyframe seams, ancestral-denoise with [`LTXEulerAncestralRFScheduler`] (`eta=0.5`), stitch by dropping the later tile's lead-in, and merge the carry-keyframe bag. Construct that scheduler yourself — the round is refused with anything else, since a deterministic step would run to completion and only return a softer canvas. Call the pipeline once per round; loop for 2x / 4x, passing `round_index`. After a round, `keyframe_positions` cannot be re-derived from the original `num_frames` and must be passed through. Each tile is handed the slice of the frozen stage-1 audio covering its own playback window. `source_seconds` is the *stage-1* duration and stays fixed across rounds, so later rounds must pass it explicitly rather than take the default. + +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.** Compose it; there is no fourth pipeline. Spatially upsample the **video only**, rebuild carry keyframes in RGB (`decode` → Lanczos ×2 → `encode` via [`~LTX2DFRPipeline.rebuild_epilogue_keyframes`]; never latent-upsample epilogue keyframes), then: + +```py +from diffusers.pipelines.ltx2.dfr_layout import epilogue_tiles, pixel_to_latent_index + +# One more doubling on top of the recipe above, so every stage below the output halves again: +# `epilogue_height` must be divisible by 128 (`4 * 32`), which is why 4K is 3840x2176. +epilogue_height, epilogue_width = height * 2, width * 2 +refined_frames = (out3.frames.shape[2] - 1) * ratio + 1 + +up_video = upsample_pipe(latents=out3.frames, output_type="latent", return_dict=False)[0] +epilogue_keyframes = pipe.rebuild_epilogue_keyframes( + out3.keyframes, + decode_timestep=0.0, + decode_noise_scale=0.0, + seed=0, + device=up_video.device, + dtype=torch.float32, +) + +# Temporal cuts land on the seams the *last* round stitched on -- the positions handed into it, doubled -- +# not on every carry keyframe, since the slots that round invented sit mid-window. +tiles = epilogue_tiles( + latent_shape=( + (refined_frames - 1) // ratio + 1, + epilogue_height // pipe.vae.spatial_compression_ratio, + epilogue_width // pipe.vae.spatial_compression_ratio, + ), + frame_tiles=2, # 2 ** number of temporal rounds + frame_seams=[pixel_to_latent_index(2 * p, ratio) for p in out2.keyframe_positions], +) + +pipe.transformer.enable_adapters() # the epilogue is a detailing pass too +out4 = pipe( + prompt=prompt, + conditions=conditions, + latents=up_video, + audio_latents=out.audio, + generate_slots=False, + guidance_keyframe_latents=epilogue_keyframes, + guidance_keyframe_positions=out3.keyframe_positions, + reference_latents=out3.frames, + height=epilogue_height, + width=epilogue_width, + num_frames=refined_frames, + frame_rate=playback_fps, + noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], + sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, + freeze_audio=True, + video_tiles=tiles, + generator=generator, + output_type="latent", +) +pipe.transformer.disable_adapters() +``` + +The two axes are seamed differently. Neither side of a spatial border holds a known answer, so those overlaps are blended with trapezoidal weights. Temporal tiles are cut on the last refine round's keyframe seams. + +**Decoding with the diffusion decoder.** For maximum detail fidelity, stay on `output_type="latent"` and hand the (already denormalized, possibly `trim_canvas`'d) latents to [`LTX2VideoDiffusionDecodePipeline`]. + +```py +from diffusers import LTX2VideoDiffusionDecodePipeline +from diffusers.models.autoencoders.ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel + +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=out3.frames, denormalize=False, output_type="np", return_dict=False)[0] +``` + ## LTX2Pipeline [[autodoc]] LTX2Pipeline @@ -1048,6 +1279,22 @@ You can see the supported workflows in the docs for each blockset (e.g. [`LTX2Au - all - __call__ +## LTX2DFRPipeline + +[[autodoc]] LTX2DFRPipeline + - all + - __call__ + +## LTX2DFRTemporalRefinePipeline + +[[autodoc]] LTX2DFRTemporalRefinePipeline + - all + - __call__ + +## LTX2DFRPipelineOutput + +[[autodoc]] pipelines.ltx2.pipeline_output.LTX2DFRPipelineOutput + ## LTX2LatentUpsamplePipeline [[autodoc]] LTX2LatentUpsamplePipeline 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 747bec2bdf84..a1f3ae52d20c 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -725,6 +725,9 @@ "LongCatImageEditPipeline", "LongCatImagePipeline", "LTX2ConditionPipeline", + "LTX2DFRPipeline", + "LTX2DFRPipelineOutput", + "LTX2DFRTemporalRefinePipeline", "LTX2HDRPipeline", "LTX2ImageToVideoPipeline", "LTX2InContextPipeline", @@ -1572,6 +1575,9 @@ LongCatImageEditPipeline, LongCatImagePipeline, LTX2ConditionPipeline, + LTX2DFRPipeline, + LTX2DFRPipelineOutput, + LTX2DFRTemporalRefinePipeline, 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..fed3e449a7e1 100644 --- a/src/diffusers/pipelines/__init__.py +++ b/src/diffusers/pipelines/__init__.py @@ -340,6 +340,9 @@ _import_structure["ltx2"] = [ "LTX2Pipeline", "LTX2ConditionPipeline", + "LTX2DFRPipeline", + "LTX2DFRPipelineOutput", + "LTX2DFRTemporalRefinePipeline", "LTX2HDRPipeline", "LTX2InContextPipeline", "LTX2ImageToVideoPipeline", @@ -802,6 +805,9 @@ ) from .ltx2 import ( LTX2ConditionPipeline, + LTX2DFRPipeline, + LTX2DFRPipelineOutput, + LTX2DFRTemporalRefinePipeline, LTX2HDRPipeline, LTX2ImageToVideoPipeline, LTX2InContextPipeline, diff --git a/src/diffusers/pipelines/ltx2/__init__.py b/src/diffusers/pipelines/ltx2/__init__.py index fe5b563a3c39..d4aa35127403 100644 --- a/src/diffusers/pipelines/ltx2/__init__.py +++ b/src/diffusers/pipelines/ltx2/__init__.py @@ -28,6 +28,9 @@ _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_dfr_temporal_refine"] = ["LTX2DFRTemporalRefinePipeline"] + _import_structure["pipeline_output"] = ["LTX2DFRPipelineOutput", "LTX2PipelineOutput", "LTX2VideoDecodeOutput"] _import_structure["pipeline_ltx2_diffusion_decode"] = ["LTX2VideoDiffusionDecodePipeline"] _import_structure["pipeline_ltx2_hdr_lora"] = ["LTX2HDRPipeline", "LTX2HDRReferenceCondition"] _import_structure["pipeline_ltx2_ic_lora"] = ["LTX2InContextPipeline", "LTX2ReferenceCondition"] @@ -49,11 +52,14 @@ 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_dfr_temporal_refine import LTX2DFRTemporalRefinePipeline from .pipeline_ltx2_diffusion_decode import LTX2VideoDiffusionDecodePipeline from .pipeline_ltx2_hdr_lora import LTX2HDRPipeline, LTX2HDRReferenceCondition from .pipeline_ltx2_ic_lora import LTX2InContextPipeline, LTX2ReferenceCondition from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline + from .pipeline_output import LTX2DFRPipelineOutput, LTX2PipelineOutput, LTX2VideoDecodeOutput from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE else: diff --git a/src/diffusers/pipelines/ltx2/dfr_core.py b/src/diffusers/pipelines/ltx2/dfr_core.py new file mode 100644 index 000000000000..2fdb8226c818 --- /dev/null +++ b/src/diffusers/pipelines/ltx2/dfr_core.py @@ -0,0 +1,1499 @@ +# 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. + +"""Shared DFR denoise core used by [`LTX2DFRPipeline`] and [`LTX2DFRTemporalRefinePipeline`].""" + +import copy +import math +from typing import Any, Callable + +import numpy as np +import PIL.Image +import torch + +from ...image_processor import PipelineImageInput +from ...schedulers import LTXEulerAncestralRFScheduler +from ...utils import is_torch_xla_available, logging +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from .latent_upsampler import LTX2LatentUpsamplerModel +from .pipeline_ltx2_condition import LTX2VideoCondition +from .pipeline_output import LTX2DFRPipelineOutput +from .prompt_enhancement import ( + _pad_inputs_for_attention_alignment, + _prepare_enhance_image, + clean_response, +) +from .utils import ( + GEMMA3_PROMPT_ENHANCEMENT_CONFIG, + GEMMA4_PROMPT_ENHANCEMENT_CONFIG, + apply_image_conditioning_crf, + resolve_default_image_crf, +) + + +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 + + +# 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 trim_canvas(latents: torch.Tensor, num_frames: int, temporal_compression_ratio: int) -> torch.Tensor: + """Drop the padded tail of a DFR canvas so it matches `num_frames` pixel frames. + + The canvas is padded to a whole number of keyframe segments; this keeps the tokens that cover + `num_frames` and is the last step before VAE decode. `output_type="latent"` returns the padded + grid so a slot on the pad (e.g. pixel 96 on an 81→97 canvas) is not dropped. + """ + keep = (num_frames - 1) // temporal_compression_ratio + 1 + return latents[:, :, :keep] + + +class LTX2DFRCoreMixin: + """Pack/unpack, latent prep, and the shared DFR denoise loop. + + Not a pipeline. [`LTX2DFRPipeline`] and [`LTX2DFRTemporalRefinePipeline`] both mix this in; they do + not subclass each other. + """ + + def _init_dfr_runtime(self) -> None: + 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 + + def _maybe_normalize_video_latents( + self, latents: torch.Tensor | None, latents_normalized: bool + ) -> torch.Tensor | None: + if latents is None or latents_normalized: + return latents + # Public latents are often bf16 from the upsampler; denoise and the old one-shot path pack in float32. + latents = latents.float() + return self._normalize_latents( + latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + + def _inference_timesteps(self, sigmas: list[float], device: torch.device) -> torch.Tensor: + """Set the scheduler schedule and return the timesteps the denoise loop should walk. + + [`FlowMatchEulerDiscreteScheduler`] takes distilled sigmas *without* a terminal 0 (it appends + one) and exposes `len(sigmas)` timesteps. [`LTXEulerAncestralRFScheduler`] takes distilled + sigmas *plus* terminal 0 and the loop must skip the last timestep, matching + [`LTXI2VLongMultiPromptPipeline`]. + """ + if isinstance(self.scheduler, LTXEulerAncestralRFScheduler): + sigmas = [float(sigma) for sigma in sigmas] + if not sigmas or sigmas[-1] != 0.0: + sigmas = [*sigmas, 0.0] + self.scheduler.set_timesteps(sigmas=sigmas, device=device) + return self.scheduler.timesteps[:-1] + self.scheduler.set_timesteps(sigmas=sigmas, device=device) + return self.scheduler.timesteps + + # 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, + latents_normalized: bool = True, + ) -> 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*): + `(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*): + `(batch_size, num_channels_latents, F, H, W)` initial content for the base tokens. Public pipeline + latents are raw (denormalized); pass `latents_normalized=False` at that boundary. Tile loops that + already sit in VAE-normalized space leave the default. + latents_normalized (`bool`, defaults to `True`): + Whether `latents`, `slot_initial_latents`, `reference_latents`, and `keyframe_latents` are already + VAE-normalized. Internal tile loops pass `True`; each pipeline `__call__` passes `False`. + 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 + + latents = self._maybe_normalize_video_latents(latents, latents_normalized) + slot_initial_latents = self._maybe_normalize_video_latents(slot_initial_latents, latents_normalized) + reference_latents = self._maybe_normalize_video_latents(reference_latents, latents_normalized) + if keyframe_latents: + keyframe_latents = [ + (position, self._maybe_normalize_video_latents(latent, latents_normalized), strength) + for position, latent, strength in keyframe_latents + ] + + 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 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, + video_tile_plan: list | None = None, + 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. + + After the x0 conditioning blend, the velocity is `(latents - denoised) / sigma` so an RF step + `x0 = x - σ v` recovers the blended `denoised`. Stage 1 / 2 / the epilogue keep + [`FlowMatchEulerDiscreteScheduler`] and take that Euler step as-is. Temporal refine swaps in + [`LTXEulerAncestralRFScheduler`] (`eta=0.5`); that step renoises every token, so the conditioning + blend is applied again afterwards or strength-0.95 seam anchors erode. + + Args: + sigmas (`list[float]`): + Noise schedule for this pass, without the terminal `0.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`, *optional*): + Per-tile token plan from [`video_tile_plan`][`~diffusers.pipelines.ltx2.dfr_layout.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. + generator (`torch.Generator`, *optional*): + Forwarded to [`LTXEulerAncestralRFScheduler.step`]. Temporal tiles pass a per-tile seed so ancestral + draws do not share a stream or consume the state the next tile's initial noising reads. Distilled Euler + does not read it. + 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] + + timesteps = self._inference_timesteps(sigmas, device) + # 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 + + ancestral = ( + isinstance(self.scheduler, LTXEulerAncestralRFScheduler) and float(self.scheduler.config.eta) > 0 + ) + if ancestral: + latents = self.scheduler.step( + (latents.float() - denoised) / sigma, t, latents, generator=generator, return_dict=False + )[0] + # Ancestral Euler noises every token, so re-apply the blend or strength-0.95 seam anchors erode. + latents = (latents.float() * (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 + + def _unpack_video_and_slots( + self, + packed: torch.Tensor, + num_frames: int, + height: int, + width: int, + slot_token_slice: slice | None, + num_slots: int, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + video = self._unpack_video_latents( + packed[:, : latent_num_frames * latent_height * latent_width], + latent_num_frames, + latent_height, + latent_width, + ) + keyframes = None + if slot_token_slice is not None and num_slots: + keyframes = self._unpack_video_latents(packed[:, slot_token_slice], num_slots, latent_height, latent_width) + return video, keyframes + + def _public_audio_from_packed(self, audio_latents: torch.Tensor, audio_num_frames: int) -> torch.Tensor: + latent_mel_bins = self.audio_mel_bins // self.audio_vae_mel_compression_ratio + audio_latents = self._denormalize_audio_latents( + audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std + ) + return self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins) + + def _pack_public_audio(self, audio_latents: torch.Tensor) -> torch.Tensor: + packed = self._pack_audio_latents(audio_latents) + return self._normalize_audio_latents(packed, self.audio_vae.latents_mean, self.audio_vae.latents_std) + + def _finalize_output( + self, + video_latents: torch.Tensor, + audio_latents: torch.Tensor, + keyframe_latents: torch.Tensor | None, + keyframe_positions: list[int] | None, + output_type: str, + return_dict: bool, + output_cls, + requested_frames: int | None, + playback_fps: float, + decode_timestep: float | list[float], + decode_noise_scale: float | list[float] | None, + generator: torch.Generator | list[torch.Generator] | None, + prompt_embeds: torch.Tensor, + ): + """Denormalize, optionally trim and decode, and wrap the DFR return value. + + `video_latents` / `keyframe_latents` are VAE-normalized 5D. `audio_latents` is unpacked denormalized 4D + (the public audio contract). Latent output is untrimmed so a slot on the padded canvas is not dropped; + decode trims to `requested_frames` when that is set. + """ + output_cls = output_cls or LTX2DFRPipelineOutput + batch_size = video_latents.shape[0] + device = video_latents.device + + if output_type != "latent" and requested_frames is not None: + video_latents = trim_canvas(video_latents, requested_frames, self.vae_temporal_compression_ratio) + num_frames = requested_frames + else: + num_frames = (video_latents.shape[2] - 1) * self.vae_temporal_compression_ratio + 1 + + video = self._denormalize_latents( + video_latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + keyframes = None + if keyframe_latents is not None: + keyframes = self._denormalize_latents( + keyframe_latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor + ) + + if output_type == "latent": + 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 / playback_fps * self.vocoder.config.output_sampling_rate) + ) + audio = audio[..., :audio_samples] + + self.maybe_free_model_hooks() + + if not return_dict: + return (video, audio) + + return output_cls(frames=video, audio=audio, keyframes=keyframes, keyframe_positions=keyframe_positions) + + @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 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..e5a3814767ff --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr.py @@ -0,0 +1,629 @@ +# 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. + +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 ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin +from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.transformers import LTX2VideoTransformer3DModel +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import replace_example_docstring +from ...utils.torch_utils import randn_tensor +from ..pipeline_utils import DiffusionPipeline +from .connectors import LTX2TextConnectors +from .dfr_core import ( + EPILOGUE_KEYFRAME_STRENGTH, + LTX2DFRCoreMixin, + _conditioning_fps, + retrieve_latents, +) +from .dfr_layout import LTX2DFREpilogueTile, resolve_canvas, video_tile_plan +from .duration_head import LTX2DurationHead +from .pipeline_ltx2_condition import LTX2VideoCondition +from .pipeline_output import LTX2DFRPipelineOutput +from .utils import ( + DISTILLED_SIGMA_VALUES, + LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT, + LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT, +) +from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTX2DFRPipeline + >>> from diffusers.utils import encode_video + + >>> pipe = LTX2DFRPipeline.from_pretrained( + ... "Lightricks/LTX-2.5-Diffusers", 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", + ... ) + ``` +""" + + +class LTX2DFRPipeline(LTX2DFRCoreMixin, DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): + r""" + Pipeline for one Diffusion Fidelity Rendering (DFR) denoise pass with LTX-2.5. + + A pass generates video *and* extra single-pixel-frame keyframe slots, or re-denoises supplied latents seeded from + those slots. Callers compose stages: this pipeline at half resolution, [`LTX2LatentUpsamplePipeline`] spatially, + this pipeline again at full resolution with the upsampled slots and an IC-LoRA reference, then + [`LTX2DFRTemporalRefinePipeline`] for each temporal refine round. See the LTX-2 docs for the 1080p recipe. + + Slot positions come from a segment grid aligned to the VAE's temporal border (`resolve_canvas`). The canvas is + padded to a whole number of segments; `output_type="latent"` returns that padded grid so a slot on the pad is not + dropped. Trim with [`trim_canvas`] before VAE decode. + + Requires a transformer whose config sets `use_keyframes_abs_pos_embedding` (LTX-2.5 and later) when + `generate_slots=True`. + + 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. + 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->vae->audio_vae->vocoder" + ) + _optional_components = ["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, + 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, + processor=processor, + prompt_enhancer=prompt_enhancer, + duration_head=duration_head, + ) + self._init_dfr_runtime() + + 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`): + Raw `(batch_size, C, K, H, W)` carry keyframes, as a previous pass returned them. + 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`: Raw `(batch_size, C, K, 2H, 2W)` latents, ready to pass straight back in as + `guidance_keyframe_latents`. + """ + if keyframe_latents.ndim != 5: + raise ValueError(f"Expected carry keyframes (B, C, K, H, W), got {tuple(keyframe_latents.shape)}") + + keyframe_latents = self._maybe_normalize_video_latents(keyframe_latents, False) + + 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(retrieve_latents(self.vae.encode(stretched), sample_mode="argmax")) + return torch.cat(encoded, dim=2).to(device=device, dtype=dtype) + + @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, + latents: torch.Tensor | None = None, + audio_latents: torch.Tensor | None = None, + keyframes_latents: torch.Tensor | None = None, + keyframe_positions: list[int] | None = None, + reference_latents: torch.Tensor | None = None, + reference_downscale_factor: int = 2, + guidance_keyframe_latents: torch.Tensor | None = None, + guidance_keyframe_positions: list[int] | None = None, + guidance_keyframe_strength: float = EPILOGUE_KEYFRAME_STRENGTH, + generate_slots: bool = True, + sigmas: list[float] = DISTILLED_SIGMA_VALUES, + noise_scale: float | None = None, + freeze_audio: bool = False, + video_tiles: list[LTX2DFREpilogueTile] | None = None, + 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. + + One denoise pass at `height` × `width`. Compose stages in the caller: this pipeline at half-res, spatial + upsample, this pipeline again with `latents` / `keyframes_latents` / `reference_latents`, then + [`LTX2DFRTemporalRefinePipeline`] for each temporal round. `return_dict=False` stays `(frames, audio)` so the + diffusion-decoder path does not break; composition uses `return_dict=True` for keyframes. + + 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 on this pass's `num_frames`. + height (`int`, *optional*, defaults to `704`): + The height in pixels of **this pass**, not the final composed output. + width (`int`, *optional*, defaults to `1216`): + The width in pixels of this pass. + num_frames (`int`, *optional*): + Pixel frame count of this pass, before internal canvas padding. If not supplied, the duration is + predicted from the prompt by the `duration_head`. Must satisfy `(num_frames - 1) % 8 == 0`. + frame_rate (`float`, *optional*, defaults to `24.0`): + Playback fps of this pass. RoPE time snaps to 60 whenever this is above 30. + 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. + latents (`torch.Tensor`, *optional*): + Raw `(batch_size, channels, frames, height, width)` video latents to re-denoise (stage 2 / epilogue). + audio_latents (`torch.Tensor`, *optional*): + Raw unpacked `(batch_size, channels, length, mel_bins)` audio latents. Stage 2 still runs a joint + audio pass; the shipped waveform of a composed recipe is stage 1's, which the caller keeps. + keyframes_latents (`torch.Tensor`, *optional*): + Raw `(batch_size, channels, num_slots, height, width)` slot initials, used when `generate_slots=True`. + keyframe_positions (`list[int]`, *optional*): + Pixel-frame indices of the generated slots. Defaults to `resolve_canvas(num_frames)`. + reference_latents (`torch.Tensor`, *optional*): + Raw IC-LoRA reference video (typically stage 1's frames). + reference_downscale_factor (`int`, *optional*, defaults to `2`): + Ratio between this pass and the reference resolution, used to scale reference token coordinates. + guidance_keyframe_latents (`torch.Tensor`, *optional*): + Raw pinned guidance keyframes `(batch_size, channels, K, height, width)` — the epilogue path. Not + generated slots. + guidance_keyframe_positions (`list[int]`, *optional*): + Pixel-frame indices for `guidance_keyframe_latents`. + guidance_keyframe_strength (`float`, *optional*, defaults to `1.0`): + Conditioning strength for pinned guidance keyframes. + generate_slots (`bool`, *optional*, defaults to `True`): + Append generated keyframe-slot tokens. The epilogue sets this to `False` and pins + `guidance_keyframe_latents` instead. + sigmas (`list[float]`, *optional*): + Noise schedule for this pass, without the terminal `0.0`. + noise_scale (`float`, *optional*): + Noise level unconditioned tokens start at. Defaults to `sigmas[0]`. + freeze_audio (`bool`, *optional*, defaults to `False`): + Hold audio at sigma 0 (epilogue / when following a frozen stage-1 waveform). + video_tiles (`list[LTX2DFREpilogueTile]`, *optional*): + Epilogue tiling from [`~diffusers.pipelines.ltx2.dfr_layout.epilogue_tiles`], so a resolution too + large for one forward pass can still step a single canvas: each step runs the transformer once per + tile and blends the predictions. Resolved into a token plan against this pass's own RoPE coordinates, + which is why the layout is passed rather than the plan — the coordinates only exist once + [`~LTX2DFRPipeline.prepare_latents`] has run. + 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"`. Latent output is the untrimmed canvas. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`LTX2DFRPipelineOutput`] or a plain `(frames, audio)` 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. + 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: + [`LTX2DFRPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`LTX2DFRPipelineOutput`] 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 + + 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, + latents=latents, + audio_latents=audio_latents, + stg_scale=0.0, + audio_stg_scale=0.0, + ) + if generate_slots and 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." + ) + if keyframes_latents is not None and keyframes_latents.ndim != 5: + raise ValueError( + f"`keyframes_latents` must be unpacked 5D `(batch, channels, slots, height, width)`, got " + f"{keyframes_latents.ndim} dims." + ) + if guidance_keyframe_latents is not None and guidance_keyframe_latents.ndim != 5: + raise ValueError( + f"`guidance_keyframe_latents` must be unpacked 5D, got {guidance_keyframe_latents.ndim} dims." + ) + if not generate_slots and keyframes_latents is not None: + raise ValueError("`keyframes_latents` seeds generated slots; set `generate_slots=True` or omit it.") + if guidance_keyframe_latents is not None and guidance_keyframe_positions is None: + raise ValueError("`guidance_keyframe_positions` is required when `guidance_keyframe_latents` is passed.") + + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + 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] + + device = self._execution_device + noise_scale = sigmas[0] if noise_scale is None else noise_scale + + 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_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, + ) + + requested_frames = num_frames + canvas_frames, _, resolved_positions = resolve_canvas(num_frames, self.vae_temporal_compression_ratio) + slot_frame_indices = None + if generate_slots: + slot_frame_indices = ( + list(keyframe_positions) if keyframe_positions is not None else list(resolved_positions) + ) + + pinned_keyframes = None + if guidance_keyframe_latents is not None: + pinned_keyframes = [ + (position, guidance_keyframe_latents[:, :, index : index + 1], guidance_keyframe_strength) + for index, position in enumerate(guidance_keyframe_positions) + ] + + num_channels_latents = self.transformer.config.in_channels + 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) + conditioning_fps = _conditioning_fps(frame_rate) + + self._num_timesteps = len(sigmas) + progress_bar = self.progress_bar(total=self._num_timesteps) + + packed, conditioning_mask, clean_latents, video_coords, keyframes_mask, slot_token_slice = ( + self.prepare_latents( + conditions=conditions, + keyframe_latents=pinned_keyframes, + slot_frame_indices=slot_frame_indices, + slot_initial_latents=keyframes_latents, + reference_latents=reference_latents, + reference_downscale_factor=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=noise_scale, + dtype=torch.float32, + device=device, + generator=generator, + latents=latents, + latents_normalized=False, + ) + ) + + resolved_tile_plan = None + if video_tiles is not None: + resolved_tile_plan = video_tile_plan( + video_tiles, + video_coords, + (canvas_frames - 1) // self.vae_temporal_compression_ratio + 1, + height // self.vae_spatial_compression_ratio, + width // self.vae_spatial_compression_ratio, + ) + + if audio_latents is None: + audio_packed = 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, + ) + else: + audio_packed = self.prepare_audio_latents( + batch_size=batch_size, + num_channels_latents=self.audio_latent_channels, + audio_latent_length=audio_latents.shape[2], + num_mel_bins=self.audio_mel_bins, + noise_scale=0.0 if freeze_audio else noise_scale, + dtype=torch.float32, + device=device, + generator=generator, + latents=audio_latents, + ) + + packed, audio_packed = self.denoise( + latents=packed, + 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=sigmas, + frame_rate=conditioning_fps, + audio_latents=audio_packed, + freeze_audio=freeze_audio, + video_tile_plan=resolved_tile_plan, + generator=generator, + use_cross_timestep=use_cross_timestep, + attention_kwargs=attention_kwargs, + progress_bar=progress_bar, + callback_on_step_end=callback_on_step_end, + callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, + ) + progress_bar.close() + + num_slots = len(slot_frame_indices) if slot_frame_indices else 0 + video_latents, slot_latents = self._unpack_video_and_slots( + packed, canvas_frames, height, width, slot_token_slice, num_slots + ) + public_audio = self._public_audio_from_packed(audio_packed, audio_packed.shape[1]) + + out_positions = slot_frame_indices + if slot_latents is None and pinned_keyframes is not None: + # Epilogue: no generated slots; surface the guidance the pass was given, still normalized until finalize. + slot_latents = self._maybe_normalize_video_latents(guidance_keyframe_latents, False) + out_positions = list(guidance_keyframe_positions) + + return self._finalize_output( + video_latents=video_latents, + audio_latents=public_audio, + keyframe_latents=slot_latents, + keyframe_positions=out_positions, + output_type=output_type, + return_dict=return_dict, + output_cls=LTX2DFRPipelineOutput, + requested_frames=requested_frames, + playback_fps=frame_rate, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + generator=generator, + prompt_embeds=prompt_embeds, + ) diff --git a/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr_temporal_refine.py b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr_temporal_refine.py new file mode 100644 index 000000000000..40c3ca375c0c --- /dev/null +++ b/src/diffusers/pipelines/ltx2/pipeline_ltx2_dfr_temporal_refine.py @@ -0,0 +1,515 @@ +# 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. + +from typing import Any, Callable + +import torch +from transformers import ( + Gemma3ForConditionalGeneration, + Gemma4UnifiedForConditionalGeneration, + GemmaTokenizer, + GemmaTokenizerFast, +) + +from ...callbacks import MultiPipelineCallbacks, PipelineCallback +from ...loaders import FromSingleFileMixin, LTX2LoraLoaderMixin +from ...models.autoencoders import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video +from ...models.transformers import LTX2VideoTransformer3DModel +from ...schedulers import LTXEulerAncestralRFScheduler +from ...utils import replace_example_docstring +from ..pipeline_utils import DiffusionPipeline +from .connectors import LTX2TextConnectors +from .dfr_core import ( + ANCHOR_KEYFRAME_STRENGTH, + TEMPORAL_ANCESTRAL_ETA, + LTX2DFRCoreMixin, + _audio_window_for_tile, + _conditioning_fps, +) +from .dfr_layout import temporal_tile_plan +from .latent_upsampler import LTX2LatentUpsamplerModel +from .pipeline_ltx2_condition import LTX2VideoCondition +from .pipeline_output import LTX2DFRPipelineOutput +from .utils import TEMPORAL_ROUND_DISTILLED_SIGMA_VALUES +from .vocoder import LTX2Vocoder, LTX2VocoderWithBWE + + +EXAMPLE_DOC_STRING = """ + Examples: + ```py + >>> import torch + >>> from diffusers import LTX2DFRPipeline, LTX2DFRTemporalRefinePipeline, LTXEulerAncestralRFScheduler + >>> from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel + + >>> pipe = LTX2DFRPipeline.from_pretrained("Lightricks/LTX-2.5-Diffusers", torch_dtype=torch.bfloat16) + >>> temporal_upsampler = LTX2LatentUpsamplerModel.from_pretrained( + ... "path/to/converted/temporal_latent_upsampler", torch_dtype=torch.bfloat16 + ... ) + >>> temporal_pipe = LTX2DFRTemporalRefinePipeline( + ... scheduler=LTXEulerAncestralRFScheduler(eta=0.5), + ... vae=pipe.vae, + ... audio_vae=pipe.audio_vae, + ... text_encoder=pipe.text_encoder, + ... tokenizer=pipe.tokenizer, + ... connectors=pipe.connectors, + ... transformer=pipe.transformer, + ... vocoder=pipe.vocoder, + ... temporal_latent_upsampler=temporal_upsampler, + ... ) + >>> # `out` is a prior DFR pass at the same spatial size, `return_dict=True`. + >>> out = temporal_pipe( + ... latents=out.frames, + ... keyframes_latents=out.keyframes, + ... keyframe_positions=out.keyframe_positions, + ... audio_latents=out.audio, + ... prompt="A tabby cat stretching in a sunlit window", + ... height=1088, + ... width=1920, + ... num_frames=121, + ... output_type="latent", + ... ) + ``` +""" + + +class LTX2DFRTemporalRefinePipeline(LTX2DFRCoreMixin, DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): + r""" + One temporal DFR refine round: interpolate the canvas x2 in time, tile on keyframe seams, ancestral-denoise each + tile, stitch by dropping the later tile's lead-in, and merge the carry-keyframe bag. + + The scheduler is [`LTXEulerAncestralRFScheduler`] (`eta=0.5` by default). Stage 1 / 2 / the spatial epilogue stay + on [`FlowMatchEulerDiscreteScheduler`] via [`LTX2DFRPipeline`]. Call this pipeline once per round; the caller + loops for 2x / 4x. + + Incoming `latents` / `keyframes_latents` / `audio_latents` are raw (denormalized), matching [`LTX2Pipeline`]. + `keyframe_positions` must be the positions returned by the previous pass — they cannot be re-derived from the + original `num_frames` after a round has run. + + Args: + scheduler ([`LTXEulerAncestralRFScheduler`]): + Ancestral Euler scheduler in the rectified-flow parameterization. Construct with `eta=0.5` to match the + DFR temporal recipe. + 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. + temporal_latent_upsampler ([`LTX2LatentUpsamplerModel`]): + Temporal x2 latent upsampler applied at the start of the round. + """ + + model_cpu_offload_seq = "text_encoder->connectors->transformer->temporal_latent_upsampler->vae->audio_vae->vocoder" + _callback_tensor_inputs = ["latents", "prompt_embeds"] + + def __init__( + self, + scheduler: LTXEulerAncestralRFScheduler, + vae: AutoencoderKLLTX2Video, + audio_vae: AutoencoderKLLTX2Audio, + text_encoder: Gemma3ForConditionalGeneration | Gemma4UnifiedForConditionalGeneration, + tokenizer: GemmaTokenizer | GemmaTokenizerFast, + connectors: LTX2TextConnectors, + transformer: LTX2VideoTransformer3DModel, + vocoder: LTX2Vocoder | LTX2VocoderWithBWE, + temporal_latent_upsampler: LTX2LatentUpsamplerModel, + ): + 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, + temporal_latent_upsampler=temporal_latent_upsampler, + ) + self._init_dfr_runtime() + + @torch.no_grad() + @replace_example_docstring(EXAMPLE_DOC_STRING) + def __call__( + self, + prompt: str | list[str] = None, + latents: torch.Tensor = None, + keyframes_latents: torch.Tensor = None, + keyframe_positions: list[int] = None, + audio_latents: torch.Tensor = None, + conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, + height: int = 704, + width: int = 1216, + num_frames: int = 121, + frame_rate: float = 24.0, + source_seconds: float | None = None, + condition_num_frames: int | None = None, + round_index: int = 1, + sigmas: list[float] = TEMPORAL_ROUND_DISTILLED_SIGMA_VALUES, + noise_scale: float | None = None, + 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, + 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""" + Run one temporal refine round. + + 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`. + latents (`torch.Tensor`): + Raw video latents of the **input** canvas, `(batch_size, channels, frames, height, width)`. + keyframes_latents (`torch.Tensor`): + Raw carry keyframes `(batch_size, channels, K, height, width)` from the previous pass. + keyframe_positions (`list[int]`): + Pixel-frame indices of `keyframes_latents` on the **input** canvas. + audio_latents (`torch.Tensor`): + Frozen stage-1 audio, unpacked and denormalized. Each tile is handed the slice covering its playback + window; the returned audio is this waveform, not a per-tile re-denoise. + conditions (`LTX2VideoCondition` or `List[LTX2VideoCondition]`, *optional*): + Frame-level conditions indexed on `condition_num_frames` (the original request). Each round scales a + condition's pixel position by `2 ** round_index`. + height (`int`, *optional*, defaults to `704`): + Pixel height of this pass (same as the incoming video). + width (`int`, *optional*, defaults to `1216`): + Pixel width of this pass. + num_frames (`int`, *optional*, defaults to `121`): + Pixel frame count of the **input** canvas (untrimmed). + frame_rate (`float`, *optional*, defaults to `24.0`): + Playback fps of the **input**. The round doubles it. + source_seconds (`float`, *optional*): + Duration of the frozen stage-1 audio. Defaults to `num_frames / frame_rate`, which is correct for the + first round; later rounds must pass the original stage-1 duration so tiles do not drift. + condition_num_frames (`int`, *optional*): + Original generation `num_frames` used to encode `conditions`. Defaults to `num_frames`. After padding + or a prior round, pass the original request so `index=-1` does not wrap to the padded tail. + round_index (`int`, *optional*, defaults to `1`): + 1-based round number. Tiles seed ancestral noise as `seed + 1000 * round_index + tile`, and conditions + are scaled by `2 ** round_index`. + sigmas (`list[float]`, *optional*): + Distilled schedule for this round's tiles, without the terminal `0.0`. + noise_scale (`float`, *optional*): + Noise level unconditioned tokens start at. Defaults to `sigmas[0]`. + generator (`torch.Generator` or `list[torch.Generator]`, *optional*): + Random generator(s) for reproducibility. Ancestral draws use a separate per-tile seed derived from + this generator's `initial_seed()`. + 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+. + output_type (`str`, *optional*, defaults to `"pil"`): + Output format. Choose `"pil"`, `"np"`, `"pt"` or `"latent"`. Latent output is the untrimmed canvas. + return_dict (`bool`, *optional*, defaults to `True`): + Whether to return a [`LTX2DFRPipelineOutput`] or a plain `(frames, audio)` 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 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: + [`LTX2DFRPipelineOutput`] or `tuple`: + If `return_dict` is `True`, [`LTX2DFRPipelineOutput`] 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 + + if latents is None or keyframes_latents is None or keyframe_positions is None or audio_latents is None: + raise ValueError( + "`latents`, `keyframes_latents`, `keyframe_positions`, and `audio_latents` are required. They come " + "from the previous DFR pass (`return_dict=True`)." + ) + if getattr(self, "temporal_latent_upsampler", None) is None: + raise ValueError("LTX2DFRTemporalRefinePipeline requires the `temporal_latent_upsampler` component.") + if not self.transformer.config.use_keyframes_abs_pos_embedding: + raise ValueError( + "Temporal DFR invents mid-segment keyframe slots, which requires a transformer whose config sets " + "`use_keyframes_abs_pos_embedding` (LTX-2.5 and later)." + ) + if round_index < 1: + raise ValueError(f"`round_index` must be >= 1, got {round_index}") + # A round densifies freshly interpolated frames, which needs the ancestral renoise. `denoise` takes a plain + # Euler step for any other scheduler, so an unchecked component here would run to completion and just return + # a softer canvas -- the one failure mode with no signal at all. + if not isinstance(self.scheduler, LTXEulerAncestralRFScheduler): + raise ValueError( + f"Temporal refine needs `LTXEulerAncestralRFScheduler(eta={TEMPORAL_ANCESTRAL_ETA})`, got " + f"{type(self.scheduler).__name__}. `from_pretrained` picks up the repo's flow-match scheduler, which " + f"steps deterministically and silently loses the refine round's detail; construct the scheduler " + f"explicitly." + ) + if float(self.scheduler.config.eta) <= 0: + raise ValueError( + f"Temporal refine needs a stochastic step, but the scheduler has `eta={self.scheduler.config.eta}`. " + f"Use `LTXEulerAncestralRFScheduler(eta={TEMPORAL_ANCESTRAL_ETA})`." + ) + + 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, + num_frames=num_frames, + latents=latents, + audio_latents=audio_latents, + stg_scale=0.0, + audio_stg_scale=0.0, + ) + + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + + 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] + + if conditions is not None and not isinstance(conditions, list): + conditions = [conditions] + + device = self._execution_device + noise_scale = sigmas[0] if noise_scale is None else noise_scale + source_seconds = num_frames / frame_rate if source_seconds is None else source_seconds + condition_num_frames = num_frames if condition_num_frames is None else condition_num_frames + 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 + + prompt_embeds, prompt_attention_mask = self.encode_prompt( + prompt=prompt, + num_videos_per_prompt=1, + 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 + ) + + num_channels_latents = self.transformer.config.in_channels + temporal_ratio = self.vae_temporal_compression_ratio + audio_latents_per_second = ( + self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) + ) + + video_latents = self._maybe_normalize_video_latents(latents.to(device=device), False) + carry_keyframes = self._maybe_normalize_video_latents(keyframes_latents.to(device=device), False) + stage_1_audio = self._pack_public_audio(audio_latents.to(device=device, dtype=torch.float32)) + + video_latents = self.upsample_latents(video_latents, self.temporal_latent_upsampler) + canvas_frames = 2 * (num_frames - 1) + 1 + playback_fps = 2 * frame_rate + conditioning_fps = _conditioning_fps(playback_fps) + seam_positions = [2 * position for position in keyframe_positions] + tiles = temporal_tile_plan(seam_positions, canvas_frames, 2**round_index, temporal_ratio) + pixel_scale = 2**round_index + + self._num_timesteps = len(sigmas) * len(tiles) + progress_bar = self.progress_bar(total=self._num_timesteps) + step_offset = 0 + + round_conditions = self.encode_conditions( + conditions, height, width, condition_num_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)} + latent_height = height // self.vae_spatial_compression_ratio + latent_width = width // self.vae_spatial_compression_ratio + tokens_per_latent_frame = latent_height * latent_width + + 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] + + 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 + ] + + tile_keyframe_latents = [ + ( + position - tile.pixel_start, + carry_keyframes[:, :, seam_to_index[position] : seam_to_index[position] + 1], + ANCHOR_KEYFRAME_STRENGTH, + ) + for position in tile.anchors + ] + + 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) + if seed_indices + else None + ) + + ( + 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 or None, + slot_initial_latents=tile_slot_initials, + batch_size=batch_size, + num_channels_latents=num_channels_latents, + height=height, + width=width, + num_frames=tile_frames, + frame_rate=conditioning_fps, + noise_scale=noise_scale, + dtype=torch.float32, + device=device, + generator=generator, + latents=tile_video_latents, + ) + ancestral_generator = torch.Generator(device=device).manual_seed( + ancestral_seed_base + 1000 * round_index + tile_index + ) + tile_audio_latents = _audio_window_for_tile( + stage_1_audio, + pixel_start=tile.pixel_start, + tile_frames=tile_frames, + playback_fps=playback_fps, + source_seconds=source_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=sigmas, + frame_rate=conditioning_fps, + audio_latents=tile_audio_latents, + freeze_audio=True, + 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(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, + ) + ) + + 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}" + ) + + carry: dict[int, torch.Tensor] = { + position: carry_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) + + progress_bar.close() + + public_audio = self._public_audio_from_packed(stage_1_audio, stage_1_audio.shape[1]) + return self._finalize_output( + video_latents=video_latents, + audio_latents=public_audio, + keyframe_latents=carry_keyframes, + keyframe_positions=carry_positions, + output_type=output_type, + return_dict=return_dict, + output_cls=LTX2DFRPipelineOutput, + requested_frames=None, + playback_fps=playback_fps, + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + generator=generator, + prompt_embeds=prompt_embeds, + ) diff --git a/src/diffusers/pipelines/ltx2/pipeline_output.py b/src/diffusers/pipelines/ltx2/pipeline_output.py index aec3096f03b5..bbf2ce816f7e 100644 --- a/src/diffusers/pipelines/ltx2/pipeline_output.py +++ b/src/diffusers/pipelines/ltx2/pipeline_output.py @@ -23,6 +23,29 @@ class LTX2PipelineOutput(BaseOutput): audio: torch.Tensor +@dataclass +class LTX2DFRPipelineOutput(LTX2PipelineOutput): + r""" + Output class for DFR pipelines. + + Args: + frames (`torch.Tensor`, `np.ndarray`, or list[list[PIL.Image.Image]]): + Denoised video. Latent output is the untrimmed canvas, shape + `(batch_size, num_channels, latent_frames, latent_height, latent_width)`. + audio (`torch.Tensor`, `np.ndarray`): + Accompanying audio latents or waveform. + keyframes (`torch.Tensor`, *optional*): + Generated or carried keyframe latents of shape `(batch_size, num_channels, num_keyframes, latent_height, + latent_width)`. `None` when the pass did not produce slots (e.g. a tiled epilogue). + keyframe_positions (`list[int]`, *optional*): + Pixel-frame index of each keyframe on this pass's canvas. After a temporal round these cannot be + re-derived from the original `num_frames` and must be passed into the next stage. + """ + + keyframes: torch.Tensor | None = None + keyframe_positions: list[int] | None = None + + @dataclass class LTX2VideoDecodeOutput(BaseOutput): r""" 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..73beb676b869 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -3122,6 +3122,51 @@ 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 LTX2DFRPipelineOutput(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 LTX2DFRTemporalRefinePipeline(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/dfr_dummies.py b/tests/pipelines/ltx2/dfr_dummies.py new file mode 100644 index 000000000000..bd395ededfec --- /dev/null +++ b/tests/pipelines/ltx2/dfr_dummies.py @@ -0,0 +1,196 @@ +# 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 torch +from transformers import AutoTokenizer, Gemma3ForConditionalGeneration + +from diffusers import ( + AutoencoderKLLTX2Audio, + AutoencoderKLLTX2Video, + FlowMatchEulerDiscreteScheduler, + LTX2VideoTransformer3DModel, + LTXEulerAncestralRFScheduler, +) +from diffusers.pipelines.ltx2 import LTX2LatentUpsamplerModel, LTX2TextConnectors +from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder + + +BASE_TEXT_ENCODER_CKPT_ID = "hf-internal-testing/tiny-gemma3" + + +def get_dfr_dummy_components(*, spatial_upsampler: bool = False, temporal_upsampler: bool = False): + tokenizer = AutoTokenizer.from_pretrained(BASE_TEXT_ENCODER_CKPT_ID) + text_encoder = Gemma3ForConditionalGeneration.from_pretrained(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, + ) + 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, + ) + + components = { + "transformer": transformer, + "vae": vae, + "audio_vae": audio_vae, + "scheduler": FlowMatchEulerDiscreteScheduler(), + "text_encoder": text_encoder, + "tokenizer": tokenizer, + "connectors": connectors, + "vocoder": vocoder, + "processor": None, + "prompt_enhancer": None, + "duration_head": None, + } + if spatial_upsampler: + torch.manual_seed(0) + components["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, + ) + if temporal_upsampler: + torch.manual_seed(0) + components["temporal_latent_upsampler"] = LTX2LatentUpsamplerModel( + in_channels=4, + mid_channels=32, + num_blocks_per_stage=1, + dims=3, + spatial_upsample=False, + temporal_upsample=True, + ) + return components + + +def get_temporal_dummy_components(): + components = get_dfr_dummy_components(temporal_upsampler=True) + components["scheduler"] = LTXEulerAncestralRFScheduler(eta=0.5) + # A refine round never enhances a prompt -- stage 1 already did, and re-enhancing would + # denoise the canvas under a different prompt than the one that generated it. + for name in ("duration_head", "processor", "prompt_enhancer"): + components.pop(name) + return components + + +def get_dfr_dummy_inputs(**overrides): + inputs = { + "prompt": "a robot dancing", + "height": 32, + "width": 32, + "num_frames": 9, + "frame_rate": 25.0, + "sigmas": [1.0, 0.5], + "use_cross_timestep": False, + "max_sequence_length": 16, + "output_type": "pt", + } + inputs.update(overrides) + return inputs 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..9d352a9e76ad --- /dev/null +++ b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr.py @@ -0,0 +1,526 @@ +# 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 pytest +import torch + +from diffusers import LTX2DFRPipeline, LTX2LatentUpsamplePipeline, LTXEulerAncestralRFScheduler +from diffusers.pipelines.ltx2.dfr_core import ( + EPILOGUE_KEYFRAME_STRENGTH, + MAX_CONDITIONING_FPS, + _audio_window_for_tile, + trim_canvas, +) +from diffusers.pipelines.ltx2.dfr_layout import LTX2DFREpilogueTile, epilogue_tiles, video_tile_plan + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BasePipelineTesterConfig, MemoryTesterMixin, PipelineTesterMixin +from .dfr_dummies import get_dfr_dummy_components, get_dfr_dummy_inputs + + +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"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (9, 3, 32, 32) + optional_input_params = BasePipelineTesterConfig.optional_input_params - { + "num_inference_steps", + "num_images_per_prompt", + "latents", + } + + def get_dummy_components(self): + return get_dfr_dummy_components() + + def get_dummy_inputs(self): + inputs = get_dfr_dummy_inputs() + inputs["generator"] = self.get_generator(0) + return inputs + + +class TestLTX2DFRPipeline(LTX2DFRPipelineTesterConfig, PipelineTesterMixin): + def test_inference_batch_single_identical(self, batch_size=3, expected_max_diff=1e-3): + 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() + inputs["num_frames"] = 11 + assert pipe(**inputs).frames.shape[1] == 11 + + def test_latent_output_keeps_the_padded_canvas(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + inputs = self.get_dummy_inputs() + inputs["num_frames"] = 11 + inputs["output_type"] = "latent" + output = pipe(**inputs) + trimmed = trim_canvas(output.frames, 11, pipe.vae_temporal_compression_ratio) + assert trimmed.shape[2] < output.frames.shape[2] + + def test_the_pass_conditions_at_the_snapped_fps(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + captured = [] + original = pipe.prepare_latents + + def capture(**kwargs): + result = original(**kwargs) + captured.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(captured) == 1 + coords, num_frames, height, width = captured[0] + 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_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_public_latents_are_normalized_on_the_way_in(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + raw = torch.randn(1, 4, 5, 16, 16, device=torch_device) + packed_raw = pipe._pack_latents(raw) + _, _, clean, _, _, _ = pipe.prepare_latents( + keyframe_latents=[(8, raw[:, :, :1], 1.0)], + num_channels_latents=4, + height=32, + width=32, + num_frames=9, + noise_scale=0.0, + dtype=torch.float32, + device=torch_device, + latents=raw, + latents_normalized=False, + ) + assert not torch.allclose(clean[:, : packed_raw.shape[1]], packed_raw) + + 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_stage1_then_spatial_upsample_then_stage2(self): + components = get_dfr_dummy_components(spatial_upsampler=True) + pipe = self.get_pipeline(**{k: v for k, v in components.items() if k != "latent_upsampler"}).to(torch_device) + upsample_pipe = LTX2LatentUpsamplePipeline( + vae=components["vae"], latent_upsampler=components["latent_upsampler"] + ).to(torch_device) + + stage1 = pipe(**{**self.get_dummy_inputs(), "output_type": "latent"}) + up_video = upsample_pipe(latents=stage1.frames, height=32, width=32, output_type="latent", return_dict=False)[ + 0 + ] + up_keyframes = upsample_pipe( + latents=stage1.keyframes, height=32, width=32, output_type="latent", return_dict=False + )[0] + stage2_height = up_video.shape[-2] * pipe.vae_spatial_compression_ratio + stage2_width = up_video.shape[-1] * pipe.vae_spatial_compression_ratio + stage2_inputs = self.get_dummy_inputs() + stage2_inputs.update( + latents=up_video, + audio_latents=stage1.audio, + keyframes_latents=up_keyframes, + keyframe_positions=stage1.keyframe_positions, + reference_latents=stage1.frames, + height=stage2_height, + width=stage2_width, + sigmas=[0.75, 0.25], + noise_scale=0.75, + output_type="latent", + ) + stage2 = pipe(**stage2_inputs) + assert stage2.frames.shape[-2] == up_video.shape[-2] + assert stage2.keyframes.shape[2] == up_keyframes.shape[2] + assert stage2.keyframe_positions == stage1.keyframe_positions + + def _epilogue_plan(self, pipe, keyframe_positions=(8, 16, 24, 32, 40, 48, 56, 64)): + 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): + 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): + 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 + 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): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) + positions = [8, 16, 24, 32, 40, 48, 56, 64] + guidance = torch.cat([keyframe] * len(positions), dim=2) + _, conditioning_mask, _, _, keyframes_mask, slot_token_slice = pipe.prepare_latents( + keyframe_latents=[ + (position, guidance[:, :, index : index + 1], EPILOGUE_KEYFRAME_STRENGTH) + for index, position in enumerate(positions) + ], + num_channels_latents=4, + height=32, + width=32, + num_frames=65, + dtype=torch.float32, + device=torch_device, + generator=self.get_generator(0), + ) + assert slot_token_slice is None + 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 + assert torch.allclose(appended, torch.full_like(appended, EPILOGUE_KEYFRAME_STRENGTH)) + assert torch.count_nonzero(keyframes_mask[:, base_tokens * tokens_per_frame :]) == 0 + + def test_composed_epilogue_pins_guidance_and_tiles(self): + # `__call__` takes the tiling, not a resolved token plan: the plan needs the RoPE coordinates + # `prepare_latents` builds inside the call, so a caller cannot produce one that is guaranteed to match. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) + output = pipe( + **self.get_dummy_inputs(), + generate_slots=False, + guidance_keyframe_latents=keyframe, + guidance_keyframe_positions=[8], + freeze_audio=True, + video_tiles=[ + LTX2DFREpilogueTile( + frames=slice(0, 5), + heights=slice(0, 16), + widths=slice(0, 16), + blend_weight=torch.ones(5, 16, 16), + ) + ], + ) + assert output.frames.shape[1] == 9 + assert pipe.num_timesteps == 2 + + def test_a_single_call_tile_reproduces_the_untiled_call(self): + # The whole canvas as one tile with unit weights must be a no-op, which is what pins the token plan + # `__call__` resolves against its own coordinates. + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) + shared = { + "generate_slots": False, + "guidance_keyframe_latents": keyframe, + "guidance_keyframe_positions": [8], + "freeze_audio": True, + } + tiled = pipe( + **self.get_dummy_inputs(), + **shared, + video_tiles=[ + LTX2DFREpilogueTile( + frames=slice(0, 5), + heights=slice(0, 16), + widths=slice(0, 16), + blend_weight=torch.ones(5, 16, 16), + ) + ], + ) + untiled = pipe(**self.get_dummy_inputs(), **shared) + assert torch.allclose(tiled.frames, untiled.frames, atol=1e-4) + + @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): + 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): + 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) + latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, _ = 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), + ) + 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 + ) + plan = video_tile_plan( + [ + LTX2DFREpilogueTile( + frames=slice(0, 5), + heights=slice(0, 16), + widths=slice(0, 16), + blend_weight=torch.ones(5, 16, 16), + ) + ], + video_coords, + 5, + 16, + 16, + ) + assert len(plan) == 1 + assert plan[0].keep.numel() == latents.shape[1] + 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_distilled_euler_keeps_the_scheduler_step(self): + # Stage 1 / 2 stay on FlowMatch Euler. Re-pinning after every step is ancestral-only; doing it here + # would snap IC-LoRA reference tokens and first-frame anchors every step and change the canvas. + 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) + latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, _ = pipe.prepare_latents( + keyframe_latents=[(8, keyframe, 0.95)], + 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), + ) + 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 + ) + + def zeros_step(model_output, timestep, sample, **kwargs): + return (torch.zeros_like(sample),) + + pipe.scheduler.step = zeros_step + 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], + 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), + ), + ) + assert torch.allclose(out, torch.zeros_like(out)) + + def test_ancestral_step_does_not_erode_conditioning(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + pipe.scheduler = LTXEulerAncestralRFScheduler(eta=0.5) + pipe._interrupt = False + pipe._current_timestep = None + pipe._attention_kwargs = None + keyframe = torch.randn(1, 4, 1, 16, 16, device=torch_device) + latents, conditioning_mask, clean_latents, video_coords, keyframes_mask, _ = pipe.prepare_latents( + keyframe_latents=[(8, keyframe, 0.95)], + 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), + ), + generator=torch.Generator(device=torch_device).manual_seed(1), + ) + cos = torch.nn.functional.cosine_similarity( + out[:, block].float().flatten(), packed_keyframe.float().flatten(), dim=0 + ) + assert cos > 0.9, f"anchor tokens drifted from their conditioned content (cos={cos:.3f})" + + def test_the_epilogue_keeps_every_batch_element_distinct(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + rebuilt = pipe.rebuild_epilogue_keyframes( + torch.randn(2, 4, 2, 8, 8, device=torch_device), + decode_timestep=0.0, + decode_noise_scale=0.0, + seed=0, + device=torch.device(torch_device), + dtype=torch.float32, + ) + assert rebuilt.shape[0] == 2 + assert not torch.allclose(rebuilt[0], rebuilt[1]) + + +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 diff --git a/tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py new file mode 100644 index 000000000000..db47f77711db --- /dev/null +++ b/tests/pipelines/ltx2/test_pipeline_ltx2_dfr_temporal_refine.py @@ -0,0 +1,310 @@ +# 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 diffusers import ( + FlowMatchEulerDiscreteScheduler, + LTX2DFRPipeline, + LTX2DFRTemporalRefinePipeline, + LTXEulerAncestralRFScheduler, +) +from diffusers.pipelines.ltx2.dfr_core import ANCHOR_KEYFRAME_STRENGTH +from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition + +from ...testing_utils import enable_full_determinism, torch_device +from ..testing_utils import BasePipelineTesterConfig, PipelineTesterMixin +from .dfr_dummies import get_dfr_dummy_components, get_dfr_dummy_inputs, get_temporal_dummy_components + + +enable_full_determinism() + + +class LTX2DFRTemporalRefinePipelineTesterConfig(BasePipelineTesterConfig): + pipeline_class = LTX2DFRTemporalRefinePipeline + required_input_params_in_call_signature = frozenset( + ["prompt", "height", "width", "num_frames", "frame_rate", "prompt_embeds"] + ) + batch_input_params = frozenset(["prompt"]) + output_shape = (17, 3, 32, 32) + optional_input_params = BasePipelineTesterConfig.optional_input_params - { + "num_inference_steps", + "num_images_per_prompt", + "latents", + } + + def get_dummy_components(self): + return get_temporal_dummy_components() + + def get_dummy_inputs(self): + generator = torch.Generator("cpu").manual_seed(1) + return { + "prompt": "a robot dancing", + "generator": self.get_generator(0), + "latents": torch.randn(1, 4, 5, 16, 16, generator=generator), + "keyframes_latents": torch.randn(1, 4, 1, 16, 16, generator=generator), + "keyframe_positions": [8], + "audio_latents": torch.randn(1, 2, 8, 2, generator=generator), + "height": 32, + "width": 32, + "num_frames": 9, + "frame_rate": 25.0, + "sigmas": [0.75, 0.25], + "use_cross_timestep": False, + "max_sequence_length": 16, + "output_type": "pt", + } + + +class TestLTX2DFRTemporalRefinePipeline(LTX2DFRTemporalRefinePipelineTesterConfig, PipelineTesterMixin): + @pytest.mark.skip("Temporal refine takes a 5D latent canvas that is not prompt-batched.") + def test_inference_batch_consistent(self, *args, **kwargs): + pass + + @pytest.mark.skip("Temporal refine takes a 5D latent canvas that is not prompt-batched.") + def test_inference_batch_single_identical(self, *args, **kwargs): + pass + + def test_temporal_upsample_round_doubles_the_frame_count(self): + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + output = pipe(**self.get_dummy_inputs()) + assert output.frames.shape[1] == (9 - 1) * 2 + 1 + assert pipe.num_timesteps == 2 + + @pytest.mark.parametrize( + ("scheduler", "message"), + [ + (FlowMatchEulerDiscreteScheduler(), "LTXEulerAncestralRFScheduler"), + (LTXEulerAncestralRFScheduler(eta=0.0), "stochastic step"), + ], + ) + def test_a_non_ancestral_scheduler_is_refused(self, scheduler, message): + # `denoise` falls back to a deterministic Euler step for anything else, so an unchecked scheduler here + # would run to completion and just return a softer canvas. + components = self.get_dummy_components() + components["scheduler"] = scheduler + pipe = self.get_pipeline(**components).to(torch_device) + with pytest.raises(ValueError, match=message): + pipe(**self.get_dummy_inputs()) + + def test_temporal_round_tiles_get_distinct_ancestral_noise(self): + dfr = LTX2DFRPipeline(**get_dfr_dummy_components()).to(torch_device) + stage = dfr(**get_dfr_dummy_inputs(generator=self.get_generator(0), num_frames=17, output_type="latent")) + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + seeds = [] + original = pipe.denoise + + def capture(**kwargs): + generator = kwargs.get("generator") + seeds.append(generator.initial_seed() if generator is not None else None) + return original(**kwargs) + + pipe.denoise = capture + try: + pipe( + prompt="a robot dancing", + latents=stage.frames, + keyframes_latents=stage.keyframes, + keyframe_positions=stage.keyframe_positions, + audio_latents=stage.audio, + height=32, + width=32, + num_frames=17, + frame_rate=25.0, + sigmas=[0.75, 0.25], + use_cross_timestep=False, + max_sequence_length=16, + output_type="latent", + generator=self.get_generator(0), + ) + finally: + del pipe.denoise + + assert seeds == [1000 * 1 + 0, 1000 * 1 + 1] + + def test_a_condition_keeps_its_moment_through_the_refine_round(self): + dfr = LTX2DFRPipeline(**get_dfr_dummy_components()).to(torch_device) + image = PIL.Image.fromarray(np.zeros((32, 32, 3), dtype=np.uint8)) + stage = dfr( + **get_dfr_dummy_inputs( + generator=self.get_generator(0), + num_frames=9, + output_type="latent", + conditions=[LTX2VideoCondition(frames=image, index=4, strength=1.0, crf=0)], + ) + ) + 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 + + pipe.prepare_latents = capture + try: + pipe( + prompt="a robot dancing", + latents=stage.frames, + keyframes_latents=stage.keyframes, + keyframe_positions=stage.keyframe_positions, + audio_latents=stage.audio, + conditions=[LTX2VideoCondition(frames=image, index=4, strength=1.0, crf=0)], + height=32, + width=32, + num_frames=9, + condition_num_frames=9, + frame_rate=25.0, + sigmas=[0.75, 0.25], + use_cross_timestep=False, + max_sequence_length=16, + output_type="latent", + generator=self.get_generator(0), + ) + finally: + del pipe.prepare_latents + + 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) == 14, f"condition landed at pixel {start}, expected 14" + end = video_coords[0, 0, condition_token, 1].item() * frame_rate + assert round(end - start) == 1 + + def test_last_frame_condition_stays_at_the_end_across_a_temporal_round(self): + dfr = LTX2DFRPipeline(**get_dfr_dummy_components()).to(torch_device) + frame = np.full((32, 32, 3), 200, dtype=np.uint8) + last_stage = dfr( + **get_dfr_dummy_inputs( + generator=self.get_generator(0), + num_frames=17, + output_type="latent", + conditions=LTX2VideoCondition(frames=frame, index=-1, strength=1.0, crf=0), + ) + ) + first_stage = dfr( + **get_dfr_dummy_inputs( + generator=self.get_generator(0), + num_frames=17, + output_type="latent", + conditions=LTX2VideoCondition(frames=frame, index=0, strength=1.0, crf=0), + ) + ) + pipe = self.get_pipeline(**self.get_dummy_components()).to(torch_device) + shared = { + "prompt": "a robot dancing", + "height": 32, + "width": 32, + "num_frames": 17, + "frame_rate": 25.0, + "sigmas": [0.75, 0.25], + "use_cross_timestep": False, + "max_sequence_length": 16, + "output_type": "pt", + "generator": self.get_generator(1), + } + last_out = pipe( + **shared, + latents=last_stage.frames, + keyframes_latents=last_stage.keyframes, + keyframe_positions=last_stage.keyframe_positions, + audio_latents=last_stage.audio, + conditions=LTX2VideoCondition(frames=frame, index=-1, strength=1.0, crf=0), + condition_num_frames=17, + ).frames + first_out = pipe( + **shared, + latents=first_stage.frames, + keyframes_latents=first_stage.keyframes, + keyframe_positions=first_stage.keyframe_positions, + audio_latents=first_stage.audio, + conditions=LTX2VideoCondition(frames=frame, index=0, strength=1.0, crf=0), + condition_num_frames=17, + ).frames + assert last_out.shape[1] == (17 - 1) * 2 + 1 + assert not torch.allclose(last_out, first_out) + + def test_a_carried_slot_is_the_copy_the_stitched_canvas_kept(self): + dfr = LTX2DFRPipeline(**get_dfr_dummy_components()).to(torch_device) + stage = dfr(**get_dfr_dummy_inputs(generator=self.get_generator(0), num_frames=17, output_type="latent")) + 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 + + pipe.prepare_latents, pipe.denoise = capture_prepare, capture_denoise + try: + round1 = pipe( + prompt="a robot dancing", + latents=stage.frames, + keyframes_latents=stage.keyframes, + keyframe_positions=stage.keyframe_positions, + audio_latents=stage.audio, + height=32, + width=32, + num_frames=17, + frame_rate=25.0, + round_index=1, + sigmas=[0.75, 0.25], + use_cross_timestep=False, + max_sequence_length=16, + output_type="latent", + generator=self.get_generator(0), + ) + pipe( + prompt="a robot dancing", + latents=round1.frames, + keyframes_latents=round1.keyframes, + keyframe_positions=round1.keyframe_positions, + audio_latents=stage.audio, + height=32, + width=32, + num_frames=(round1.frames.shape[2] - 1) * pipe.vae_temporal_compression_ratio + 1, + frame_rate=50.0, + source_seconds=17 / 25.0, + round_index=2, + sigmas=[0.75, 0.25], + use_cross_timestep=False, + max_sequence_length=16, + output_type="latent", + generator=self.get_generator(0), + ) + finally: + del pipe.prepare_latents, pipe.denoise + + assert len(denoised) == 2 + 4 + first_tile_slot = denoised[0][:, slot_slices[0]] + second_tile_slot = denoised[1][:, slot_slices[1]][:, : first_tile_slot.shape[1]] + assert not torch.allclose(first_tile_slot, second_tile_slot) + round_2_anchors = anchors[2] + 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)