# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ Fully independent LTX-2.3 pipeline for vLLM-Omni. This pipeline does NOT inherit from LTX2Pipeline because: - LTX-2.3 uses a different text encoding strategy (flatten ALL 49 hidden states vs. LTX-2's _pack_text_embeds with per-layer normalization and pooling) - LTX-2.3 connectors expect the padding_side API (not additive_mask) - LTX-2.3 uses a BWE vocoder outputting 48kHz audio (not 16kHz) - LTX-2.3 transformer requires the sigma parameter for prompt modulation - CPU offloading is required for the 22B transformer (~44GB VRAM) """ from __future__ import annotations import copy import json import os from collections.abc import Iterable from contextlib import nullcontext from typing import Any import numpy as np import torch from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler from diffusers.pipelines.ltx2 import LTX2TextConnectors from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps from diffusers.utils.torch_utils import randn_tensor from diffusers.video_processor import VideoProcessor from huggingface_hub import hf_hub_download from torch import nn from transformers import AutoTokenizer, Gemma3ForConditionalGeneration from vllm.logger import init_logger from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from .pipeline_ltx2 import ( _get_prompt_field, calculate_shift, create_transformer_from_config, load_transformer_config, ) logger = init_logger(__name__) # Try to import LTX2VocoderWithBWE (diffusers >= 0.38.0) try: from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE except ImportError: LTX2VocoderWithBWE = None def _detect_vocoder_output_sample_rate(model: str) -> int | None: """Detect the vocoder output sample rate from vocoder/config.json. This runs at factory time (engine process) so the rate is captured in the post-process closure and doesn't need cross-process communication. Returns: Output sample rate (e.g. 48000 for LTX-2.3 BWE vocoder) or None. """ vocoder_config_path = os.path.join(model, "vocoder", "config.json") if not os.path.exists(vocoder_config_path): try: vocoder_config_path = hf_hub_download(model, "vocoder/config.json") except Exception: return None try: with open(vocoder_config_path) as f: cfg = json.load(f) return cfg.get("output_sampling_rate") except Exception: return None def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): """Factory for the LTX-2.3 post-process function. Detects the vocoder output sample rate at factory time and captures it in the closure so that the audio_sample_rate flows through DiffusionEngine -> OmniRequestOutput -> serving_video. """ output_sr = _detect_vocoder_output_sample_rate(od_config.model) def post_process_func(output: tuple[torch.Tensor, torch.Tensor] | torch.Tensor): if isinstance(output, tuple) and len(output) == 2: video, audio = output if isinstance(audio, torch.Tensor): audio = audio.detach().cpu() result: dict[str, Any] = {"video": video, "audio": audio} if output_sr is not None: result["audio_sample_rate"] = output_sr return result return output return post_process_func class LTX23Pipeline(nn.Module, ProgressBarMixin): """Fully independent LTX-2.3 pipeline. Key differences from LTX2Pipeline: - Text encoding: uses ALL 49 hidden states from Gemma-3-12B, flattened - Connectors: uses padding_side API (not additive_mask) - Vocoder: uses LTX2VocoderWithBWE (48kHz output) - Transformer: passes sigma for prompt_adaln - CPU offloading: text encoder, connectors, VAE, vocoder stay on CPU """ def __init__( self, *, od_config: OmniDiffusionConfig, prefix: str = "", ): super().__init__() self.od_config = od_config self.device = get_local_device() dtype = getattr(od_config, "dtype", torch.bfloat16) model = od_config.model local_files_only = os.path.exists(model) # Weight sources for transformer (loaded via AutoWeightsLoader) self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=od_config.model, subfolder="transformer", revision=None, prefix="transformer.", fall_back_to_pt=True, ), ] # --- Tokenizer (lightweight, stays wherever) --- self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) # --- Text encoder: load on CPU, move to GPU only during encoding --- with torch.device("cpu"): self.text_encoder = Gemma3ForConditionalGeneration.from_pretrained( model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ) # --- Connectors: CPU (LTX-2.3 connectors include caption projection) --- self.connectors = LTX2TextConnectors.from_pretrained( model, subfolder="connectors", torch_dtype=dtype, local_files_only=local_files_only ) # --- VAE, Audio VAE: CPU --- self.vae = AutoencoderKLLTX2Video.from_pretrained( model, subfolder="vae", torch_dtype=dtype, local_files_only=local_files_only ) self.audio_vae = AutoencoderKLLTX2Audio.from_pretrained( model, subfolder="audio_vae", torch_dtype=dtype, local_files_only=local_files_only ) # --- Vocoder: prefer BWE vocoder (48kHz) for LTX-2.3 --- vocoder_cls = LTX2VocoderWithBWE or LTX2Vocoder try: self.vocoder = vocoder_cls.from_pretrained( model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only ) except (TypeError, OSError, ValueError): self.vocoder = LTX2Vocoder.from_pretrained( model, subfolder="vocoder", torch_dtype=dtype, local_files_only=local_files_only ) # --- Transformer: created empty, weights loaded via AutoWeightsLoader --- transformer_config = load_transformer_config(model, "transformer", local_files_only) self.transformer = create_transformer_from_config(transformer_config) # --- Scheduler --- self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( model, subfolder="scheduler", local_files_only=local_files_only ) # --- Derived compression ratios --- self.vae_spatial_compression_ratio = self.vae.spatial_compression_ratio if self.vae is not None else 32 self.vae_temporal_compression_ratio = self.vae.temporal_compression_ratio if self.vae is not None else 8 self.audio_vae_mel_compression_ratio = self.audio_vae.mel_compression_ratio if self.audio_vae is not None else 4 self.audio_vae_temporal_compression_ratio = ( self.audio_vae.temporal_compression_ratio if self.audio_vae is not None else 4 ) self.transformer_spatial_patch_size = self.transformer.config.patch_size if self.transformer is not None else 1 self.transformer_temporal_patch_size = ( self.transformer.config.patch_size_t if self.transformer is not None else 1 ) self.audio_sampling_rate = self.audio_vae.config.sample_rate if self.audio_vae is not None else 16000 self.audio_hop_length = self.audio_vae.config.mel_hop_length if self.audio_vae is not None else 160 self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio) # Tokenizer max length tokenizer_max_length = 1024 if self.tokenizer is not None: tokenizer_max_length = self.tokenizer.model_max_length if tokenizer_max_length is None or tokenizer_max_length > 100000: encoder_config = getattr(self.text_encoder, "config", None) config_max_len = getattr(encoder_config, "max_position_embeddings", None) if config_max_len is None: config_max_len = getattr(encoder_config, "max_seq_len", None) tokenizer_max_length = config_max_len or 1024 self.tokenizer_max_length = int(tokenizer_max_length) # Pipeline state self._guidance_scale = None self._attention_kwargs = None self._interrupt = False self._num_timesteps = None self._current_timestep = None # ------------------------------------------------------------------ # Text Encoding (LTX-2.3 specific) # ------------------------------------------------------------------ def _get_gemma_prompt_embeds( self, prompt: str | list[str], num_videos_per_prompt: int = 1, max_sequence_length: int = 1024, device: torch.device | None = None, dtype: torch.dtype | None = None, ): """Encode prompts using Gemma-3-12B, returning ALL 49 hidden states flattened. LTX-2.3 differs from LTX-2 in text encoding: - LTX-2: uses _pack_text_embeds (layer selection + pooling) - LTX-2.3: stacks ALL 49 hidden states and flattens to [B, seq, 188160] The connectors unflatten, apply per_token_rms_norm, and project internally. """ device = device or self.device dtype = dtype or self.text_encoder.dtype prompt = [prompt] if isinstance(prompt, str) else prompt batch_size = len(prompt) if self.tokenizer is not None: 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.to(device) prompt_attention_mask = text_inputs.attention_mask.to(device) # Move text encoder to GPU for encoding self.text_encoder.to(device) text_encoder_outputs = self.text_encoder( input_ids=text_input_ids, attention_mask=prompt_attention_mask, output_hidden_states=True, ) # Move text encoder back to CPU immediately self.text_encoder.to("cpu") torch.cuda.empty_cache() hidden_states = text_encoder_outputs.hidden_states # LTX-2.3: Stack ALL 49 hidden states and flatten # [49 x (B, seq, 3840)] -> [B, seq, 3840, 49] -> [B, seq, 188160] prompt_embeds = torch.stack(hidden_states, dim=-1).flatten(2, 3).to(dtype=dtype) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_videos_per_prompt, 1) prompt_embeds = prompt_embeds.view(batch_size * num_videos_per_prompt, seq_len, -1) prompt_attention_mask = prompt_attention_mask.view(batch_size, -1) prompt_attention_mask = prompt_attention_mask.repeat(num_videos_per_prompt, 1) return prompt_embeds, prompt_attention_mask def encode_prompt( self, prompt: str | list[str], 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, device: torch.device | None = None, dtype: torch.dtype | None = None, ): device = device or self.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, 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 as `prompt`, but got {type(negative_prompt)} !=" f" {type(prompt)}." ) if isinstance(negative_prompt, list) and 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, device=device, dtype=dtype, ) return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask # ------------------------------------------------------------------ # Latent utilities (shared with LTX2Pipeline) # ------------------------------------------------------------------ @staticmethod def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor: 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 def _unpack_latents( latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1, ) -> torch.Tensor: 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 def _normalize_latents( latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 ) -> torch.Tensor: 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 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 def _denormalize_latents( latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0 ) -> torch.Tensor: 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 def _denormalize_audio_latents(latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor): latents_mean = latents_mean.to(latents.device, latents.dtype) latents_std = latents_std.to(latents.device, latents.dtype) return (latents * latents_std) + latents_mean @staticmethod def _pack_audio_latents( latents: torch.Tensor, patch_size: int | None = None, patch_size_t: int | None = None ) -> torch.Tensor: if patch_size is not None and patch_size_t is not None: 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: latents = latents.transpose(1, 2).flatten(2, 3) return latents @staticmethod 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: 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: latents = latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2) return latents @staticmethod def _unpad_audio_latents(latents: torch.Tensor, num_frames: int) -> torch.Tensor: return latents[:, :num_frames] # ------------------------------------------------------------------ # Latent preparation # ------------------------------------------------------------------ def prepare_latents( self, batch_size: int = 1, num_channels_latents: int = 128, height: int = 512, width: int = 768, num_frames: int = 121, 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: if latents.ndim == 5: latents = self._normalize_latents( latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor ) latents = self._pack_latents( latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size ) if latents.ndim != 3: raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) latents = noise_scale * noise + (1 - noise_scale) * latents return latents.to(device=device, dtype=dtype) height = height // self.vae_spatial_compression_ratio width = width // self.vae_spatial_compression_ratio num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1 shape = (batch_size, num_channels_latents, num_frames, height, width) latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) latents = self._pack_latents(latents, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size) return latents def prepare_audio_latents( self, batch_size: int = 1, num_channels_latents: int = 8, audio_latent_length: int = 1, num_mel_bins: int = 64, noise_scale: float = 0.0, dtype: torch.dtype | None = None, device: torch.device | None = None, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.Tensor | None = None, ) -> tuple[torch.Tensor, int, int]: original_latent_length = audio_latent_length padded_latent_length = original_latent_length latent_mel_bins = num_mel_bins // self.audio_vae_mel_compression_ratio if latents is not None: if latents.ndim == 4: latents = self._pack_audio_latents(latents) if latents.ndim != 3: raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].") latents = self._normalize_audio_latents(latents, self.audio_vae.latents_mean, self.audio_vae.latents_std) noise = randn_tensor(latents.shape, generator=generator, device=latents.device, dtype=latents.dtype) latents = noise_scale * noise + (1 - noise_scale) * latents return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins) latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) latents = self._pack_audio_latents(latents) return latents, original_latent_length, padded_latent_length # ------------------------------------------------------------------ # Properties # ------------------------------------------------------------------ @property def guidance_scale(self): return self._guidance_scale @property def do_classifier_free_guidance(self): return self._guidance_scale is not None and 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 interrupt(self): return self._interrupt # ------------------------------------------------------------------ # Input validation # ------------------------------------------------------------------ def check_inputs( self, prompt, height, width, prompt_embeds=None, negative_prompt_embeds=None, prompt_attention_mask=None, negative_prompt_attention_mask=None, ): if height % 32 != 0 or width % 32 != 0: raise ValueError(f"`height` and `width` must be divisible by 32 but are {height} and {width}.") if prompt is not None and prompt_embeds is not None: raise ValueError("Cannot forward both `prompt` and `prompt_embeds`.") elif prompt is None and prompt_embeds is None: raise ValueError("Provide either `prompt` or `prompt_embeds`.") elif prompt is not None and not isinstance(prompt, (str, 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 got: `prompt_attention_mask` " f"{prompt_attention_mask.shape} != `negative_prompt_attention_mask` " f"{negative_prompt_attention_mask.shape}." ) # ------------------------------------------------------------------ # Cache context # ------------------------------------------------------------------ def _transformer_cache_context(self, context_name: str): cache_context = getattr(self.transformer, "cache_context", None) if callable(cache_context): return cache_context(context_name) return nullcontext() # ------------------------------------------------------------------ # Forward pass # ------------------------------------------------------------------ @torch.no_grad() def forward( self, req: OmniDiffusionRequest, prompt: str | list[str] | None = None, negative_prompt: str | list[str] | None = None, height: int | None = None, width: int | None = None, num_frames: int | None = None, frame_rate: float | None = None, num_inference_steps: int | None = None, sigmas: list[float] | None = None, timesteps: list[int] | None = None, guidance_scale: float = 4.0, noise_scale: float = 0.0, 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, negative_prompt_embeds: torch.Tensor | None = None, prompt_attention_mask: 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 = "np", return_dict: bool = True, attention_kwargs: dict[str, Any] | None = None, max_sequence_length: int | None = None, ) -> DiffusionOutput: # ---- Extract from request ---- prompt = [p if isinstance(p, str) else (p.get("prompt") or "") for p in req.prompts] or prompt if all(isinstance(p, str) or p.get("negative_prompt") is None for p in req.prompts): negative_prompt = None elif req.prompts: negative_prompt = ["" if isinstance(p, str) else (p.get("negative_prompt") or "") for p in req.prompts] height = req.sampling_params.height or height or 512 width = req.sampling_params.width or width or 768 num_frames = req.sampling_params.num_frames or num_frames or 121 frame_rate = req.sampling_params.resolved_frame_rate or frame_rate or 24.0 num_inference_steps = req.sampling_params.num_inference_steps or num_inference_steps or 40 # Enforce minimum of 2 timesteps for flow matching scheduler if timesteps is None: num_inference_steps = max(int(num_inference_steps), 2) elif len(timesteps) < 2: raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.") num_videos_per_prompt = ( req.sampling_params.num_outputs_per_prompt if req.sampling_params.num_outputs_per_prompt > 0 else num_videos_per_prompt or 1 ) max_sequence_length = ( req.sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length ) if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale if generator is None: generator = req.sampling_params.generator if generator is None and req.sampling_params.seed is not None: generator = torch.Generator(device=self.device).manual_seed(req.sampling_params.seed) latents = req.sampling_params.latents if req.sampling_params.latents is not None else latents audio_latents = ( req.sampling_params.audio_latents if req.sampling_params.audio_latents is not None else req.sampling_params.extra_args.get("audio_latents", audio_latents) ) # Override with pre-computed embeddings if provided in request req_prompt_embeds = [_get_prompt_field(p, "prompt_embeds") for p in req.prompts] if any(p is not None for p in req_prompt_embeds): prompt_embeds = torch.stack(req_prompt_embeds) req_negative_prompt_embeds = [_get_prompt_field(p, "negative_prompt_embeds") for p in req.prompts] if any(p is not None for p in req_negative_prompt_embeds): negative_prompt_embeds = torch.stack(req_negative_prompt_embeds) req_prompt_attention_masks = [ _get_prompt_field(p, "prompt_attention_mask") or _get_prompt_field(p, "attention_mask") for p in req.prompts ] if any(m is not None for m in req_prompt_attention_masks): prompt_attention_mask = torch.stack(req_prompt_attention_masks) req_negative_attention_masks = [ _get_prompt_field(p, "negative_prompt_attention_mask") or _get_prompt_field(p, "negative_attention_mask") for p in req.prompts ] if any(m is not None for m in req_negative_attention_masks): negative_prompt_attention_mask = torch.stack(req_negative_attention_masks) if req.sampling_params.decode_timestep is not None: decode_timestep = req.sampling_params.decode_timestep if req.sampling_params.decode_noise_scale is not None: decode_noise_scale = req.sampling_params.decode_noise_scale if req.sampling_params.output_type is not None: output_type = req.sampling_params.output_type self.check_inputs( prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, prompt_attention_mask=prompt_attention_mask, negative_prompt_attention_mask=negative_prompt_attention_mask, ) self._guidance_scale = guidance_scale self._attention_kwargs = attention_kwargs self._interrupt = False self._current_timestep = None if prompt is not None and isinstance(prompt, str): batch_size = 1 elif prompt is not None and isinstance(prompt, list): batch_size = len(prompt) else: batch_size = prompt_embeds.shape[0] device = self.device # ---- Encode prompts ---- ( 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, ) # ---- Connectors (LTX-2.3: padding_side API) ---- # Concatenate negative + positive embeddings BEFORE connector call, # matching diffusers which calls connectors once with batch=2. # This ensures batch-dependent operations produce identical results. 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) self.connectors.to(device) tokenizer_padding_side = getattr(self.tokenizer, "padding_side", "left") connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( prompt_embeds, prompt_attention_mask, padding_side=tokenizer_padding_side ) self.connectors.to("cpu") if torch.cuda.is_available(): torch.cuda.empty_cache() # ---- Prepare latents ---- 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 and latents.ndim == 5: _, _, latent_num_frames, latent_height, latent_width = latents.shape num_channels_latents = self.transformer.config.in_channels latents = self.prepare_latents( batch_size * num_videos_per_prompt, num_channels_latents, height, width, num_frames, noise_scale, torch.float32, device, generator, latents, ) 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 and audio_latents.ndim == 4: _, _, audio_num_frames, _ = audio_latents.shape num_mel_bins = self.audio_vae.config.mel_bins if self.audio_vae 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 self.audio_vae is not None else 8 audio_latents, original_audio_num_frames, padded_audio_num_frames = 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, ) # ---- Scheduler setup ---- sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas # Use max_image_seq_len (not actual video_sequence_length) for mu calculation, # matching diffusers' LTX2Pipeline which hardcodes this value. mu = calculate_shift( self.scheduler.config.get("max_image_seq_len", 4096), 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), ) 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, ) self._num_timesteps = len(timesteps) # ---- RoPE coordinates ---- 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], padded_audio_num_frames, audio_latents.device, ) # ---- CFG: duplicate coords for batch=2 ---- # Connector outputs are already batch=2 (neg+pos concatenated before connector call) if self.do_classifier_free_guidance: video_coords = video_coords.repeat((2,) + (1,) * (video_coords.ndim - 1)) audio_coords = audio_coords.repeat((2,) + (1,) * (audio_coords.ndim - 1)) # ---- Denoising loop ---- # Uses x0-space CFG (delta formulation) matching diffusers' LTX2Pipeline. # The velocity predictions are converted to x0, guidance is applied in x0 # space, then converted back to velocity for the scheduler step. with self.progress_bar(total=len(timesteps)) as pbar: for i, t in enumerate(timesteps): if self.interrupt: continue self._current_timestep = t # Duplicate latents for CFG (uncond + cond) latent_model_input = torch.cat([latents] * 2) if self.do_classifier_free_guidance else latents latent_model_input = latent_model_input.to(connector_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(connector_prompt_embeds.dtype) ts = t.expand(latent_model_input.shape[0]) 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=ts, sigma=ts, # LTX-2.3: sigma for prompt_adaln 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=padded_audio_num_frames, video_coords=video_coords, audio_coords=audio_coords, attention_kwargs=attention_kwargs, return_dict=False, ) noise_pred_video = noise_pred_video.float() noise_pred_audio = noise_pred_audio.float() # CFG in x0-space (delta formulation matching diffusers) if self.do_classifier_free_guidance: noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) # Convert velocity to x0: x0 = sample - velocity * sigma x0_video_cond = latents - noise_pred_video_cond * self.scheduler.sigmas[i] x0_video_uncond = latents - noise_pred_video_uncond * self.scheduler.sigmas[i] video_cfg_delta = (guidance_scale - 1) * (x0_video_cond - x0_video_uncond) x0_video_guided = x0_video_cond + video_cfg_delta noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) x0_audio_cond = audio_latents - noise_pred_audio_cond * audio_scheduler.sigmas[i] x0_audio_uncond = audio_latents - noise_pred_audio_uncond * audio_scheduler.sigmas[i] audio_cfg_delta = (guidance_scale - 1) * (x0_audio_cond - x0_audio_uncond) x0_audio_guided = x0_audio_cond + audio_cfg_delta # Convert x0 back to velocity: v = (sample - x0) / sigma noise_pred_video = (latents - x0_video_guided) / self.scheduler.sigmas[i] noise_pred_audio = (audio_latents - x0_audio_guided) / audio_scheduler.sigmas[i] latents = self.scheduler.step(noise_pred_video, t, latents, return_dict=False)[0] audio_latents = audio_scheduler.step(noise_pred_audio, t, audio_latents, return_dict=False)[0] pbar.update() # ---- Unpack and denormalize ---- 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._unpad_audio_latents(audio_latents, original_audio_num_frames) 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, original_audio_num_frames, num_mel_bins=latent_mel_bins, ) # ---- Decode ---- if output_type == "latent": video = latents audio = audio_latents else: latents = latents.to(connector_prompt_embeds.dtype) if not self.vae.config.timestep_conditioning: timestep_decode = 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_decode = torch.tensor(decode_timestep, device=device, dtype=latents.dtype) decode_noise_scale_t = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[ :, None, None, None, None ] latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise # Move VAE, audio_vae, vocoder to GPU for decoding self.vae.to(device) latents = latents.to(self.vae.dtype) video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] video = self.video_processor.postprocess_video(video, output_type=output_type) self.vae.to("cpu") self.audio_vae.to(device) audio_latents = audio_latents.to(self.audio_vae.dtype) generated_mel_spectrograms = self.audio_vae.decode(audio_latents, return_dict=False)[0] self.audio_vae.to("cpu") self.vocoder.to(device) audio = self.vocoder(generated_mel_spectrograms) self.vocoder.to("cpu") torch.cuda.empty_cache() return DiffusionOutput(output=(video, audio)) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return loader.load_weights(weights) class LTX23ImageToVideoPipeline(nn.Module): """LTX-2.3 image-to-video pipeline placeholder.""" def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): super().__init__() raise NotImplementedError( "LTX23ImageToVideoPipeline is not yet implemented. " "Use LTX23Pipeline for single-stage text-to-video generation." )