# Copyright 2025 Lightricks and The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import copy import inspect from dataclasses import dataclass from typing import Any, Callable import numpy as np import PIL.Image import torch from transformers import Gemma3ForConditionalGeneration, 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 FlowMatchEulerDiscreteScheduler from ...utils import is_torch_xla_available, logging, replace_example_docstring from ...utils.torch_utils import randn_tensor from ...video_processor import VideoProcessor from ..pipeline_utils import DiffusionPipeline from .connectors import LTX2TextConnectors from .pipeline_output import LTX2PipelineOutput from .vocoder import LTX2Vocoder 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 EXAMPLE_DOC_STRING = """ Examples: ```py >>> import torch >>> from diffusers import LTX2ConditionPipeline >>> from diffusers.pipelines.ltx2.export_utils import encode_video >>> from diffusers.pipelines.ltx2.pipeline_ltx2_condition import LTX2VideoCondition >>> from diffusers.utils import load_image >>> pipe = LTX2ConditionPipeline.from_pretrained("Lightricks/LTX-2", torch_dtype=torch.bfloat16) >>> pipe.enable_model_cpu_offload() >>> first_image = load_image( ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png" ... ) >>> last_image = load_image( ... "https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png" ... ) >>> first_cond = LTX2VideoCondition(frames=first_image, index=0, strength=1.0) >>> last_cond = LTX2VideoCondition(frames=last_image, index=-1, strength=1.0) >>> conditions = [first_cond, last_cond] >>> prompt = "CG animation style, a small blue bird takes off from the ground, flapping its wings." >>> negative_prompt = "worst quality, inconsistent motion, blurry, jittery, distorted, static" >>> frame_rate = 24.0 >>> video = pipe( ... conditions=conditions, ... prompt=prompt, ... negative_prompt=negative_prompt, ... width=768, ... height=512, ... num_frames=121, ... frame_rate=frame_rate, ... num_inference_steps=40, ... guidance_scale=4.0, ... output_type="np", ... return_dict=False, ... ) >>> video = (video * 255).round().astype("uint8") >>> video = torch.from_numpy(video) >>> encode_video( ... video[0], ... fps=frame_rate, ... audio=audio[0].float().cpu(), ... audio_sample_rate=pipe.vocoder.config.output_sampling_rate, # should be 24000 ... output_path="video.mp4", ... ) ``` """ @dataclass class LTX2VideoCondition: """ Defines a single frame-conditioning item for LTX-2 Video - a single frame or a sequence of frames. Attributes: frames (`PIL.Image.Image` or `List[PIL.Image.Image]` or `np.ndarray` or `torch.Tensor`): The image (or video) to condition the video on. Accepts any type that can be handled by VideoProcessor.preprocess_video. index (`int`, defaults to `0`): The index at which the image or video will conditionally affect the video generation. strength (`float`, defaults to `1.0`): The strength of the conditioning effect. A value of `1.0` means the conditioning effect is fully applied. """ frames: PIL.Image.Image | list[PIL.Image.Image] | np.ndarray | torch.Tensor index: int = 0 strength: float = 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") # Copied from diffusers.pipelines.flux.pipeline_flux.calculate_shift def calculate_shift( image_seq_len, base_seq_len: int = 256, max_seq_len: int = 4096, base_shift: float = 0.5, max_shift: float = 1.15, ): m = (max_shift - base_shift) / (max_seq_len - base_seq_len) b = base_shift - m * base_seq_len mu = image_seq_len * m + b return mu # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.retrieve_timesteps def retrieve_timesteps( scheduler, num_inference_steps: int | None = None, device: str | torch.device | None = None, timesteps: list[int] | None = None, sigmas: list[float] | None = None, **kwargs, ): r""" Calls the scheduler's `set_timesteps` method and retrieves timesteps from the scheduler after the call. Handles custom timesteps. Any kwargs will be supplied to `scheduler.set_timesteps`. Args: scheduler (`SchedulerMixin`): The scheduler to get timesteps from. num_inference_steps (`int`): The number of diffusion steps used when generating samples with a pre-trained model. If used, `timesteps` must be `None`. device (`str` or `torch.device`, *optional*): The device to which the timesteps should be moved to. If `None`, the timesteps are not moved. timesteps (`list[int]`, *optional*): Custom timesteps used to override the timestep spacing strategy of the scheduler. If `timesteps` is passed, `num_inference_steps` and `sigmas` must be `None`. sigmas (`list[float]`, *optional*): Custom sigmas used to override the timestep spacing strategy of the scheduler. If `sigmas` is passed, `num_inference_steps` and `timesteps` must be `None`. Returns: `tuple[torch.Tensor, int]`: A tuple where the first element is the timestep schedule from the scheduler and the second element is the number of inference steps. """ if timesteps is not None and sigmas is not None: raise ValueError("Only one of `timesteps` or `sigmas` can be passed. Please choose one to set custom values") if timesteps is not None: accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) if not accepts_timesteps: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" timestep schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) elif sigmas is not None: accept_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) if not accept_sigmas: raise ValueError( f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" f" sigmas schedules. Please check whether you are using the correct scheduler." ) scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) timesteps = scheduler.timesteps num_inference_steps = len(timesteps) else: scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) timesteps = scheduler.timesteps return timesteps, num_inference_steps # Copied from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion.rescale_noise_cfg def rescale_noise_cfg(noise_cfg, noise_pred_text, guidance_rescale=0.0): r""" Rescales `noise_cfg` tensor based on `guidance_rescale` to improve image quality and fix overexposure. Based on Section 3.4 from [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Args: noise_cfg (`torch.Tensor`): The predicted noise tensor for the guided diffusion process. noise_pred_text (`torch.Tensor`): The predicted noise tensor for the text-guided diffusion process. guidance_rescale (`float`, *optional*, defaults to 0.0): A rescale factor applied to the noise predictions. Returns: noise_cfg (`torch.Tensor`): The rescaled noise prediction tensor. """ std_text = noise_pred_text.std(dim=list(range(1, noise_pred_text.ndim)), keepdim=True) std_cfg = noise_cfg.std(dim=list(range(1, noise_cfg.ndim)), keepdim=True) # rescale the results from guidance (fixes overexposure) noise_pred_rescaled = noise_cfg * (std_text / std_cfg) # mix with the original results from guidance by factor guidance_rescale to avoid "plain looking" images noise_cfg = guidance_rescale * noise_pred_rescaled + (1 - guidance_rescale) * noise_cfg return noise_cfg class LTX2ConditionPipeline(DiffusionPipeline, FromSingleFileMixin, LTX2LoraLoaderMixin): r""" Pipeline for video generation which allows image conditions to be inserted at arbitary parts of the video. Reference: https://github.com/Lightricks/LTX-Video TODO """ model_cpu_offload_seq = "text_encoder->connectors->transformer->vae->audio_vae->vocoder" _optional_components = [] _callback_tensor_inputs = ["latents", "prompt_embeds", "negative_prompt_embeds"] def __init__( self, scheduler: FlowMatchEulerDiscreteScheduler, vae: AutoencoderKLLTX2Video, audio_vae: AutoencoderKLLTX2Audio, text_encoder: Gemma3ForConditionalGeneration, tokenizer: GemmaTokenizer | GemmaTokenizerFast, connectors: LTX2TextConnectors, transformer: LTX2VideoTransformer3DModel, vocoder: LTX2Vocoder, ): 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, ) 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") 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.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 ) @staticmethod # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._pack_text_embeds def _pack_text_embeds( text_hidden_states: torch.Tensor, sequence_lengths: torch.Tensor, device: str | torch.device, padding_side: str = "left", scale_factor: int = 8, eps: float = 1e-6, ) -> torch.Tensor: """ Packs and normalizes text encoder hidden states, respecting padding. Normalization is performed per-batch and per-layer in a masked fashion (only over non-padded positions). Args: text_hidden_states (`torch.Tensor` of shape `(batch_size, seq_len, hidden_dim, num_layers)`): Per-layer hidden_states from a text encoder (e.g. `Gemma3ForConditionalGeneration`). sequence_lengths (`torch.Tensor of shape `(batch_size,)`): The number of valid (non-padded) tokens for each batch instance. device: (`str` or `torch.device`, *optional*): torch device to place the resulting embeddings on padding_side: (`str`, *optional*, defaults to `"left"`): Whether the text tokenizer performs padding on the `"left"` or `"right"`. scale_factor (`int`, *optional*, defaults to `8`): Scaling factor to multiply the normalized hidden states by. eps (`float`, *optional*, defaults to `1e-6`): A small positive value for numerical stability when performing normalization. Returns: `torch.Tensor` of shape `(batch_size, seq_len, hidden_dim * num_layers)`: Normed and flattened text encoder hidden states. """ batch_size, seq_len, hidden_dim, num_layers = text_hidden_states.shape original_dtype = text_hidden_states.dtype # Create padding mask token_indices = torch.arange(seq_len, device=device).unsqueeze(0) if padding_side == "right": # For right padding, valid tokens are from 0 to sequence_length-1 mask = token_indices < sequence_lengths[:, None] # [batch_size, seq_len] elif padding_side == "left": # For left padding, valid tokens are from (T - sequence_length) to T-1 start_indices = seq_len - sequence_lengths[:, None] # [batch_size, 1] mask = token_indices >= start_indices # [B, T] else: raise ValueError(f"padding_side must be 'left' or 'right', got {padding_side}") mask = mask[:, :, None, None] # [batch_size, seq_len] --> [batch_size, seq_len, 1, 1] # Compute masked mean over non-padding positions of shape (batch_size, 1, 1, seq_len) masked_text_hidden_states = text_hidden_states.masked_fill(~mask, 0.0) num_valid_positions = (sequence_lengths * hidden_dim).view(batch_size, 1, 1, 1) masked_mean = masked_text_hidden_states.sum(dim=(1, 2), keepdim=True) / (num_valid_positions + eps) # Compute min/max over non-padding positions of shape (batch_size, 1, 1 seq_len) x_min = text_hidden_states.masked_fill(~mask, float("inf")).amin(dim=(1, 2), keepdim=True) x_max = text_hidden_states.masked_fill(~mask, float("-inf")).amax(dim=(1, 2), keepdim=True) # Normalization normalized_hidden_states = (text_hidden_states - masked_mean) / (x_max - x_min + eps) normalized_hidden_states = normalized_hidden_states * scale_factor # Pack the hidden states to a 3D tensor (batch_size, seq_len, hidden_dim * num_layers) normalized_hidden_states = normalized_hidden_states.flatten(2) mask_flat = mask.squeeze(-1).expand(-1, -1, hidden_dim * num_layers) normalized_hidden_states = normalized_hidden_states.masked_fill(~mask_flat, 0.0) normalized_hidden_states = normalized_hidden_states.to(dtype=original_dtype) return normalized_hidden_states # 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) sequence_lengths = prompt_attention_mask.sum(dim=-1) prompt_embeds = self._pack_text_embeds( text_encoder_hidden_states, sequence_lengths, device=device, padding_side=self.tokenizer.padding_side, scale_factor=scale_factor, ) prompt_embeds = prompt_embeds.to(dtype=dtype) # 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 # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline.encode_prompt def encode_prompt( self, prompt: str | list[str], negative_prompt: str | list[str] | None = None, do_classifier_free_guidance: bool = True, num_videos_per_prompt: int = 1, prompt_embeds: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, prompt_attention_mask: torch.Tensor | None = None, negative_prompt_attention_mask: torch.Tensor | None = None, 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 negative_prompt (`str` or `list[str]`, *optional*): The prompt or prompts not to guide the image generation. If not defined, one has to pass `negative_prompt_embeds` instead. Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`). do_classifier_free_guidance (`bool`, *optional*, defaults to `True`): Whether to use classifier free guidance or not. num_videos_per_prompt (`int`, *optional*, defaults to 1): Number of videos that should be generated per prompt. torch device to place the resulting embeddings on 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. negative_prompt_embeds (`torch.Tensor`, *optional*): Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. device: (`torch.device`, *optional*): torch device dtype: (`torch.dtype`, *optional*): torch dtype """ device = device or self._execution_device prompt = [prompt] if isinstance(prompt, str) else prompt if prompt is not None: batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] 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, scale_factor=scale_factor, device=device, dtype=dtype, ) if do_classifier_free_guidance and negative_prompt_embeds is None: negative_prompt = negative_prompt or "" negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt if prompt is not None and type(prompt) is not type(negative_prompt): raise TypeError( f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !=" f" {type(prompt)}." ) elif batch_size != len(negative_prompt): raise ValueError( f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:" f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches" " the batch size of `prompt`." ) negative_prompt_embeds, negative_prompt_attention_mask = self._get_gemma_prompt_embeds( prompt=negative_prompt, num_videos_per_prompt=num_videos_per_prompt, max_sequence_length=max_sequence_length, scale_factor=scale_factor, device=device, dtype=dtype, ) return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask 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, 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 {latents.ndim} dims. If you have packed (3D) latents, please unpack them (e.g." f" using the `_unpack_audio_latents` method)." ) @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._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 @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 @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 # Copied from diffusers.pipelines.ltx.pipeline_ltx_condition.LTXConditionPipeline.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 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]]: """ 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]]`: Returns a 3-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 indices in latent space to insert the corresponding condition. """ conditioning_frames, conditioning_strengths, conditioning_indices = [], [], [] 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): if isinstance(condition.frames, PIL.Image.Image): # Single image, convert to List[PIL.Image.Image] video_like_cond = [condition.frames] elif isinstance(condition.frames, np.ndarray) and condition.frames.ndim == 3: # Image-like ndarray of shape (H, W, C), insert frame dim in first axis video_like_cond = np.expand_dims(condition.frames, axis=0) elif isinstance(condition.frames, torch.Tensor) and condition.frames.ndim == 3: # Image-like tensor of shape (C, H, W), insert frame dim in first dim video_like_cond = condition.frames.unsqueeze(0) else: # Treat all other as videos. Note that this means 4D ndarrays and tensors will be treated as videos of # shape (F, H, W, C) and (F, C, H, W), respectively. video_like_cond = condition.frames condition_pixels = self.video_processor.preprocess_video( video_like_cond, height, width, resize_mode="crop" ) # 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) return conditioning_frames, conditioning_strengths, conditioning_indices def apply_visual_conditioning( self, latents: torch.Tensor, conditioning_mask: torch.Tensor, condition_latents: list[torch.Tensor], condition_strengths: list[float], condition_indices: list[int], latent_height: int, latent_width: int, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: """ Applies visual conditioning frames to an initial latent. Args: latents (`torch.Tensor`): Initial packed (patchified) latents of shape [batch_size, patch_seq_len, hidden_dim]. conditioning_mask (`torch.Tensor`, *optional*): Initial packed (patchified) conditioning mask of shape [batch_size, patch_seq_len, 1] with values in [0, 1] where 0 means that the denoising model output will be fully used and 1 means that the condition will be fully used (with intermediate values specifying a blend of the denoised and latent values). Returns: `Tuple[torch.Tensor, torch.Tensor, torch.Tensor]`: Returns a 3-tuple of tensors where: 1. The first element is the packed video latents (with unchanged shape [batch_size, patch_seq_len, hidden_dim]) with the conditions applied 2. The second element is the packed conditioning mask with conditioning strengths applied 3. The third element holds the clean conditioning latents. """ # Latents-like tensor which holds the clean conditioning latents clean_latents = torch.zeros_like(latents) for cond, strength, latent_idx in zip(condition_latents, condition_strengths, condition_indices): num_cond_tokens = cond.size(1) start_token_idx = latent_idx * latent_height * latent_width end_token_idx = start_token_idx + num_cond_tokens # Overwrite the portion of latents starting with start_token_idx with the condition latents[:, start_token_idx:end_token_idx] = cond conditioning_mask[:, start_token_idx:end_token_idx] = strength clean_latents[:, start_token_idx:end_token_idx] = cond return latents, conditioning_mask, clean_latents def prepare_latents( self, conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, batch_size: int = 1, num_channels_latents: int = 128, height: int = 512, width: int = 768, num_frames: int = 121, noise_scale: float = 1.0, dtype: torch.dtype | None = None, device: torch.device | None = None, generator: torch.Generator | None = None, latents: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: 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 shape = (batch_size, num_channels_latents, latent_num_frames, latent_height, latent_width) mask_shape = (batch_size, 1, latent_num_frames, latent_height, latent_width) if latents is not None: # Latents are expected to be unpacked (5D) with shape [B, F, C, H, W] latents = self._normalize_latents( latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor ) else: # NOTE: we set the initial latents to zeros rather a sample from the standard Gaussian prior because we # will sample from the prior later once we have calculated the conditioning mask latents = torch.zeros(shape, device=device, dtype=dtype) conditioning_mask = latents.new_zeros(mask_shape) latents = self._pack_latents( latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size ) conditioning_mask = self._pack_latents( conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size ) # [B, seq_len, 1] if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape[:2]: raise ValueError( f"Provided `latents` tensor has shape {latents.shape}, but the expected shape is {conditioning_mask.shape[:2] + (num_channels_latents,)}." ) if isinstance(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." ) generator = generator[0] condition_frames, condition_strengths, condition_indices = self.preprocess_conditions( conditions, height, width, num_frames, device=device ) condition_latents = [] for condition_tensor in condition_frames: condition_latent = retrieve_latents( self.vae.encode(condition_tensor), generator=generator, sample_mode="argmax" ) condition_latent = self._normalize_latents( condition_latent, self.vae.latents_mean, self.vae.latents_std ).to(device=device, dtype=dtype) condition_latent = self._pack_latents( condition_latent, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size ) condition_latents.append(condition_latent) # NOTE: following the I2V pipeline, we return a conditioning mask. The original LTX 2 code uses a denoising # mask, which is the inverse of the conditioning mask (`denoise_mask = 1 - conditioning_mask`) latents, conditioning_mask, clean_latents = self.apply_visual_conditioning( latents, conditioning_mask, condition_latents, condition_strengths, condition_indices, latent_height=latent_height, latent_width=latent_width, ) # Sample from the standard Gaussian prior (or an intermediate Gaussian distribution if noise_scale < 1.0). noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) scaled_mask = (1.0 - conditioning_mask) * noise_scale # Add noise to the `latents` so that it is at the noise level specified by `noise_scale`. latents = noise * scaled_mask + latents * (1 - scaled_mask) return latents, conditioning_mask, clean_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 shape = (batch_size, num_channels_latents, audio_latent_length, latent_mel_bins) 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." ) latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) latents = self._pack_audio_latents(latents) return latents @property def guidance_scale(self): return self._guidance_scale @property def guidance_rescale(self): return self._guidance_rescale @property def do_classifier_free_guidance(self): return self._guidance_scale > 1.0 @property def num_timesteps(self): return self._num_timesteps @property def current_timestep(self): return self._current_timestep @property def attention_kwargs(self): return self._attention_kwargs @property def interrupt(self): return self._interrupt @torch.no_grad() @replace_example_docstring(EXAMPLE_DOC_STRING) def __call__( self, conditions: LTX2VideoCondition | list[LTX2VideoCondition] | None = None, prompt: str | list[str] = None, negative_prompt: str | list[str] | None = None, height: int = 512, width: int = 768, num_frames: int = 121, frame_rate: float = 24.0, num_inference_steps: int = 40, sigmas: list[float] | None = None, timesteps: list[float] | None = None, guidance_scale: float = 4.0, guidance_rescale: float = 0.0, noise_scale: float | None = None, num_videos_per_prompt: int | None = 1, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.Tensor | None = None, audio_latents: torch.Tensor | None = None, prompt_embeds: torch.Tensor | None = None, prompt_attention_mask: torch.Tensor | None = None, negative_prompt_embeds: torch.Tensor | None = None, negative_prompt_attention_mask: torch.Tensor | None = None, decode_timestep: float | list[float] = 0.0, decode_noise_scale: float | list[float] | None = None, output_type: str = "pil", return_dict: bool = True, attention_kwargs: dict[str, Any] | None = None, callback_on_step_end: Callable[[int, int], None] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 1024, ): r""" Function invoked when calling the pipeline for generation. Args: conditions (`List[LTXVideoCondition], *optional*`): The list of frame-conditioning items for the video generation. prompt (`str` or `List[str]`, *optional*): The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`. instead. height (`int`, *optional*, defaults to `512`): The height in pixels of the generated image. This is set to 480 by default for the best results. width (`int`, *optional*, defaults to `768`): The width in pixels of the generated image. This is set to 848 by default for the best results. num_frames (`int`, *optional*, defaults to `121`): The number of video frames to generate frame_rate (`float`, *optional*, defaults to `24.0`): The frames per second (FPS) of the generated video. num_inference_steps (`int`, *optional*, defaults to 40): The number of denoising steps. More denoising steps usually lead to a higher quality image at the expense of slower inference. sigmas (`List[float]`, *optional*): Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. timesteps (`List[int]`, *optional*): Custom timesteps to use for the denoising process with schedulers which support a `timesteps` argument in their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed will be used. Must be in descending order. guidance_scale (`float`, *optional*, defaults to `4.0`): Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://huggingface.co/papers/2207.12598). `guidance_scale` is defined as `w` of equation 2. of [Imagen Paper](https://huggingface.co/papers/2205.11487). Guidance scale is enabled by setting `guidance_scale > 1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`, usually at the expense of lower image quality. guidance_rescale (`float`, *optional*, defaults to 0.0): Guidance rescale factor proposed by [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891) `guidance_scale` is defined as `φ` in equation 16. of [Common Diffusion Noise Schedules and Sample Steps are Flawed](https://huggingface.co/papers/2305.08891). Guidance rescale factor should fix overexposure when using zero terminal SNR. noise_scale (`float`, *optional*, defaults to `None`): The interpolation factor between random noise and denoised latents at each timestep. Applying noise to the `latents` and `audio_latents` before continue denoising. If not set, will be inferred from the sigma schedule. 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*): One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html) to make generation deterministic. latents (`torch.Tensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for video generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will be generated by sampling using the supplied random `generator`. audio_latents (`torch.Tensor`, *optional*): Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for audio generation. Can be used to tweak the same generation with different prompts. If not provided, a latents tensor will be generated by sampling using the supplied random `generator`. 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 text embeddings. negative_prompt_embeds (`torch.FloatTensor`, *optional*): Pre-generated negative text embeddings. For PixArt-Sigma this negative prompt should be "". If not provided, negative_prompt_embeds will be generated from `negative_prompt` input argument. negative_prompt_attention_mask (`torch.FloatTensor`, *optional*): Pre-generated attention mask for negative text embeddings. decode_timestep (`float`, defaults to `0.0`): The timestep at which generated video is decoded. decode_noise_scale (`float`, defaults to `None`): The interpolation factor between random noise and denoised latents at the decode timestep. output_type (`str`, *optional*, defaults to `"pil"`): The output format of the generate image. Choose between [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`. return_dict (`bool`, *optional*, defaults to `True`): Whether or not to return a [`~pipelines.ltx.LTX2PipelineOutput`] instead of a plain tuple. attention_kwargs (`dict`, *optional*): A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under `self.processor` in [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py). callback_on_step_end (`Callable`, *optional*): A function that calls at the end of each denoising steps during the inference. The function is called with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int, callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by `callback_on_step_end_tensor_inputs`. callback_on_step_end_tensor_inputs (`List`, *optional*): The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the `._callback_tensor_inputs` attribute of your pipeline class. max_sequence_length (`int`, *optional*, defaults to `1024`): Maximum sequence length to use with the `prompt`. Examples: Returns: [`~pipelines.ltx.LTX2PipelineOutput`] or `tuple`: If `return_dict` is `True`, [`~pipelines.ltx.LTX2PipelineOutput`] is returned, otherwise a `tuple` is returned where the first element is a list with the generated images. """ if isinstance(callback_on_step_end, (PipelineCallback, MultiPipelineCallbacks)): callback_on_step_end_tensor_inputs = callback_on_step_end.tensor_inputs # 1. Check inputs. Raise error if not correct 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, negative_prompt_embeds=negative_prompt_embeds, prompt_attention_mask=prompt_attention_mask, negative_prompt_attention_mask=negative_prompt_attention_mask, latents=latents, audio_latents=audio_latents, ) self._guidance_scale = guidance_scale self._guidance_rescale = guidance_rescale self._attention_kwargs = attention_kwargs self._interrupt = False self._current_timestep = None # 2. Define call parameters if prompt is not None and isinstance(prompt, str): batch_size = 1 elif prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] if conditions is not None and not isinstance(conditions, list): conditions = [conditions] # Infer noise scale: first (largest) sigma value if using custom sigmas, else 1.0 if noise_scale is None: noise_scale = sigmas[0] if sigmas is not None else 1.0 device = self._execution_device # 3. Prepare text embeddings ( prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask, ) = self.encode_prompt( prompt=prompt, negative_prompt=negative_prompt, do_classifier_free_guidance=self.do_classifier_free_guidance, num_videos_per_prompt=num_videos_per_prompt, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, prompt_attention_mask=prompt_attention_mask, negative_prompt_attention_mask=negative_prompt_attention_mask, max_sequence_length=max_sequence_length, device=device, ) if self.do_classifier_free_guidance: prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds], dim=0) prompt_attention_mask = torch.cat([negative_prompt_attention_mask, prompt_attention_mask], dim=0) additive_attention_mask = (1 - prompt_attention_mask.to(prompt_embeds.dtype)) * -1000000.0 connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( prompt_embeds, additive_attention_mask, additive_mask=True ) # 4. Prepare latent variables 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 if latents is not None: logger.info( "Got latents of shape [batch_size, latent_dim, latent_frames, latent_height, latent_width], `latent_num_frames`, `latent_height`, `latent_width` will be inferred." ) _, _, latent_num_frames, latent_height, latent_width = latents.shape # [B, C, F, H, W] video_sequence_length = latent_num_frames * latent_height * latent_width num_channels_latents = self.transformer.config.in_channels latents, conditioning_mask, clean_latents = self.prepare_latents( conditions, batch_size * num_videos_per_prompt, num_channels_latents, height, width, num_frames, noise_scale, torch.float32, device, generator, latents, ) if self.do_classifier_free_guidance: conditioning_mask = torch.cat([conditioning_mask, conditioning_mask]) duration_s = num_frames / frame_rate audio_latents_per_second = ( self.audio_sampling_rate / self.audio_hop_length / float(self.audio_vae_temporal_compression_ratio) ) audio_num_frames = round(duration_s * audio_latents_per_second) if audio_latents is not None: logger.info( "Got audio_latents of shape [batch_size, num_channels, audio_num_frames, mel_bins], `audio_num_frames` will be inferred." ) _, _, audio_num_frames, _ = audio_latents.shape # [B, C, L, M] num_mel_bins = self.audio_vae.config.mel_bins if getattr(self, "audio_vae", None) is not None else 64 latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio num_channels_latents_audio = ( self.audio_vae.config.latent_channels if getattr(self, "audio_vae", None) is not None else 8 ) audio_latents = self.prepare_audio_latents( batch_size * num_videos_per_prompt, num_channels_latents=num_channels_latents_audio, audio_latent_length=audio_num_frames, num_mel_bins=num_mel_bins, noise_scale=noise_scale, dtype=torch.float32, device=device, generator=generator, latents=audio_latents, ) # 5. Prepare timesteps sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas mu = calculate_shift( video_sequence_length, self.scheduler.config.get("base_image_seq_len", 1024), self.scheduler.config.get("max_image_seq_len", 4096), self.scheduler.config.get("base_shift", 0.95), self.scheduler.config.get("max_shift", 2.05), ) # For now, duplicate the scheduler for use with the audio latents audio_scheduler = copy.deepcopy(self.scheduler) _, _ = retrieve_timesteps( audio_scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu, ) timesteps, num_inference_steps = retrieve_timesteps( self.scheduler, num_inference_steps, device, timesteps, sigmas=sigmas, mu=mu, ) num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0) self._num_timesteps = len(timesteps) # 6. Prepare micro-conditions rope_interpolation_scale = ( self.vae_temporal_compression_ratio / frame_rate, self.vae_spatial_compression_ratio, self.vae_spatial_compression_ratio, ) # Pre-compute video and audio positional ids as they will be the same at each step of the denoising loop video_coords = self.transformer.rope.prepare_video_coords( latents.shape[0], latent_num_frames, latent_height, latent_width, latents.device, fps=frame_rate ) audio_coords = self.transformer.audio_rope.prepare_audio_coords( audio_latents.shape[0], audio_num_frames, audio_latents.device ) # Duplicate the positional ids as well if using CFG if self.do_classifier_free_guidance: video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) # Repeat twice in batch dim audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) # 7. Denoising loop with self.progress_bar(total=num_inference_steps) as progress_bar: for i, t in enumerate(timesteps): if self.interrupt: continue self._current_timestep = t latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents latent_model_input = latent_model_input.to(prompt_embeds.dtype) audio_latent_model_input = ( torch.cat([audio_latents] * 2) if self.do_classifier_free_guidance else audio_latents ) audio_latent_model_input = audio_latent_model_input.to(prompt_embeds.dtype) timestep = t.expand(latent_model_input.shape[0]) video_timestep = timestep.unsqueeze(-1) * (1 - conditioning_mask.squeeze(-1)) with self.transformer.cache_context("cond_uncond"): noise_pred_video, noise_pred_audio = self.transformer( hidden_states=latent_model_input, audio_hidden_states=audio_latent_model_input, encoder_hidden_states=connector_prompt_embeds, audio_encoder_hidden_states=connector_audio_prompt_embeds, timestep=video_timestep, audio_timestep=timestep, encoder_attention_mask=connector_attention_mask, audio_encoder_attention_mask=connector_attention_mask, num_frames=latent_num_frames, height=latent_height, width=latent_width, fps=frame_rate, audio_num_frames=audio_num_frames, video_coords=video_coords, audio_coords=audio_coords, # rope_interpolation_scale=rope_interpolation_scale, attention_kwargs=attention_kwargs, return_dict=False, ) noise_pred_video = noise_pred_video.float() noise_pred_audio = noise_pred_audio.float() if self.do_classifier_free_guidance: noise_pred_video_uncond, noise_pred_video_text = noise_pred_video.chunk(2) noise_pred_video = noise_pred_video_uncond + self.guidance_scale * ( noise_pred_video_text - noise_pred_video_uncond ) noise_pred_audio_uncond, noise_pred_audio_text = noise_pred_audio.chunk(2) noise_pred_audio = noise_pred_audio_uncond + self.guidance_scale * ( noise_pred_audio_text - noise_pred_audio_uncond ) if self.guidance_rescale > 0: # Based on 3.4. in https://huggingface.co/papers/2305.08891 noise_pred_video = rescale_noise_cfg( noise_pred_video, noise_pred_video_text, guidance_rescale=self.guidance_rescale ) noise_pred_audio = rescale_noise_cfg( noise_pred_audio, noise_pred_audio_text, guidance_rescale=self.guidance_rescale ) # NOTE: use only the first chunk of conditioning mask in case it is duplicated for CFG bsz = noise_pred_video.size(0) sigma = self.scheduler.sigmas[i] # Convert the noise_pred_video velocity model prediction into a sample (x0) prediction denoised_sample = latents - noise_pred_video * sigma # Apply the (packed) conditioning mask to the denoised (x0) sample and clean conditioning. The # conditioning mask contains conditioning strengths from 0 (always use denoised sample) to 1 (always # use conditions), with intermediate values specifying how strongly to follow the conditions. denoised_sample_cond = ( denoised_sample * (1 - conditioning_mask[:bsz]) + clean_latents.float() * conditioning_mask[:bsz] ).to(noise_pred_video.dtype) # Convert the denoised (x0) sample back to a velocity for the scheduler denoised_latents_cond = ((latents - denoised_sample_cond) / sigma).to(noise_pred_video.dtype) # Compute the previous noisy sample x_t -> x_t-1 latents = self.scheduler.step(denoised_latents_cond, t, latents, return_dict=False)[0] # NOTE: for now duplicate scheduler for audio latents in case self.scheduler sets internal state in # the step method (such as _step_index) audio_latents = audio_scheduler.step(noise_pred_audio, 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: callback_kwargs[k] = locals()[k] callback_outputs = callback_on_step_end(self, i, t, callback_kwargs) latents = callback_outputs.pop("latents", latents) prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds) # call the callback, if provided if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0): progress_bar.update() if XLA_AVAILABLE: xm.mark_step() latents = self._unpack_latents( latents, latent_num_frames, latent_height, latent_width, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ) latents = self._denormalize_latents( latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor ) audio_latents = self._denormalize_audio_latents( audio_latents, self.audio_vae.latents_mean, self.audio_vae.latents_std ) audio_latents = self._unpack_audio_latents(audio_latents, audio_num_frames, num_mel_bins=latent_mel_bins) if output_type == "latent": video = latents audio = audio_latents else: latents = latents.to(prompt_embeds.dtype) if not self.vae.config.timestep_conditioning: timestep = None else: noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=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=latents.dtype) decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ :, None, None, None, None ] latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise latents = latents.to(self.vae.dtype) video = self.vae.decode(latents, timestep, return_dict=False)[0] video = self.video_processor.postprocess_video(video, output_type=output_type) audio_latents = audio_latents.to(self.audio_vae.dtype) generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] audio = self.vocoder(generated_mel_spectrograms) # Offload all models self.maybe_free_model_hooks() if not return_dict: return (video, audio) return LTX2PipelineOutput(frames=video, audio=audio)