diff --git a/docs/source/en/api/pipelines/ltx2.md b/docs/source/en/api/pipelines/ltx2.md index e73689b4af52..9924825a6428 100644 --- a/docs/source/en/api/pipelines/ltx2.md +++ b/docs/source/en/api/pipelines/ltx2.md @@ -1028,6 +1028,100 @@ encode_video( ) ``` +#### Diffusion Fidelity Rendering (DFR) + +`LTX2DFRBlocks` trades wall-clock time for detail fidelity. It generates on a canvas padded to a whole number of keyframe segments and spends one extra latent frame of tokens per segment border on a **keyframe slot** — a single-pixel-frame latent the model fills in. Relaxing the effective temporal compression at those positions means the surrounding video is conditioned on genuinely new frames rather than interpolated ones. This needs a transformer whose config sets `use_keyframes_abs_pos_embedding`, which LTX-2.5 checkpoints ship. + +The recipe is two passes of the same blocks: a base pass at half resolution, then a detailing pass at full resolution seeded from it. Both the video latents and the keyframe slots are upsampled in between, and the spatial detailing IC-LoRA applies to the second pass only — switched on in the seam between the two calls, like the [stage 2 distilled LoRA](#stage-2-with-the-distilled-lora). The DFR schedules are distilled and run without guidance, so these blocks carry no guider and take no `guidance_scale`. + +```py +import torch +from diffusers import ComponentsManager +from diffusers.modular_pipelines import LTX2DFRBlocks +from diffusers.pipelines.ltx2 import LTX2LatentUpsamplePipeline +from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel +from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES +from diffusers.utils import encode_video + +device, model_path = "cuda", "Lightricks/LTX-2.5-Diffusers" +height, width, frame_rate = 704, 1216, 24.0 +prompt = "A cinematic shot of a red fox walking through a snowy forest at dawn." + +cm = ComponentsManager() +pipe = LTX2DFRBlocks().init_pipeline(model_path, components_manager=cm) +pipe.load_components(dtype=torch.bfloat16) + +# Load the detailing IC-LoRA *before* offload is enabled. `load_lora_adapter` restores offload +# afterwards through `DiffusionPipeline.enable_model_cpu_offload`, which a `ModularPipeline` does +# not implement, and the upsample pipeline below shares this `vae` -- so offloading first leaves +# accelerate hooks on it and the LoRA load then fails. +pipe.load_lora_weights("Lightricks/LTX-2.5-22b-IC-LoRA-Pixel-Spatial-Upscaler", adapter_name="detailing") +pipe.disable_lora() # pass 1 runs on the base model +cm.enable_auto_cpu_offload(device=device, memory_reserve_margin="20GB") + +# One generator across both passes, so pass 2 continues the noise stream. +generator = torch.Generator(device).manual_seed(42) +common = dict(prompt=prompt, frame_rate=frame_rate, generator=generator) + +# Pass 1: half resolution. Returns the video latents plus one keyframe slot per segment border. +# `num_frames` is omitted, so the duration head picks the length. +first = pipe( + **common, height=height // 2, width=width // 2, sigmas=DISTILLED_SIGMA_VALUES, output_type="latent" +) +video_latents, keyframes_latents = first.get("videos"), first.get("keyframes_latents") +num_frames = first.get("num_frames") + +# Upsample the video *and* the keyframe slots. `latents_normalized=False`: `output_type="latent"` +# already applied the latent statistics. +upsampler = LTX2LatentUpsamplerModel.from_pretrained( + model_path, subfolder="latent_upsampler", dtype=torch.bfloat16 +) +upsample_pipe = LTX2LatentUpsamplePipeline(vae=pipe.vae, latent_upsampler=upsampler) +upsample_pipe.enable_model_cpu_offload(device=device) + + +def upsample(latents): + return upsample_pipe( + latents=latents, latents_normalized=False, output_type="latent", return_dict=False + )[0] + + +# The detailing IC-LoRA belongs to pass 2 alone, which is why this is one blockset run twice rather +# than one call: the adapter is switched on in the seam between them. Calibrated for strength 0.5. +pipe.enable_lora() +pipe.set_adapters(["detailing"], adapter_weights=[0.5]) + +# Pass 2: full resolution, seeded from pass 1 -- video, audio, the keyframe slots, and the +# half-resolution result as the detailing adapter's in-context reference. +pipe.vae.enable_tiling() +out = pipe( + **common, + height=height, + width=width, + num_frames=num_frames, + latents=upsample(video_latents), + keyframes_latents=upsample(keyframes_latents), + audio_latents=first.get("audio"), + detailing_reference_latents=video_latents, + detailing_reference_downscale_factor=2, + sigmas=STAGE_2_DISTILLED_SIGMA_VALUES, + noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0], + output_type="np", +) + +encode_video( + out.get("videos")[0], + fps=frame_rate, + audio=out.get("audio")[0].float().cpu(), + audio_sample_rate=pipe.vocoder.config.output_sampling_rate, + output_path="ltx2_5_dfr.mp4", +) +``` + +`height` and `width` are the output resolution and must be divisible by twice the VAE's spatial compression ratio (64 for LTX-2.5), since the base pass runs at half of each axis. Whatever `num_frames` asks for, the canvas is padded onto the segment grid internally and trimmed back before decoding, so the caller always gets the frame count it requested. + +DFR decodes with the convolutional `vae`, matching the reference implementation. For maximum detail fidelity, run the second pass with `output_type="latent"` and hand the latents to [`LTX2VideoDiffusionDecodePipeline`] instead. + You can see the supported workflows in the docs for each blockset (e.g. [`LTX2AutoBlocks`], [`LTX25AutoBlocks`]). ## LTX2Pipeline @@ -1086,6 +1180,14 @@ You can see the supported workflows in the docs for each blockset (e.g. [`LTX2Au [[autodoc]] LTX25AutoBlocks +## LTX2DFRModularPipeline + +[[autodoc]] LTX2DFRModularPipeline + +## LTX2DFRBlocks + +[[autodoc]] LTX2DFRBlocks + ## LTX2Guidance [[autodoc]] modular_pipelines.ltx2.guider.LTX2Guidance diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index b7d79b8ee97d..eb4aa8688b73 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -545,6 +545,8 @@ "LTX25AutoBlocks", "LTX25ModularPipeline", "LTX2AutoBlocks", + "LTX2DFRBlocks", + "LTX2DFRModularPipeline", "LTX2ModularPipeline", "LTXAutoBlocks", "LTXModularPipeline", @@ -1393,6 +1395,8 @@ Krea2TurboAutoBlocks, Krea2TurboModularPipeline, LTX2AutoBlocks, + LTX2DFRBlocks, + LTX2DFRModularPipeline, LTX2ModularPipeline, LTX25AutoBlocks, LTX25ModularPipeline, 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/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 81b93f88f515..0f1ffa7e5be7 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -124,8 +124,10 @@ _import_structure["ltx2"] = [ "LTX2AutoBlocks", "LTX25AutoBlocks", + "LTX2DFRBlocks", "LTX2ModularPipeline", "LTX25ModularPipeline", + "LTX2DFRModularPipeline", ] _import_structure["minimax_h3"] = [ "MiniMaxH3Blocks", @@ -189,7 +191,14 @@ Krea2TurboModularPipeline, ) from .ltx import LTXAutoBlocks, LTXModularPipeline - from .ltx2 import LTX2AutoBlocks, LTX2ModularPipeline, LTX25AutoBlocks, LTX25ModularPipeline + from .ltx2 import ( + LTX2AutoBlocks, + LTX2DFRBlocks, + LTX2DFRModularPipeline, + LTX2ModularPipeline, + LTX25AutoBlocks, + LTX25ModularPipeline, + ) from .minimax_h3 import ( MiniMaxH3Blocks, MiniMaxH3ModularPipeline, diff --git a/src/diffusers/modular_pipelines/ltx2/__init__.py b/src/diffusers/modular_pipelines/ltx2/__init__.py index caf44b9a6179..e4a097fea570 100644 --- a/src/diffusers/modular_pipelines/ltx2/__init__.py +++ b/src/diffusers/modular_pipelines/ltx2/__init__.py @@ -29,7 +29,12 @@ "LTX2InContextBlocks", ] _import_structure["modular_blocks_ltx25"] = ["LTX25AutoBlocks"] - _import_structure["modular_pipeline"] = ["LTX2ModularPipeline", "LTX25ModularPipeline"] + _import_structure["modular_blocks_ltx2_dfr"] = ["LTX2DFRBlocks"] + _import_structure["modular_pipeline"] = [ + "LTX2DFRModularPipeline", + "LTX2ModularPipeline", + "LTX25ModularPipeline", + ] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: try: @@ -45,8 +50,9 @@ LTX2ImageToVideoBlocks, LTX2InContextBlocks, ) + from .modular_blocks_ltx2_dfr import LTX2DFRBlocks from .modular_blocks_ltx25 import LTX25AutoBlocks - from .modular_pipeline import LTX2ModularPipeline, LTX25ModularPipeline + from .modular_pipeline import LTX2DFRModularPipeline, LTX2ModularPipeline, LTX25ModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/ltx2/before_denoise.py b/src/diffusers/modular_pipelines/ltx2/before_denoise.py index 81ffc28188ea..e40b37ef6fce 100644 --- a/src/diffusers/modular_pipelines/ltx2/before_denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/before_denoise.py @@ -30,6 +30,7 @@ from ...utils.torch_utils import randn_tensor from ..modular_pipeline import ModularPipelineBlocks, PipelineState from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam +from .utils import resolve_canvas logger = logging.get_logger(__name__) @@ -1389,6 +1390,481 @@ def __call__(self, components, state: PipelineState) -> PipelineState: return components, state +class LTX2DFRPlanStep(ModularPipelineBlocks): + model_name = "ltx2.5-dfr" + + @property + def description(self) -> str: + return ( + "Resolves the DFR keyframe segment grid: pads `num_frames` up to a whole number of keyframe segments " + "and reports the pixel-frame positions the pipeline generates keyframe slots at. The padding is trimmed " + "back by `LTX2DFRSplitKeyframesStep` before decoding, so the caller always gets the frame count it asked " + "for." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", LTX2VideoTransformer3DModel), + ComponentSpec("vae", AutoencoderKLLTX2Video), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "num_frames", + type_hint=int, + default=None, + description=( + "The number of frames the caller asked for, before the canvas is padded onto the segment grid. " + "Omit to auto-predict via the `duration_head` (see `LTX2AutoDurationStep`)." + ), + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "num_frames", + type_hint=int, + description=( + "The padded canvas frame count every later block generates on, a whole number of keyframe " + "segments plus one." + ), + ), + OutputParam( + "requested_num_frames", + type_hint=int, + description="The frame count the caller asked for, restored before decoding.", + ), + OutputParam( + "slot_frame_indices", + type_hint=list, + description=( + "Pixel-frame positions of the generated keyframe slots, `[S, 2S, ..., num_frames - 1]` for the " + "chosen segment length `S`." + ), + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + if not components.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 components.transformer_temporal_patch_size != 1: + raise ValueError( + "DFR appends one latent frame of tokens per keyframe slot, which a temporal patch size above 1 " + f"cannot represent, but the transformer patchifies time by " + f"{components.transformer_temporal_patch_size}." + ) + if block_state.num_frames is None: + raise ValueError( + "`num_frames` must be a concrete integer here. Pass `num_frames`, or use a blockset that runs " + "`LTX2AutoDurationStep` on a checkpoint shipping a `duration_head`." + ) + + block_state.requested_num_frames = block_state.num_frames + block_state.num_frames, _, block_state.slot_frame_indices = resolve_canvas( + block_state.num_frames, components.vae_temporal_compression_ratio + ) + + self.set_block_state(state, block_state) + return components, state + + +class LTX2DFRPrepareLatentsStep(ModularPipelineBlocks): + model_name = "ltx2.5-dfr" + + @property + def description(self) -> str: + return ( + "Prepares the packed video latents for one DFR pass. The sequence is laid out as " + "`[base | keyframes | slots | reference]`: base tokens cover the target latent grid (seeded from " + "`latents`), frame conditions are placed exactly as in `LTX2ConditionPrepareLatentsStep`, each entry of " + "`slot_frame_indices` appends one latent frame's worth of *generated* keyframe tokens spanning a single " + "pixel frame, and `detailing_reference_latents` appends a fully clean in-context reference for the " + "spatial detailing IC-LoRA. Slots are what buy DFR its extra frames: they carry conditioning mask 0 and " + "are marked in `video_keyframes_mask` so the transformer adds its learned single-frame embedding." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", LTX2VideoTransformer3DModel), + ComponentSpec("vae", AutoencoderKLLTX2Video), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "slot_frame_indices", + type_hint=list, + required=True, + description="Pixel-frame positions of the generated keyframe slots, from `LTX2DFRPlanStep`.", + ), + InputParam( + "keyframes_latents", + type_hint=torch.Tensor, + description=( + "`[B, C, num_slots, H, W]` content seeding the keyframe slots: a previous DFR pass's " + "`keyframes_latents` upsampled to this pass's resolution. Denormalized, like every latent " + "crossing the pipeline boundary. Slots start from noise when omitted." + ), + ), + InputParam( + "detailing_reference_latents", + type_hint=torch.Tensor, + description=( + "`[B, C, F, H, W]` latents appended as a fully clean in-context reference for the spatial " + "detailing IC-LoRA: the previous pass's output at its own resolution, denormalized. Only " + "meaningful with that adapter loaded." + ), + ), + InputParam( + "detailing_reference_downscale_factor", + type_hint=int, + default=2, + description=( + "Ratio between this pass's resolution and the reference's, used to scale the reference tokens' " + "spatial coordinates into the target coordinate space. Must match the factor the IC-LoRA was " + "trained with." + ), + ), + InputParam( + "condition_latents", + type_hint=list, + description="Per-condition normalized VAE latents of shape [1, C, F, H, W].", + ), + InputParam( + "condition_strengths", + type_hint=list, + description="Per-condition conditioning strengths.", + ), + InputParam( + "condition_indices", + type_hint=list, + description="Per-condition latent frame index at which the condition is applied.", + ), + InputParam( + "condition_pixel_frames", + type_hint=list, + description="Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords.", + ), + InputParam.template("latents"), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam( + "num_frames", + type_hint=int, + required=True, + description="The padded canvas frame count, from `LTX2DFRPlanStep`.", + ), + InputParam( + "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." + ), + InputParam( + "noise_scale", + type_hint=float, + default=None, + description=( + "Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` " + "when custom `sigmas` are supplied, else 1.0." + ), + ), + InputParam.template("sigmas"), + InputParam.template("num_images_per_prompt", name="num_videos_per_prompt"), + InputParam( + "batch_size", + type_hint=int, + required=True, + description="The number of prompts being denoised, used to expand conditioning per prompt.", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description="Packed noisy video latents, with keyframe, slot and reference tokens appended.", + ), + OutputParam( + "conditioning_mask", + type_hint=torch.Tensor, + description=( + "Packed per-token conditioning strengths of shape [B, S, 1] in [0, 1]: 1 at fully-conditioned " + "positions, 0 at free positions, including every keyframe slot." + ), + ), + OutputParam( + "clean_latents", + type_hint=torch.Tensor, + description="Clean condition latents at conditioned positions, zeros elsewhere; same shape as `latents`.", + ), + OutputParam( + "video_keyframes_mask", + type_hint=torch.Tensor, + kwargs_type="denoiser_input_fields", + description=( + "Packed [B, S, 1] marker, 1 on tokens whose latent frame encodes a single pixel frame -- the " + "causal first frame and every generated keyframe slot. Those tokens receive the transformer's " + "`keyframes_abs_pos_embedding`." + ), + ), + OutputParam( + "appended_coords", + type_hint=torch.Tensor, + description=( + "RoPE coordinates of shape [B, 3, num_appended_tokens, 2] for the appended keyframe, slot and " + "reference tokens, in the order they were appended." + ), + ), + OutputParam( + "base_token_count", + type_hint=int, + description="Number of generated-video tokens, i.e. the sequence length before appended tokens.", + ), + OutputParam( + "slot_token_slice", + type_hint=slice, + description="Slice of the packed sequence holding the generated keyframe slot tokens.", + ), + OutputParam( + "noise_scale", + type_hint=float, + description="The resolved initial noise level, forwarded to the audio latents step.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + batch_size = block_state.batch_size * block_state.num_videos_per_prompt + spatial_patch = components.transformer_spatial_patch_size + temporal_patch = components.transformer_temporal_patch_size + frame_scale_factor = components.vae_temporal_compression_ratio + + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + latent_num_frames = (block_state.num_frames - 1) // frame_scale_factor + 1 + tokens_per_latent_frame = (latent_height // spatial_patch) * (latent_width // spatial_patch) + + noise_scale = block_state.noise_scale + if noise_scale is None: + noise_scale = block_state.sigmas[0] if block_state.sigmas is not None else 1.0 + + if isinstance(block_state.generator, list): + logger.warning( + f"{self.__class__.__name__} does not support using a list of generators. The first generator in the" + f" list will be used for all (pseudo-)random operations." + ) + + if block_state.latents is not None: + latents = _normalize_latents( + block_state.latents, + components.latents_mean, + components.latents_std, + components.vae_scaling_factor, + ) + else: + shape = ( + batch_size, + components.transformer.config.in_channels, + latent_num_frames, + latent_height, + latent_width, + ) + latents = torch.zeros(shape, device=device, dtype=torch.float32) + + conditioning_mask = latents.new_zeros((batch_size, 1, latent_num_frames, latent_height, latent_width)) + latents = _pack_latents(latents, spatial_patch, temporal_patch) + conditioning_mask = _pack_latents(conditioning_mask, spatial_patch, temporal_patch) # [B, S, 1] + + if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape[:2]: + raise ValueError( + f"Provided `latents` tensor packs to shape {latents.shape}, but the expected packed shape is " + f"{conditioning_mask.shape[:2] + (components.transformer.config.in_channels,)}." + ) + + base_token_count = latents.shape[1] + clean_latents = torch.zeros_like(latents) + # Causal encoding gives the first latent frame a temporal stride of one pixel frame, so it carries the same + # marker the generated slots do. + keyframes_mask = torch.zeros_like(conditioning_mask) + keyframes_mask[:, :tokens_per_latent_frame] = 1.0 + + condition_latents = block_state.condition_latents or [] + condition_strengths = block_state.condition_strengths or [] + condition_indices = block_state.condition_indices or [] + condition_pixel_frames = block_state.condition_pixel_frames or [] + condition_latents_packed = [_pack_latents(cond, spatial_patch, temporal_patch) for cond in condition_latents] + + # First-frame conditions (latent index 0): overwrite the tokens at the first-frame positions. Condition + # tensors carry batch 1 and broadcast across the generation batch. + for cond, strength, latent_idx in zip(condition_latents_packed, condition_strengths, condition_indices): + if latent_idx != 0: + continue + num_cond_tokens = cond.size(1) + latents[:, :num_cond_tokens] = cond + conditioning_mask[:, :num_cond_tokens] = strength + clean_latents[:, :num_cond_tokens] = cond + + scale_factors = ( + frame_scale_factor, + components.vae_spatial_compression_ratio, + components.vae_spatial_compression_ratio, + ) + appended_coords = [] + + # Non-first-frame ("keyframe") conditions (latent index > 0): appended as extra tokens carrying their content + # in both `latents` and `clean_latents`, so the noising below leaves them at (1 - strength) * noise_scale. + for cond_5d, cond_packed, strength, latent_idx, num_pixel_frames in zip( + condition_latents, + condition_latents_packed, + condition_strengths, + condition_indices, + condition_pixel_frames, + ): + if latent_idx == 0: + continue + + _, _, kf_latent_frames, kf_latent_height, kf_latent_width = cond_5d.shape + coords = _prepare_keyframe_coords( + keyframe_latent_num_frames=kf_latent_frames, + keyframe_latent_height=kf_latent_height, + keyframe_latent_width=kf_latent_width, + pixel_frame_idx=(latent_idx - 1) * frame_scale_factor + 1, + num_pixel_frames=num_pixel_frames, + fps=block_state.frame_rate, + patch_size=spatial_patch, + patch_size_t=temporal_patch, + scale_factors=scale_factors, + device=device, + ) + tokens = cond_packed.expand(batch_size, -1, -1) + latents = torch.cat([latents, 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))], dim=1 + ) + appended_coords.append(coords.expand(batch_size, -1, -1, -1)) + + # Generated keyframe slots: fully-denoised single-pixel-frame token blocks the model fills in. Their seed + # goes into `latents` only -- `clean_latents` stays zero, and mask 0 means the noising treats them like any + # other free token. + if block_state.keyframes_latents is not None: + keyframes_latents = _normalize_latents( + block_state.keyframes_latents, + components.latents_mean, + components.latents_std, + components.vae_scaling_factor, + ).to(device=device, dtype=latents.dtype) + slot_tokens = torch.cat( + [ + _pack_latents(keyframes_latents[:, :, index : index + 1], spatial_patch, temporal_patch) + for index in range(keyframes_latents.shape[2]) + ], + dim=1, + ).expand(batch_size, -1, -1) + else: + num_slot_tokens = tokens_per_latent_frame * len(block_state.slot_frame_indices) + slot_tokens = latents.new_zeros((batch_size, num_slot_tokens, latents.shape[2])) + + 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( + _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=block_state.frame_rate, + patch_size=spatial_patch, + patch_size_t=temporal_patch, + scale_factors=scale_factors, + device=device, + ).expand(batch_size, -1, -1, -1) + for position in block_state.slot_frame_indices + ) + + # Spatial detailing IC-LoRA reference, held fully clean at mask 1. + if block_state.detailing_reference_latents is not None: + reference_latents = _normalize_latents( + block_state.detailing_reference_latents, + components.latents_mean, + components.latents_std, + components.vae_scaling_factor, + ).to(device=device, dtype=latents.dtype) + reference_tokens = _pack_latents(reference_latents, spatial_patch, temporal_patch).expand( + batch_size, -1, -1 + ) + reference_coords = components.transformer.rope.prepare_video_coords( + batch_size=batch_size, + num_frames=reference_latents.shape[2], + height=reference_latents.shape[3], + width=reference_latents.shape[4], + device=device, + fps=block_state.frame_rate, + ) + reference_coords[:, 1:, :, :] = ( + reference_coords[:, 1:, :, :] * block_state.detailing_reference_downscale_factor + ) + + latents = torch.cat([latents, 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) + + # Mask semantics: 0 -> fully noised, 1 -> kept clean, in between -> noise level (1 - mask) * noise_scale. + noise = randn_tensor( + latents.shape, generator=block_state.generator, device=latents.device, dtype=latents.dtype + ) + scaled_mask = (1.0 - conditioning_mask) * noise_scale + block_state.latents = noise * scaled_mask + latents * (1 - scaled_mask) + + block_state.conditioning_mask = conditioning_mask + block_state.clean_latents = clean_latents + block_state.video_keyframes_mask = keyframes_mask + block_state.appended_coords = torch.cat(appended_coords, dim=2) + block_state.base_token_count = base_token_count + block_state.slot_token_slice = slot_token_slice + block_state.noise_scale = noise_scale + + self.set_block_state(state, block_state) + return components, state + + class LTX2BuildVideoSelfAttentionMaskStep(ModularPipelineBlocks): model_name = "ltx2" diff --git a/src/diffusers/modular_pipelines/ltx2/decoders.py b/src/diffusers/modular_pipelines/ltx2/decoders.py index fc957a3f9925..998470bc97b0 100644 --- a/src/diffusers/modular_pipelines/ltx2/decoders.py +++ b/src/diffusers/modular_pipelines/ltx2/decoders.py @@ -125,6 +125,103 @@ def __call__(self, components, state: PipelineState) -> PipelineState: return components, state +class LTX2DFRSplitKeyframesStep(ModularPipelineBlocks): + model_name = "ltx2.5-dfr" + + @property + def description(self) -> str: + return ( + "Splits a denoised DFR sequence into the parts the caller needs: the generated keyframe slots come out " + "as `keyframes_latents`, the appended tokens are dropped, and the canvas is trimmed back to the frame " + "count originally requested. Feed both outputs into the next DFR pass -- upsampled -- to run the " + "detailing stage." + ) + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("latents", required=True), + InputParam( + "base_token_count", + type_hint=int, + required=True, + description="Number of generated-video tokens, i.e. the sequence length before appended tokens.", + ), + InputParam( + "slot_token_slice", + type_hint=slice, + required=True, + description="Slice of the packed sequence holding the generated keyframe slot tokens.", + ), + InputParam( + "requested_num_frames", + type_hint=int, + required=True, + description="The frame count the caller asked for, which the canvas is trimmed back to.", + ), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", + type_hint=torch.Tensor, + description=( + "Denoised latents for the generated video, with appended tokens removed and the canvas padding " + "trimmed off." + ), + ), + OutputParam( + "keyframes_latents", + type_hint=torch.Tensor, + description=( + "Denormalized `[B, C, num_slots, H, W]` generated keyframe slots, one latent frame per slot " + "position. Upsample these alongside the video latents to seed the next DFR pass." + ), + ), + OutputParam( + "num_frames", + type_hint=int, + description="The trimmed frame count, restored to what the caller asked for.", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + spatial_patch = components.transformer_spatial_patch_size + temporal_patch = components.transformer_temporal_patch_size + + # A keyframe slot is one latent frame, so `LTX2DFRPlanStep` rejects a temporal patch size above 1 and the + # token counts here are per latent frame. + tokens_per_latent_frame = (latent_height // spatial_patch) * (latent_width // spatial_patch) + + slot_tokens = block_state.latents[:, block_state.slot_token_slice] + num_slots = slot_tokens.shape[1] // tokens_per_latent_frame + keyframes_latents = _unpack_latents( + slot_tokens, num_slots, latent_height, latent_width, spatial_patch, temporal_patch + ) + block_state.keyframes_latents = _denormalize_latents( + keyframes_latents, components.latents_mean, components.latents_std, components.vae_scaling_factor + ) + + # Drop the appended keyframe, slot and reference tokens, as `LTX2TrimConditionTokensStep` does. + latents = block_state.latents[:, : block_state.base_token_count] + # Then trim the canvas padding. `_pack_latents` is frame-major, so that is a prefix slice on what is left. + trimmed_latent_frames = (block_state.requested_num_frames - 1) // components.vae_temporal_compression_ratio + 1 + block_state.latents = latents[:, : trimmed_latent_frames * tokens_per_latent_frame] + block_state.num_frames = block_state.requested_num_frames + + self.set_block_state(state, block_state) + return components, state + + class LTX2DiffusionVaeDecoderStep(ModularPipelineBlocks): model_name = "ltx2" diff --git a/src/diffusers/modular_pipelines/ltx2/denoise.py b/src/diffusers/modular_pipelines/ltx2/denoise.py index 5cc5a4e57abc..1115373fac4c 100644 --- a/src/diffusers/modular_pipelines/ltx2/denoise.py +++ b/src/diffusers/modular_pipelines/ltx2/denoise.py @@ -27,7 +27,7 @@ ModularPipelineBlocks, PipelineState, ) -from ..modular_pipeline_utils import ComponentSpec, InputParam +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam from .guider import LTX2Guidance @@ -643,6 +643,119 @@ def __call__(self, components, state: PipelineState) -> PipelineState: return components, state +class LTX2DFRLoopDenoiser(ModularPipelineBlocks): + model_name = "ltx2.5-dfr" + + @property + def description(self) -> str: + return ( + "Joint video+audio denoiser for DFR. One transformer call per step and no guider: the DFR sigma " + "schedules are distilled and trained to run without guidance, so there is nothing to combine and no " + "negative conditioning to carry. Converts both streams' velocity to x0, which is the space " + "`LTX2ConditionLoopAfterDenoiser` expects." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", LTX2VideoTransformer3DModel), + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("connector_prompt_embeds", required=True, type_hint=torch.Tensor), + InputParam("connector_audio_prompt_embeds", required=True, type_hint=torch.Tensor), + InputParam("connector_attention_mask", required=True, type_hint=torch.Tensor), + InputParam.template("denoiser_input_fields"), + InputParam.template("height", default=512), + InputParam.template("width", default=704), + InputParam("num_frames", type_hint=int, required=True), + InputParam( + "frame_rate", type_hint=float, default=24.0, description="Frames per second of the generated video." + ), + InputParam( + "use_cross_timestep", + type_hint=bool, + default=True, + description="Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+).", + ), + InputParam.template("attention_kwargs"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("noise_pred_video", type_hint=torch.Tensor, description="Video x0 prediction for this step."), + OutputParam("noise_pred_audio", type_hint=torch.Tensor, description="Audio x0 prediction for this step."), + ] + + @torch.no_grad() + def __call__(self, components, block_state: BlockState, i: int, t: torch.Tensor): + latent_num_frames = (block_state.num_frames - 1) // components.vae_temporal_compression_ratio + 1 + latent_height = block_state.height // components.vae_spatial_compression_ratio + latent_width = block_state.width // components.vae_spatial_compression_ratio + + # Upstream-produced transformer arguments (`audio_num_frames`, `video_coords`, `audio_coords`, and DFR's + # `video_keyframes_mask`) arrive tagged `denoiser_input_fields`; filter them against the signature. The + # latent dims are computed here so their names don't clash with the pixel-space ones in state. + transformer_args = set(inspect.signature(components.transformer.forward).parameters) + transformer_kwargs = {k: v for k, v in block_state.denoiser_input_fields.items() if k in transformer_args} + transformer_kwargs.update( + num_frames=latent_num_frames, + height=latent_height, + width=latent_width, + fps=block_state.frame_rate, + use_cross_timestep=block_state.use_cross_timestep, + attention_kwargs=block_state.attention_kwargs, + perturbation_mask=None, + # No STG and no modality isolation: both are guidance perturbations, and this pass is the plain + # conditional forward. + spatio_temporal_guidance_blocks=None, + isolate_modalities=False, + encoder_hidden_states=block_state.connector_prompt_embeds, + audio_encoder_hidden_states=block_state.connector_audio_prompt_embeds, + encoder_attention_mask=block_state.connector_attention_mask, + audio_encoder_attention_mask=block_state.connector_attention_mask, + ) + + with components.transformer.cache_context("cond"): + noise_pred_video, noise_pred_audio = components.transformer( + hidden_states=block_state.latent_model_input, + audio_hidden_states=block_state.audio_latent_model_input, + timestep=block_state.video_timestep, + audio_timestep=block_state.audio_timestep, + sigma=block_state.audio_timestep, # plain (unmasked) timestep, used by LTX-2.3 + return_dict=False, + **transformer_kwargs, + ) + + # The after block converts back to velocity, so hand it x0 -- the same space the guided path combines in. + block_state.noise_pred_video = convert_velocity_to_x0( + block_state.latents, noise_pred_video.float(), i, components.scheduler + ) + block_state.noise_pred_audio = convert_velocity_to_x0( + block_state.audio_latents, noise_pred_audio.float(), i, block_state.audio_scheduler + ) + return components, block_state + + +class LTX2DFRDenoiseStep(LTX2DenoiseLoopWrapper): + model_name = "ltx2.5-dfr" + block_classes = [LTX2ConditionLoopBeforeDenoiser, LTX2DFRLoopDenoiser, LTX2ConditionLoopAfterDenoiser] + block_names = ["before_denoiser", "denoiser", "after_denoiser"] + + @property + def description(self) -> str: + return ( + "DFR denoise step. Identical to `LTX2ConditionDenoiseStep` except that the denoiser carries no guider, " + "since the distilled DFR schedules run without guidance. Iterates " + "`LTX2DenoiseLoopWrapper.__call__`, running per step:\n" + " - `LTX2ConditionLoopBeforeDenoiser`\n - `LTX2DFRLoopDenoiser`\n - `LTX2ConditionLoopAfterDenoiser`" + ) + + class LTX2DenoiseStep(LTX2DenoiseLoopWrapper): block_classes = [LTX2LoopBeforeDenoiser, LTX2LoopDenoiser, LTX2LoopAfterDenoiser] block_names = ["before_denoiser", "denoiser", "after_denoiser"] diff --git a/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2_dfr.py b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2_dfr.py new file mode 100644 index 000000000000..777197a047b2 --- /dev/null +++ b/src/diffusers/modular_pipelines/ltx2/modular_blocks_ltx2_dfr.py @@ -0,0 +1,454 @@ +# Copyright 2026 The HuggingFace Team. All rights reserved. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +import torch + +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import OutputParam +from .before_denoise import ( + LTX2ConditionPrepareAudioLatentsStep, + LTX2ConditionPrepareCoordsStep, + LTX2ConditionSetTimestepsStep, + LTX2DFRPlanStep, + LTX2DFRPrepareLatentsStep, + LTX2TextInputStep, +) +from .decoders import LTX2AudioDecoderStep, LTX2DFRSplitKeyframesStep, LTX2VaeDecoderStep +from .denoise import LTX2DFRDenoiseStep +from .modular_blocks_ltx2 import ( + LTX2AutoConditionEncoderStep, + LTX2AutoDurationStep, + LTX2AutoPromptEnhancerStep, + LTX2TextConditioningStep, +) + + +# auto_docstring +class LTX2DFRCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise stage for one DFR pass. Identical to `LTX2ConditionCoreDenoiseStep` except for the prepare-latents + block, which appends the generated keyframe slots and the optional spatial detailing reference. Everything + downstream is unchanged: the slot marker rides to the transformer as a `denoiser_input_fields` output, and the slot + coordinates ride in `appended_coords`. + + Components: + transformer (`LTX2VideoTransformer3DModel`) vae (`AutoencoderKLLTX2Video`) scheduler + (`FlowMatchEulerDiscreteScheduler`) audio_vae (`AutoencoderKLLTX2Audio`) + + Inputs: + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond). + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond). + connector_attention_mask (`Tensor`): + Binary text attention mask (cond). + negative_connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (uncond). + negative_connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (uncond). + negative_connector_attention_mask (`Tensor`): + Binary text attention mask (uncond). + slot_frame_indices (`list`): + Pixel-frame positions of the generated keyframe slots, from `LTX2DFRPlanStep`. + keyframes_latents (`Tensor`, *optional*): + `[B, C, num_slots, H, W]` content seeding the keyframe slots: a previous DFR pass's `keyframes_latents` + upsampled to this pass's resolution. Denormalized, like every latent crossing the pipeline boundary. + Slots start from noise when omitted. + detailing_reference_latents (`Tensor`, *optional*): + `[B, C, F, H, W]` latents appended as a fully clean in-context reference for the spatial detailing + IC-LoRA: the previous pass's output at its own resolution, denormalized. Only meaningful with that + adapter loaded. + detailing_reference_downscale_factor (`int`, *optional*, defaults to 2): + Ratio between this pass's resolution and the reference's, used to scale the reference tokens' spatial + coordinates into the target coordinate space. Must match the factor the IC-LoRA was trained with. + condition_latents (`list`, *optional*): + Per-condition normalized VAE latents of shape [1, C, F, H, W]. + condition_strengths (`list`, *optional*): + Per-condition conditioning strengths. + condition_indices (`list`, *optional*): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`, *optional*): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_frames (`int`): + The padded canvas frame count, from `LTX2DFRPlanStep`. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + noise_scale (`float`, *optional*): + Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom + `sigmas` are supplied, else 1.0. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + batch_size (`int`): + The number of prompts being denoised, used to expand conditioning per prompt. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + num_inference_steps (`int`, *optional*, defaults to 30): + The number of denoising steps. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + audio_latents (`Tensor`, *optional*): + Optional pre-encoded audio latents; random noise is used when not provided. + dtype (`dtype`): + The dtype the model inputs are cast to. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + use_cross_timestep (`bool`, *optional*, defaults to True): + Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + + Outputs: + connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (cond), expanded per prompt. + connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (cond), expanded per prompt. + connector_attention_mask (`Tensor`): + Binary text attention mask (cond), expanded per prompt. + negative_connector_prompt_embeds (`Tensor`): + Video-branch text conditioning (uncond), expanded per prompt. + negative_connector_audio_prompt_embeds (`Tensor`): + Audio-branch text conditioning (uncond), expanded per prompt. + negative_connector_attention_mask (`Tensor`): + Binary text attention mask (uncond), expanded per prompt. + latents (`Tensor`): + Packed noisy video latents, with keyframe, slot and reference tokens appended. + conditioning_mask (`Tensor`): + Packed per-token conditioning strengths of shape [B, S, 1] in [0, 1]: 1 at fully-conditioned positions, 0 + at free positions, including every keyframe slot. + clean_latents (`Tensor`): + Clean condition latents at conditioned positions, zeros elsewhere; same shape as `latents`. + video_keyframes_mask (`Tensor`): + Packed [B, S, 1] marker, 1 on tokens whose latent frame encodes a single pixel frame -- the causal first + frame and every generated keyframe slot. Those tokens receive the transformer's + `keyframes_abs_pos_embedding`. + appended_coords (`Tensor`): + RoPE coordinates of shape [B, 3, num_appended_tokens, 2] for the appended keyframe, slot and reference + tokens, in the order they were appended. + base_token_count (`int`): + Number of generated-video tokens, i.e. the sequence length before appended tokens. + slot_token_slice (`slice`): + Slice of the packed sequence holding the generated keyframe slot tokens. + noise_scale (`float`): + The resolved initial noise level, forwarded to the audio latents step. + timesteps (`Tensor`): + TODO: Add description. + num_inference_steps (`int`): + TODO: Add description. + audio_scheduler (`None`): + Independent deep copy of `scheduler` used to update the audio latents in the loop. + audio_latents (`Tensor`): + Packed noisy audio latents. + audio_num_frames (`int`): + Number of audio latent frames. + video_coords (`Tensor`): + Video RoPE patch coordinates, with the keyframe-condition coordinates appended. + audio_coords (`Tensor`): + Audio RoPE patch coordinates. + noise_pred_video (`Tensor`): + Video x0 prediction for this step. + noise_pred_audio (`Tensor`): + Audio x0 prediction for this step. + """ + + model_name = "ltx2.5-dfr" + block_classes = [ + LTX2TextInputStep, + LTX2DFRPrepareLatentsStep, + LTX2ConditionSetTimestepsStep, + LTX2ConditionPrepareAudioLatentsStep, + LTX2ConditionPrepareCoordsStep, + LTX2DFRDenoiseStep, + ] + block_names = [ + "input", + "prepare_latents", + "set_timesteps", + "prepare_audio_latents", + "prepare_coords", + "denoise", + ] + + @property + def description(self): + return ( + "Core denoise stage for one DFR pass. Identical to `LTX2ConditionCoreDenoiseStep` except for the " + "prepare-latents block, which appends the generated keyframe slots and the optional spatial detailing " + "reference. Everything downstream is unchanged: the slot marker rides to the transformer as a " + "`denoiser_input_fields` output, and the slot coordinates ride in `appended_coords`." + ) + + +# auto_docstring +class LTX2DFRDecoderStep(SequentialPipelineBlocks): + """ + Decode stage for DFR: splits the generated keyframe slots out of the denoised sequence and trims the canvas + padding, then decodes the video latents with the convolutional VAE and vocodes the audio latents (or returns + latents). The convolutional VAE is what the reference DFR implementation decodes with. For maximum detail fidelity, + run with `output_type="latent"` and hand the latents to [`LTX2VideoDiffusionDecodePipeline`] instead. + + Components: + vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`) + vocoder (`LTX2Vocoder`) + + Inputs: + latents (`Tensor`): + Pre-generated noisy latents for image generation. + base_token_count (`int`): + Number of generated-video tokens, i.e. the sequence length before appended tokens. + slot_token_slice (`slice`): + Slice of the packed sequence holding the generated keyframe slot tokens. + requested_num_frames (`int`): + The frame count the caller asked for, which the canvas is trimmed back to. + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + decode_timestep (`None`, *optional*, defaults to 0.0): + The timestep at which the VAE decodes the final latents. + decode_noise_scale (`None`, *optional*): + Noise interpolation factor applied to the latents at the decode timestep. + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + batch_size (`int`, *optional*, defaults to 1): + Number of prompts, the final batch size of model inputs should be batch_size * num_images_per_prompt. Can + be generated in input step. + dtype (`dtype`): + The dtype of the model inputs, can be generated in input step. + audio_latents (`Tensor`): + Denoised audio latents. + audio_num_frames (`int`): + Number of audio latent frames, used to unpack the audio latent sequence. + + Outputs: + videos (`list`): + The generated videos. + audio (`Tensor`): + The generated audio waveform. + keyframes_latents (`Tensor`): + Denormalized `[B, C, num_slots, H, W]` generated keyframe slots. Upsample these alongside the video + latents to seed the detailing pass. + """ + + model_name = "ltx2.5-dfr" + block_classes = [LTX2DFRSplitKeyframesStep, LTX2VaeDecoderStep, LTX2AudioDecoderStep] + block_names = ["split_keyframes", "video_decode", "audio_decode"] + + @property + def description(self): + return ( + "Decode stage for DFR: splits the generated keyframe slots out of the denoised sequence and trims the " + "canvas padding, then decodes the video latents with the convolutional VAE and vocodes the audio " + "latents (or returns latents). The convolutional VAE is what the reference DFR implementation decodes " + 'with. For maximum detail fidelity, run with `output_type="latent"` and hand the latents to ' + "[`LTX2VideoDiffusionDecodePipeline`] instead." + ) + + @property + def outputs(self): + return [ + OutputParam.template("videos"), + OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform."), + OutputParam( + "keyframes_latents", + type_hint=torch.Tensor, + description=( + "Denormalized `[B, C, num_slots, H, W]` generated keyframe slots. Upsample these alongside the " + "video latents to seed the detailing pass." + ), + ), + ] + + +# auto_docstring +class LTX2DFRBlocks(SequentialPipelineBlocks): + """ + Modular pipeline blocks for one LTX-2.5 Diffusion Fidelity Rendering (DFR) pass (joint video + audio). + DFR generates on a canvas padded to a whole number of keyframe segments and spends one extra latent frame of + tokens per segment border on a *keyframe slot*: a single-pixel-frame latent the model fills in. Relaxing the + effective temporal compression at those positions means the surrounding video is conditioned on genuinely new + frames rather than interpolated ones. The full recipe is two passes of these blocks. The first runs at half the + target resolution and returns `videos`/`keyframes_latents` as latents; upsample both with + `LTX2LatentUpsamplePipeline`, load the spatial detailing IC-LoRA, and run the blocks again at full resolution + with `latents`, `keyframes_latents` and `detailing_reference_latents` supplied. Needs a transformer whose config + sets `use_keyframes_abs_pos_embedding`, which LTX-2.5 checkpoints ship. + + Supported workflows: + - `text2video`: requires `prompt` + - `condition`: requires `conditions`, `prompt` + + Components: + prompt_enhancer (`PreTrainedModel`) processor (`ProcessorMixin`) text_encoder (`PreTrainedModel`) tokenizer + (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) duration_head (`LTX2DurationHead`) transformer + (`LTX2VideoTransformer3DModel`) vae (`AutoencoderKLLTX2Video`) scheduler (`FlowMatchEulerDiscreteScheduler`) + audio_vae (`AutoencoderKLLTX2Audio`) video_processor (`VideoProcessor`) vocoder (`LTX2Vocoder`) + + Inputs: + prompt (`str`, *optional*): + The prompt or prompts to guide image generation. + conditions (`list`, *optional*): + `LTX2VideoCondition` (or list of them) placing image/video conditions at latent frame indices of the + generated video. + enable_prompt_enhancement (`bool`, *optional*, defaults to False): + Whether to run the prompt enhancer. Opt-in, matching the Lightricks reference pipelines. + system_prompt (`str`, *optional*): + System prompt for enhancement. Defaults to `LTX2_5_I2V_DEFAULT_SYSTEM_PROMPT` when a `PIL.Image.Image` + condition frame is available, else `LTX2_5_T2V_DEFAULT_SYSTEM_PROMPT`. + prompt_max_new_tokens (`int`, *optional*): + Maximum number of new tokens to generate during prompt enhancement. Defaults to 600, the LTX-2.5 Gemma-4 + enhancer's budget. + prompt_enhancement_kwargs (`dict`, *optional*): + Keyword arguments for the enhancer's `.generate` call. Defaults to greedy decoding. + prompt_enhancement_seed (`int`, *optional*, defaults to 10): + Random seed for prompt enhancement (inert under LTX-2.5's greedy decoding). + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + image (`Image | list`, *optional*): + Reference image(s) for denoising. Can be a single image or list of images. + negative_prompt (`str`, *optional*): + The prompt or prompts not to guide the image generation. + max_sequence_length (`int`, *optional*, defaults to 1024): + Maximum sequence length for prompt encoding. + min_seconds (`float`, *optional*, defaults to 1.0): + Lower bound on the auto-predicted duration. + max_seconds (`float`, *optional*, defaults to 20.0): + Upper bound on the auto-predicted duration. Must be strictly greater than `min_seconds`. + frame_rate (`float`, *optional*, defaults to 24.0): + Frames per second of the generated video. + num_frames (`int`, *optional*): + The number of frames the caller asked for, before the canvas is padded onto the segment grid. Omit to + auto-predict via the `duration_head` (see `LTX2AutoDurationStep`). + height (`int`, *optional*, defaults to 512): + The height in pixels of the generated image. + width (`int`, *optional*, defaults to 704): + The width in pixels of the generated image. + num_videos_per_prompt (`int`, *optional*, defaults to 1): + The number of images to generate per prompt. + keyframes_latents (`Tensor`, *optional*): + `[B, C, num_slots, H, W]` content seeding the keyframe slots: a previous DFR pass's `keyframes_latents` + upsampled to this pass's resolution. Denormalized, like every latent crossing the pipeline boundary. + Slots start from noise when omitted. + detailing_reference_latents (`Tensor`, *optional*): + `[B, C, F, H, W]` latents appended as a fully clean in-context reference for the spatial detailing + IC-LoRA: the previous pass's output at its own resolution, denormalized. Only meaningful with that + adapter loaded. + detailing_reference_downscale_factor (`int`, *optional*, defaults to 2): + Ratio between this pass's resolution and the reference's, used to scale the reference tokens' spatial + coordinates into the target coordinate space. Must match the factor the IC-LoRA was trained with. + condition_latents (`list`, *optional*): + Per-condition normalized VAE latents of shape [1, C, F, H, W]. + condition_strengths (`list`, *optional*): + Per-condition conditioning strengths. + condition_indices (`list`, *optional*): + Per-condition latent frame index at which the condition is applied. + condition_pixel_frames (`list`, *optional*): + Per-condition trimmed pixel frame count, used to clamp single-frame keyframe coords. + latents (`Tensor`, *optional*): + Pre-generated noisy latents for image generation. + noise_scale (`float`, *optional*): + Initial noise level for the un-conditioned tokens. `None` (default) resolves to `sigmas[0]` when custom + `sigmas` are supplied, else 1.0. + sigmas (`list`, *optional*): + Custom sigmas for the denoising process. + num_inference_steps (`int`, *optional*, defaults to 30): + The number of denoising steps. + timesteps (`Tensor`, *optional*): + Timesteps for the denoising process. + audio_latents (`Tensor`, *optional*): + Optional pre-encoded audio latents; random noise is used when not provided. + **denoiser_input_fields (`None`, *optional*): + conditional model inputs for the denoiser: e.g. prompt_embeds, negative_prompt_embeds, etc. + use_cross_timestep (`bool`, *optional*, defaults to True): + Whether to condition the transformer on a separate per-token cross timestep (LTX-2.3+). + attention_kwargs (`dict`, *optional*): + Additional kwargs for attention processors. + output_type (`str`, *optional*, defaults to pil): + Output format: 'pil', 'np', 'pt'. + decode_timestep (`None`, *optional*, defaults to 0.0): + The timestep at which the VAE decodes the final latents. + decode_noise_scale (`None`, *optional*): + Noise interpolation factor applied to the latents at the decode timestep. + + Outputs: + videos (`list`): + The generated videos. + audio (`Tensor`): + The generated audio waveform. + keyframes_latents (`Tensor`): + Denormalized `[B, C, num_slots, H, W]` generated keyframe slots. Upsample these alongside the video + latents to seed the detailing pass. + """ + + model_name = "ltx2.5-dfr" + block_classes = [ + LTX2AutoPromptEnhancerStep, + LTX2TextConditioningStep, + LTX2AutoDurationStep, + LTX2DFRPlanStep, + LTX2AutoConditionEncoderStep, + LTX2DFRCoreDenoiseStep, + LTX2DFRDecoderStep, + ] + block_names = [ + "prompt_enhancer", + "text_encoder", + "duration", + "plan", + "condition_encoder", + "denoise", + "decode", + ] + _workflow_map = { + "text2video": {"prompt": True}, + "condition": {"conditions": True, "prompt": True}, + } + + @property + def description(self): + return ( + "Modular pipeline blocks for one LTX-2.5 Diffusion Fidelity Rendering (DFR) pass (joint video + audio).\n" + "DFR generates on a canvas padded to a whole number of keyframe segments and spends one extra latent " + "frame of tokens per segment border on a *keyframe slot*: a single-pixel-frame latent the model fills " + "in. Relaxing the effective temporal compression at those positions means the surrounding video is " + "conditioned on genuinely new frames rather than interpolated ones.\n" + "The full recipe is two passes of these blocks. The first runs at half the target resolution and " + "returns `videos`/`keyframes_latents` as latents; upsample both with `LTX2LatentUpsamplePipeline`, load " + "the spatial detailing IC-LoRA, and run the blocks again at full resolution with `latents`, " + "`keyframes_latents` and `detailing_reference_latents` supplied. Needs a transformer whose config sets " + "`use_keyframes_abs_pos_embedding`, which LTX-2.5 checkpoints ship." + ) + + @property + def outputs(self): + return [ + OutputParam.template("videos"), + OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform."), + OutputParam( + "keyframes_latents", + type_hint=torch.Tensor, + description=( + "Denormalized `[B, C, num_slots, H, W]` generated keyframe slots. Upsample these alongside the " + "video latents to seed the detailing pass." + ), + ), + ] diff --git a/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py b/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py index c2e3409bbf37..7554cd409d37 100644 --- a/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/ltx2/modular_pipeline.py @@ -132,3 +132,16 @@ class LTX25ModularPipeline(LTX2ModularPipeline): """ default_blocks_name = "LTX25AutoBlocks" + + +class LTX2DFRModularPipeline(LTX2ModularPipeline): + """ + A ModularPipeline for one LTX-2.5 Diffusion Fidelity Rendering (DFR) pass (joint video + audio). + + Generates on a canvas padded to a whole number of keyframe segments, spending one extra latent frame of tokens per + segment border on a generated single-pixel-frame keyframe slot. The full recipe runs these blocks twice -- see + [`LTX2DFRBlocks`]. A checkpoint routes here through `modular_model_index.json`. + + """ + + default_blocks_name = "LTX2DFRBlocks" diff --git a/src/diffusers/modular_pipelines/ltx2/utils.py b/src/diffusers/modular_pipelines/ltx2/utils.py index 814af7ccdff4..1f933d9b8f12 100644 --- a/src/diffusers/modular_pipelines/ltx2/utils.py +++ b/src/diffusers/modular_pipelines/ltx2/utils.py @@ -96,3 +96,57 @@ 1.3125, 1.289062, 1.296875, 1.242188, 1.234375, 1.21875, 1.226562, 1.054688, ] # fmt: on + + +# 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. +DFR_SEGMENT_LATENT_CANDIDATES = (3, 4) + + +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 `DFR_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 DFR_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 slot + positions `[S, 2S, ..., N' - 1]` in pixel frames. 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 diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 3aa2c854dfe6..0db929ca56c3 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -154,6 +154,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("ltx", _create_default_map_fn("LTXModularPipeline")), ("ltx2", _create_default_map_fn("LTX2ModularPipeline")), ("ltx2.5", _create_default_map_fn("LTX25ModularPipeline")), + ("ltx2.5-dfr", _create_default_map_fn("LTX2DFRModularPipeline")), ("minimax-h3", _create_default_map_fn("MiniMaxH3ModularPipeline")), ("minimax-music3", _create_default_map_fn("MiniMaxMusic3ModularPipeline")), ("ernie-image", _create_default_map_fn("ErnieImageModularPipeline")), diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 376596d632ea..d056a2115f9d 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -497,6 +497,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +class LTX2DFRBlocks(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 LTX2DFRModularPipeline(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 LTX2ModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] diff --git a/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2_dfr.py b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2_dfr.py new file mode 100644 index 000000000000..75c037c1d18a --- /dev/null +++ b/tests/modular_pipelines/ltx2/test_modular_pipeline_ltx2_dfr.py @@ -0,0 +1,257 @@ +# coding=utf-8 +# Copyright 2026 HuggingFace Inc. +# +# 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 +import torch.nn.functional as F + +from diffusers.modular_pipelines import LTX2DFRBlocks, LTX2DFRModularPipeline +from diffusers.modular_pipelines.ltx2.utils import resolve_canvas +from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition + +from ..testing_utils import ( + BaseModularPipelineTesterConfig, + ModularLoadingTesterMixin, + ModularMemoryTesterMixin, + ModularPipelineTesterMixin, + ModularWorkflowTesterMixin, +) + + +# Differs from `hf-internal-testing/tiny-ltx2-5-modular-pipe` in one place: its transformer sets +# `use_keyframes_abs_pos_embedding`, which DFR requires and the LTX-2.5 fixture leaves off. +LTX2DFR_REPO_ID = "hf-internal-testing/tiny-ltx2-5-dfr-modular-pipe" + +LTX2DFR_WORKFLOWS = { + "text2video": [ + ("text_encoder.text_encoder", "LTX2TextEncoderStep"), + ("text_encoder.connectors", "LTX2TextConnectorStep"), + ("duration", "LTX2DurationStep"), + ("plan", "LTX2DFRPlanStep"), + ("denoise.input", "LTX2TextInputStep"), + ("denoise.prepare_latents", "LTX2DFRPrepareLatentsStep"), + ("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"), + ("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"), + ("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"), + ("denoise.denoise", "LTX2DFRDenoiseStep"), + ("decode.split_keyframes", "LTX2DFRSplitKeyframesStep"), + ("decode.video_decode", "LTX2VaeDecoderStep"), + ("decode.audio_decode", "LTX2AudioDecoderStep"), + ], + "condition": [ + ("text_encoder.text_encoder", "LTX2TextEncoderStep"), + ("text_encoder.connectors", "LTX2TextConnectorStep"), + ("duration", "LTX2DurationStep"), + ("plan", "LTX2DFRPlanStep"), + ("condition_encoder", "LTX2ConditionEncoderStep"), + ("denoise.input", "LTX2TextInputStep"), + ("denoise.prepare_latents", "LTX2DFRPrepareLatentsStep"), + ("denoise.set_timesteps", "LTX2ConditionSetTimestepsStep"), + ("denoise.prepare_audio_latents", "LTX2ConditionPrepareAudioLatentsStep"), + ("denoise.prepare_coords", "LTX2ConditionPrepareCoordsStep"), + ("denoise.denoise", "LTX2DFRDenoiseStep"), + ("decode.split_keyframes", "LTX2DFRSplitKeyframesStep"), + ("decode.video_decode", "LTX2VaeDecoderStep"), + ("decode.audio_decode", "LTX2AudioDecoderStep"), + ], +} + +# The fixture VAE compresses time by 2, so the segment grid offers 6 and 8 pixel frames. 13 frames land on the +# 6-frame grid unpadded and buy two keyframe slots, at pixel frames 6 and 12. +DUMMY_NUM_FRAMES = 13 +DUMMY_NUM_SLOTS = 2 + + +class LTX2DFRModularPipelineTesterConfig(BaseModularPipelineTesterConfig): + """Shared configuration for the DFR workflows; a variant config adds its own `params` and dummy inputs.""" + + pipeline_class = LTX2DFRModularPipeline + pipeline_blocks_class = LTX2DFRBlocks + pretrained_model_name_or_path = LTX2DFR_REPO_ID + batch_params = frozenset(["prompt"]) + optional_params = frozenset(["num_inference_steps", "num_videos_per_prompt", "latents"]) + expected_workflow_blocks = LTX2DFR_WORKFLOWS + output_name = "videos" + + def get_dummy_inputs(self, seed=0): + return { + "prompt": "a robot dancing", + "negative_prompt": "", + "generator": self.get_generator(seed), + "num_inference_steps": 2, + "height": 32, + "width": 32, + "num_frames": DUMMY_NUM_FRAMES, + "frame_rate": 25.0, + "max_sequence_length": 16, + "output_type": "pt", + } + + +class LTX2DFRModularPipelineFastTesterMixin(ModularPipelineTesterMixin): + """`ModularPipelineTesterMixin` with the two adjustments every LTX-2 workflow needs.""" + + @pytest.mark.skip(reason="num_videos_per_prompt") + def test_num_images_per_prompt(self): + pass + + def test_inference_batch_single_identical(self): + super().test_inference_batch_single_identical(expected_max_diff=1e-3) + + +class LTX2DFRText2VideoModularPipelineTesterConfig(LTX2DFRModularPipelineTesterConfig): + params = frozenset(["prompt", "height", "width", "num_frames"]) + + +class TestLTX2DFRText2VideoModularPipelineFast( + LTX2DFRText2VideoModularPipelineTesterConfig, LTX2DFRModularPipelineFastTesterMixin +): + def test_generates_a_keyframe_slot_per_segment_border(self): + pipe = self.get_pipeline().to("cpu") + + inputs = self.get_dummy_inputs() + output = pipe(**inputs, output=["videos", "keyframes_latents"]) + + assert output["videos"].shape == (1, DUMMY_NUM_FRAMES, 3, 32, 32) + # One latent frame of content per slot, at the target's latent resolution. + keyframes = output["keyframes_latents"] + assert keyframes.shape[2] == DUMMY_NUM_SLOTS + assert keyframes.shape[3:] == (32 // pipe.vae_spatial_compression_ratio,) * 2 + assert torch.isnan(keyframes).sum() == 0 + + def test_canvas_padding_is_trimmed_back_to_the_requested_length(self): + pipe = self.get_pipeline().to("cpu") + ratio = pipe.vae_temporal_compression_ratio + + # 11 frames divide neither segment length, so the canvas pads to 13 and is trimmed back. + requested = 11 + inputs = self.get_dummy_inputs() + inputs["num_frames"] = requested + canvas_frames, _, _ = resolve_canvas(requested, ratio) + assert canvas_frames > requested, "pick a frame count that actually pads, or this asserts nothing" + + videos = pipe(**inputs, output="videos") + assert videos.shape[1] == requested + + def test_keyframe_marker_reaches_the_transformer(self): + # `video_keyframes_mask` is plumbed to the denoiser purely by its `denoiser_input_fields` tag, with no + # DFR-specific denoise block. Zeroing the learned embedding it gates is the behavioural check that the + # tag actually arrives: if it does not, the two runs are identical. + pipe = self.get_pipeline().to("cpu") + + with_marker = pipe(**self.get_dummy_inputs(), output="videos") + with torch.no_grad(): + pipe.transformer.keyframes_abs_pos_embedding.zero_() + without_marker = pipe(**self.get_dummy_inputs(), output="videos") + + assert not torch.allclose(with_marker, without_marker) + + def test_detailing_pass_consumes_the_first_pass_output(self): + # The full DFR recipe is two passes of these blocks, with the spatial detailing IC-LoRA loaded in + # between. Here only the plumbing is exercised: the second pass has to accept the first pass's video + # latents, its keyframe slots and an in-context reference, and return the same shapes one scale up. + pipe = self.get_pipeline().to("cpu") + + base_inputs = self.get_dummy_inputs() + base_inputs["output_type"] = "latent" + first = pipe(**base_inputs, output=["videos", "keyframes_latents"]) + video_latents, keyframes_latents = first["videos"], first["keyframes_latents"] + + # Stands in for `LTX2LatentUpsamplePipeline`, which is a separate pipeline in the real recipe. + upsample = lambda latents: F.interpolate(latents, scale_factor=(1, 2, 2), mode="nearest") # noqa: E731 + + detail_inputs = self.get_dummy_inputs() + detail_inputs.update( + output_type="latent", + height=64, + width=64, + latents=upsample(video_latents), + keyframes_latents=upsample(keyframes_latents), + detailing_reference_latents=video_latents, + detailing_reference_downscale_factor=2, + noise_scale=0.4, + ) + second = pipe(**detail_inputs, output=["videos", "keyframes_latents"]) + + assert second["videos"].shape[-2:] == (2 * video_latents.shape[-2], 2 * video_latents.shape[-1]) + assert second["keyframes_latents"].shape[2] == DUMMY_NUM_SLOTS + assert torch.isnan(second["videos"]).sum() == 0 + + def test_auto_duration_lands_on_the_segment_grid(self): + # The plan step pads a duration-head prediction onto the segment grid, so the two have to agree: the + # predicted length must itself be a length `resolve_canvas` accepts. + pipe = self.get_pipeline().to("cpu") + + inputs = self.get_dummy_inputs() + inputs.pop("num_frames") + inputs["min_seconds"] = 0.5 + inputs["max_seconds"] = 2.0 + videos = pipe(**inputs, output="videos") + + num_frames = videos.shape[1] + assert (num_frames - 1) % pipe.vae_temporal_compression_ratio == 0 + assert 0 < num_frames <= round(2.0 * inputs["frame_rate"]) + + def test_rejects_a_frame_count_off_the_latent_grid(self): + pipe = self.get_pipeline().to("cpu") + + inputs = self.get_dummy_inputs() + inputs["num_frames"] = DUMMY_NUM_FRAMES + 1 + with pytest.raises(ValueError, match="num_frames"): + pipe(**inputs, output="videos") + + +class TestLTX2DFRText2VideoModularPipelineLoading( + LTX2DFRText2VideoModularPipelineTesterConfig, ModularLoadingTesterMixin +): + pass + + +class TestLTX2DFRText2VideoModularPipelineMemory( + LTX2DFRText2VideoModularPipelineTesterConfig, ModularMemoryTesterMixin +): + pass + + +# Both workflows share `LTX2DFRBlocks` and the same repo, so one workflow test class covers them. +class TestLTX2DFRModularPipelineWorkflow(LTX2DFRText2VideoModularPipelineTesterConfig, ModularWorkflowTesterMixin): + pass + + +class LTX2DFRConditionModularPipelineTesterConfig(LTX2DFRModularPipelineTesterConfig): + params = frozenset(["prompt", "conditions", "height", "width", "num_frames"]) + + def get_dummy_inputs(self, seed=0): + inputs = super().get_dummy_inputs(seed) + image = torch.rand((1, 3, 32, 32), generator=torch.Generator("cpu").manual_seed(seed)) + # Synthetic float tensors skip H.264 CRF re-compression (training path uses PIL/uint8). + inputs["conditions"] = LTX2VideoCondition(frames=image, index=0, strength=1.0, crf=0) + return inputs + + +class TestLTX2DFRConditionModularPipelineFast( + LTX2DFRConditionModularPipelineTesterConfig, LTX2DFRModularPipelineFastTesterMixin +): + pass + + +class TestLTX2DFRConditionModularPipelineLoading( + LTX2DFRConditionModularPipelineTesterConfig, ModularLoadingTesterMixin +): + pass + + +class TestLTX2DFRConditionModularPipelineMemory(LTX2DFRConditionModularPipelineTesterConfig, ModularMemoryTesterMixin): + pass