2026-03-05 23:31:48 +08:00
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
from __future__ import annotations
import copy
2026-04-03 23:17:49 +08:00
import os
from collections.abc import Iterable
2026-04-22 16:23:00 -04:00
from typing import Any , ClassVar
2026-03-05 23:31:48 +08:00
import numpy as np
import PIL.Image
import torch
2026-04-03 23:17:49 +08:00
import torch.nn as nn
from diffusers import FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES , STAGE_2_DISTILLED_SIGMA_VALUES
2026-04-08 04:44:04 +01:00
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
2026-03-05 23:31:48 +08:00
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
from diffusers.utils.torch_utils import randn_tensor
from diffusers.video_processor import VideoProcessor
2026-04-03 23:17:49 +08:00
from vllm.logger import init_logger
from vllm.model_executor.models.utils import AutoWeightsLoader
2026-03-05 23:31:48 +08:00
from vllm_omni.diffusion.data import DiffusionOutput , OmniDiffusionConfig
2026-04-03 23:17:49 +08:00
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.lora.manager import DiffusionLoRAManager
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
2026-04-20 06:21:43 -07:00
from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
2026-05-08 10:16:49 -04:00
from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
2026-03-05 23:31:48 +08:00
from vllm_omni.diffusion.request import OmniDiffusionRequest
2026-04-03 23:17:49 +08:00
from vllm_omni.lora.request import LoRARequest
2026-03-05 23:31:48 +08:00
from .pipeline_ltx2 import (
LTX2Pipeline ,
_get_prompt_field ,
calculate_shift ,
)
from .pipeline_ltx2 import (
get_ltx2_post_process_func as _get_ltx2_post_process_func ,
)
2026-04-03 23:17:49 +08:00
from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
logger = init_logger ( __name__ )
2026-03-05 23:31:48 +08:00
def get_ltx2_post_process_func ( od_config : OmniDiffusionConfig ):
return _get_ltx2_post_process_func ( od_config )
2026-04-08 04:44:04 +01:00
class _I2VVideoAudioScheduler :
"""Composite scheduler for I2V: uses _step_video_latents_i2v for video, standard step for audio."""
def __init__ ( self , pipeline , audio_scheduler , latent_num_frames , latent_height , latent_width ):
self . video_scheduler = pipeline . scheduler
self . audio_scheduler = audio_scheduler
self . _pipeline = pipeline
self . _latent_num_frames = latent_num_frames
self . _latent_height = latent_height
self . _latent_width = latent_width
def step ( self , noise_pred , t , latents , return_dict = False , generator = None ):
video_out = self . _pipeline . _step_video_latents_i2v (
noise_pred [ 0 ],
latents [ 0 ],
t [ 0 ],
self . _latent_num_frames ,
self . _latent_height ,
self . _latent_width ,
)
audio_out = self . audio_scheduler . step ( noise_pred [ 1 ], t [ 1 ], latents [ 1 ], return_dict = False , generator = generator )[
0
]
return (( video_out , audio_out ),)
2026-03-05 23:31:48 +08:00
class LTX2ImageToVideoPipeline ( LTX2Pipeline ):
support_image_input = True
def __init__ (
self ,
* ,
od_config : OmniDiffusionConfig ,
prefix : str = "" ,
):
super () . __init__ ( od_config = od_config , prefix = prefix )
self . video_processor = VideoProcessor ( vae_scale_factor = self . vae_spatial_compression_ratio , resample = "bilinear" )
@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
def prepare_latents (
self ,
image : torch . Tensor | None = None ,
batch_size : int = 1 ,
num_channels_latents : int = 128 ,
height : int = 512 ,
width : int = 768 ,
num_frames : int = 121 ,
2026-04-03 23:17:49 +08:00
noise_scale : float = 0.0 ,
2026-03-05 23:31:48 +08:00
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 , torch . Tensor ]:
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 )
mask_shape = ( batch_size , 1 , num_frames , height , width )
if latents is not None :
2026-04-03 23:17:49 +08:00
if latents . ndim == 5 :
# conditioning_mask needs to the same shape as latents in two stages generation.
batch_size , _ , num_frames , height , width = latents . shape
mask_shape = ( batch_size , 1 , num_frames , height , width )
conditioning_mask = latents . new_zeros ( mask_shape )
conditioning_mask [:, :, 0 ] = 1.0
latents = self . _normalize_latents (
latents , self . vae . latents_mean , self . vae . latents_std , self . vae . config . scaling_factor
)
latents = self . _create_noised_state ( latents , noise_scale * ( 1 - conditioning_mask ), generator )
# latents are of shape [B, C, F, H, W], need to be packed
latents = self . _pack_latents (
latents , self . transformer_spatial_patch_size , self . transformer_temporal_patch_size
)
else :
conditioning_mask = latents . new_zeros ( mask_shape )
conditioning_mask [:, :, 0 ] = 1.0
2026-03-05 23:31:48 +08:00
conditioning_mask = self . _pack_latents (
conditioning_mask , self . transformer_spatial_patch_size , self . transformer_temporal_patch_size
) . squeeze ( - 1 )
2026-04-03 23:17:49 +08:00
2026-03-05 23:31:48 +08:00
if latents . ndim != 3 or latents . shape [: 2 ] != conditioning_mask . shape :
raise ValueError (
"Provided `latents` tensor has shape"
f " { latents . shape } , but the expected shape is { conditioning_mask . shape + ( num_channels_latents ,) } ."
)
return latents . to ( device = device , dtype = dtype ), conditioning_mask
if image is None :
raise ValueError ( "`image` must be provided when `latents` is None." )
image_batch_size = image . shape [ 0 ]
if image_batch_size == 0 :
raise ValueError ( "`image` batch is empty." )
if batch_size % image_batch_size != 0 :
raise ValueError (
f "`batch_size` ( { batch_size } ) must be divisible by image batch size ( { image_batch_size } ) "
"for image-to-video outputs."
)
num_videos_per_prompt = batch_size // image_batch_size
if isinstance ( generator , list ):
if len ( generator ) != batch_size :
raise ValueError (
f "You have passed a list of generators of length { len ( generator ) } , but requested an effective"
f " batch size of { batch_size } . Make sure the batch size matches the length of the generators."
)
image_generators = [ generator [ i * num_videos_per_prompt ] for i in range ( image_batch_size )]
init_latents = [
retrieve_latents ( self . vae . encode ( image [ i ] . unsqueeze ( 0 ) . unsqueeze ( 2 )), image_generators [ i ], "argmax" )
for i in range ( image_batch_size )
]
else :
init_latents = [
retrieve_latents ( self . vae . encode ( img . unsqueeze ( 0 ) . unsqueeze ( 2 )), generator , "argmax" ) for img in image
]
init_latents = torch . cat ( init_latents , dim = 0 ) . to ( dtype )
if num_videos_per_prompt > 1 :
init_latents = init_latents . repeat_interleave ( num_videos_per_prompt , dim = 0 )
init_latents = self . _normalize_latents ( init_latents , self . vae . latents_mean , self . vae . latents_std )
init_latents = init_latents . repeat ( 1 , 1 , num_frames , 1 , 1 )
conditioning_mask = torch . zeros ( mask_shape , device = device , dtype = dtype )
conditioning_mask [:, :, 0 ] = 1.0
noise = randn_tensor ( shape , generator = generator , device = device , dtype = dtype )
latents = init_latents * conditioning_mask + noise * ( 1 - conditioning_mask )
conditioning_mask = self . _pack_latents (
conditioning_mask , self . transformer_spatial_patch_size , self . transformer_temporal_patch_size
) . squeeze ( - 1 )
latents = self . _pack_latents ( latents , self . transformer_spatial_patch_size , self . transformer_temporal_patch_size )
return latents , conditioning_mask
def check_inputs (
self ,
image ,
height ,
width ,
prompt ,
latents = None ,
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` have to be divisible by 32 but are { height } and { width } ." )
if image is None and latents is None :
raise ValueError ( "Provide either `image` or `latents`. Cannot leave both undefined." )
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."
)
if prompt is None and prompt_embeds is None :
raise ValueError (
"Provide either `prompt` or `prompt_embeds`. Cannot leave both `prompt` and `prompt_embeds` undefined."
)
if 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 } ."
)
def _step_video_latents_i2v (
self ,
noise_pred_video : torch . Tensor ,
latents : torch . Tensor ,
t : torch . Tensor ,
latent_num_frames : int ,
latent_height : int ,
latent_width : int ,
) -> torch . Tensor :
noise_pred_video = self . _unpack_latents (
noise_pred_video ,
latent_num_frames ,
latent_height ,
latent_width ,
self . transformer_spatial_patch_size ,
self . transformer_temporal_patch_size ,
)
latents_unpacked = self . _unpack_latents (
latents ,
latent_num_frames ,
latent_height ,
latent_width ,
self . transformer_spatial_patch_size ,
self . transformer_temporal_patch_size ,
)
noise_pred_video = noise_pred_video [:, :, 1 :]
noise_latents = latents_unpacked [:, :, 1 :]
pred_latents = self . scheduler . step ( noise_pred_video , t , noise_latents , return_dict = False )[ 0 ]
latents_unpacked = torch . cat ([ latents_unpacked [:, :, : 1 ], pred_latents ], dim = 2 )
latents = self . _pack_latents (
latents_unpacked , self . transformer_spatial_patch_size , self . transformer_temporal_patch_size
)
return latents
@torch.no_grad ()
def forward (
self ,
req : OmniDiffusionRequest ,
image : PIL . Image . Image | torch . Tensor | None = None ,
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 ,
2026-04-03 23:17:49 +08:00
sigmas : list [ float ] | None = None ,
2026-03-05 23:31:48 +08:00
timesteps : list [ int ] | None = None ,
guidance_scale : float = 4.0 ,
guidance_rescale : float = 0.0 ,
2026-04-03 23:17:49 +08:00
noise_scale : float = 0.0 ,
2026-03-05 23:31:48 +08:00
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 :
2026-04-08 04:44:04 +01:00
# Extract prompt/negative_prompt from request.
# Input format: req.prompts is a list of str or dict with "prompt"/"negative_prompt" keys.
2026-03-05 23:31:48 +08:00
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
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 req . sampling_params . guidance_rescale is not None :
guidance_rescale = req . sampling_params . guidance_rescale
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 )
)
2026-04-08 04:44:04 +01:00
# Override with pre-computed embeddings if provided in request.
2026-03-05 23:31:48 +08:00
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 ) # type: ignore[arg-type]
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 ) # type: ignore[arg-type]
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 ) # type: ignore[arg-type]
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 ) # type: ignore[arg-type]
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
if image is None and req . prompts :
raw_images = []
for prompt_item in req . prompts :
if isinstance ( prompt_item , str ):
raw_image = None
else :
multi_modal_data = prompt_item . get ( "multi_modal_data" ) or {}
raw_image = multi_modal_data . get ( "image" )
if raw_image is None :
additional = prompt_item . get ( "additional_information" ) or {}
raw_image = (
additional . get ( "preprocessed_image" )
or additional . get ( "pixel_values" )
or additional . get ( "image" )
)
if isinstance ( raw_image , list ):
raw_image = raw_image [ 0 ] if raw_image else None
if isinstance ( raw_image , str ):
raw_image = PIL . Image . open ( raw_image ) . convert ( "RGB" )
raw_images . append ( raw_image )
if any ( img is None for img in raw_images ):
if latents is None :
raise ValueError ( "Image is required for LTX2 I2V generation." )
if len ( raw_images ) == 1 :
image = raw_images [ 0 ]
elif raw_images :
image = raw_images
self . check_inputs (
image = image ,
height = height ,
width = width ,
prompt = prompt ,
latents = latents ,
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 . _guidance_rescale = guidance_rescale
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
(
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 ,
)
2026-04-08 04:44:04 +01:00
# Compute positive prompt connectors
2026-05-17 21:58:44 +08:00
tokenizer_padding_side = getattr ( self . tokenizer , "padding_side" , "left" )
2026-03-05 23:31:48 +08:00
connector_prompt_embeds , connector_audio_prompt_embeds , connector_attention_mask = self . connectors (
2026-05-17 21:58:44 +08:00
prompt_embeds , prompt_attention_mask , padding_side = tokenizer_padding_side
2026-03-05 23:31:48 +08:00
)
2026-04-08 04:44:04 +01:00
# Compute negative prompt connectors when CFG is enabled
2026-03-05 23:31:48 +08:00
negative_connector_prompt_embeds = None
negative_connector_audio_prompt_embeds = None
negative_connector_attention_mask = None
2026-04-08 04:44:04 +01:00
if self . do_classifier_free_guidance :
2026-03-05 23:31:48 +08:00
(
negative_connector_prompt_embeds ,
negative_connector_audio_prompt_embeds ,
negative_connector_attention_mask ,
) = self . connectors (
negative_prompt_embeds ,
2026-05-17 21:58:44 +08:00
negative_prompt_attention_mask ,
padding_side = tokenizer_padding_side ,
2026-03-05 23:31:48 +08:00
)
2026-04-03 23:17:49 +08:00
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 :
if latents . ndim == 5 :
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." # noqa
)
_ , _ , latent_num_frames , latent_height , latent_width = latents . shape # [B, C, F, H, W]
elif latents . ndim == 3 :
logger . warning (
f "You have supplied packed `latents` of shape { latents . shape } , so the latent dims cannot be"
f " inferred. Make sure the supplied `height`, `width`, and `num_frames` are correct."
)
else :
raise ValueError (
f "Provided `latents` tensor has shape { latents . shape } , but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, latent_dim, latent_frames, latent_height, latent_width]." # noqa
)
video_sequence_length = latent_num_frames * latent_height * latent_width
2026-03-05 23:31:48 +08:00
if latents is None :
if isinstance ( image , torch . Tensor ):
if image . ndim == 3 :
image = image . unsqueeze ( 0 )
elif isinstance ( image , list ) and image and isinstance ( image [ 0 ], torch . Tensor ):
image = torch . stack ( image , dim = 0 )
else :
image = self . video_processor . preprocess ( image , height = height , width = width )
image = image . to ( device = device , dtype = prompt_embeds . dtype )
num_channels_latents = self . transformer . config . in_channels
latents , conditioning_mask = self . prepare_latents (
image ,
batch_size * num_videos_per_prompt ,
num_channels_latents ,
height ,
width ,
num_frames ,
2026-04-03 23:17:49 +08:00
noise_scale ,
2026-03-05 23:31:48 +08:00
torch . float32 ,
device ,
generator ,
latents ,
)
2026-04-03 23:17:49 +08:00
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 :
if audio_latents . ndim == 4 :
logger . info (
"Got audio_latents of shape [batch_size, num_channels, audio_length, mel_bins], `audio_num_frames` will be inferred." # noqa
)
_ , _ , audio_num_frames , _ = audio_latents . shape # [B, C, L, M]
elif audio_latents . ndim == 3 :
logger . warning (
f "You have supplied packed `audio_latents` of shape { audio_latents . shape } , so the latent dims"
f " cannot be inferred. Make sure the supplied `num_frames` and `frame_rate` are correct."
)
else :
raise ValueError (
f "Provided `audio_latents` tensor has shape { audio_latents . shape } , but the expected shape is either [batch_size, seq_len, num_features] or [batch_size, num_channels, audio_length, mel_bins]." # noqa
)
2026-03-05 23:31:48 +08:00
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
)
2026-04-08 04:44:04 +01:00
audio_latents , original_audio_num_frames , padded_audio_num_frames = self . prepare_audio_latents (
2026-03-05 23:31:48 +08:00
batch_size * num_videos_per_prompt ,
num_channels_latents = num_channels_latents_audio ,
2026-04-03 23:17:49 +08:00
audio_latent_length = audio_num_frames ,
2026-03-05 23:31:48 +08:00
num_mel_bins = num_mel_bins ,
2026-04-03 23:17:49 +08:00
noise_scale = noise_scale ,
2026-03-05 23:31:48 +08:00
dtype = torch . float32 ,
device = device ,
generator = generator ,
latents = audio_latents ,
)
2026-04-03 23:17:49 +08:00
sigmas = np . linspace ( 1.0 , 1 / num_inference_steps , num_inference_steps ) if sigmas is None else sigmas
2026-03-05 23:31:48 +08:00
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 ),
)
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 )
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 (
2026-04-08 04:44:04 +01:00
audio_latents . shape [ 0 ], padded_audio_num_frames , audio_latents . device
2026-03-05 23:31:48 +08:00
)
2026-04-08 04:44:04 +01:00
i2v_scheduler = _I2VVideoAudioScheduler (
pipeline = self ,
audio_scheduler = audio_scheduler ,
latent_num_frames = latent_num_frames ,
latent_height = latent_height ,
latent_width = latent_width ,
)
# No coord duplication needed: mixin handles CFG via separate forward calls.
2026-03-05 23:31:48 +08:00
2026-04-03 23:17:49 +08:00
with self . progress_bar ( total = len ( timesteps )) as pbar :
for i , t in enumerate ( timesteps ):
if self . interrupt :
continue
2026-03-05 23:31:48 +08:00
2026-04-03 23:17:49 +08:00
self . _current_timestep = t
2026-03-05 23:31:48 +08:00
2026-04-08 04:44:04 +01:00
latent_model_input = latents . to ( prompt_embeds . dtype )
audio_latent_model_input = audio_latents . to ( prompt_embeds . dtype )
timestep = t . expand ( latent_model_input . shape [ 0 ])
video_timestep = timestep . unsqueeze ( - 1 ) * ( 1 - conditioning_mask )
do_true_cfg = self . do_classifier_free_guidance
2026-03-05 23:31:48 +08:00
2026-04-08 04:44:04 +01:00
positive_kwargs = {
"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" : padded_audio_num_frames ,
"video_coords" : video_coords ,
"audio_coords" : audio_coords ,
"attention_kwargs" : attention_kwargs ,
"return_dict" : False ,
}
negative_kwargs = (
{
** positive_kwargs ,
2026-04-03 23:17:49 +08:00
"encoder_hidden_states" : negative_connector_prompt_embeds ,
"audio_encoder_hidden_states" : negative_connector_audio_prompt_embeds ,
"encoder_attention_mask" : negative_connector_attention_mask ,
"audio_encoder_attention_mask" : negative_connector_attention_mask ,
}
2026-04-08 04:44:04 +01:00
if do_true_cfg
else None
)
2026-03-05 23:31:48 +08:00
2026-04-08 04:44:04 +01:00
noise_pred_video , noise_pred_audio = self . predict_noise_maybe_with_cfg (
do_true_cfg = do_true_cfg ,
true_cfg_scale = guidance_scale ,
positive_kwargs = positive_kwargs ,
negative_kwargs = negative_kwargs ,
cfg_normalize = False ,
)
2026-04-03 23:17:49 +08:00
2026-04-08 04:44:04 +01:00
latents , audio_latents = self . scheduler_step_maybe_with_cfg (
( noise_pred_video , noise_pred_audio ),
( t , t ),
( latents , audio_latents ),
do_true_cfg = do_true_cfg ,
per_request_scheduler = i2v_scheduler ,
)
latents , audio_latents = self . _synchronize_cfg_parallel_step_output (
( latents , audio_latents ),
do_true_cfg = do_true_cfg ,
)
2026-03-05 23:31:48 +08:00
2026-04-03 23:17:49 +08:00
pbar . update ()
2026-03-05 23:31:48 +08:00
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
)
2026-04-08 04:44:04 +01:00
audio_latents = self . _unpad_audio_latents ( audio_latents , original_audio_num_frames )
2026-03-05 23:31:48 +08:00
audio_latents = self . _denormalize_audio_latents (
audio_latents , self . audio_vae . latents_mean , self . audio_vae . latents_std
)
2026-04-08 04:44:04 +01:00
audio_latents = self . _unpack_audio_latents (
audio_latents ,
original_audio_num_frames ,
num_mel_bins = latent_mel_bins ,
)
2026-03-05 23:31:48 +08:00
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 )
if not return_dict :
return DiffusionOutput ( output = ( video , audio ))
return DiffusionOutput ( output = ( video , audio ))
2026-04-03 23:17:49 +08:00
2026-05-08 10:16:49 -04:00
class LTX2ImageToVideoTwoStagesPipeline ( nn . Module , SupportsComponentDiscovery ):
2026-04-03 23:17:49 +08:00
"""LTXImageToVideoTwoStagesPipeline is for two stages image to video generation"""
support_image_input = True
2026-04-22 16:23:00 -04:00
_dit_modules : ClassVar [ list [ str ]] = [ "pipe.transformer" ]
_encoder_modules : ClassVar [ list [ str ]] = [ "pipe.text_encoder" ]
_vae_modules : ClassVar [ list [ str ]] = [ "pipe.vae" , "pipe.audio_vae" ]
2026-04-03 23:17:49 +08:00
def __init__ (
self ,
* ,
od_config : OmniDiffusionConfig ,
prefix : str = "" ,
):
super () . __init__ ()
self . device = get_local_device ()
self . dtype = getattr ( od_config , "dtype" , torch . bfloat16 )
self . model_path = od_config . model
self . distilled = False
# User provided model path may contain '/' in the end and basename function
# will not return the expected directory name, so we need to remove it by normpath
if "distilled" in os . path . basename ( os . path . normpath ( self . model_path )):
self . distilled = True
else :
raise NotImplementedError ( f " { self . model_path } is not supported for { self . __class__ . __name__ } ." )
self . pipe = LTX2ImageToVideoPipeline ( od_config = od_config , prefix = prefix )
self . upsample_pipe = LTX2LatentUpsamplePipeline (
vae = self . pipe . vae ,
od_config = od_config ,
)
self . lora_manager = DiffusionLoRAManager (
pipeline = self . pipe ,
device = self . device ,
dtype = self . dtype ,
max_cached_adapters = od_config . max_cpu_loras ,
)
self . weights_sources = [
DiffusersPipelineLoader . ComponentSource (
model_or_path = od_config . model ,
subfolder = "transformer" ,
revision = None ,
prefix = "pipe.transformer." ,
fall_back_to_pt = True ,
),
]
@torch.no_grad ()
def forward (
self ,
req : OmniDiffusionRequest ,
image : PIL . Image . Image | torch . Tensor | None = None ,
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 ,
guidance_rescale : float = 0.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 ,
):
video_latent , audio_latent = self . pipe (
req = req ,
image = image ,
prompt = prompt ,
negative_prompt = negative_prompt ,
height = height ,
width = width ,
num_frames = num_frames ,
frame_rate = frame_rate ,
num_inference_steps = num_inference_steps ,
sigmas = DISTILLED_SIGMA_VALUES if self . distilled else None ,
timesteps = timesteps ,
guidance_scale = guidance_scale ,
guidance_rescale = guidance_rescale ,
noise_scale = noise_scale ,
num_videos_per_prompt = num_videos_per_prompt ,
generator = generator ,
latents = latents ,
audio_latents = audio_latents ,
prompt_embeds = prompt_embeds ,
negative_prompt_embeds = negative_prompt_embeds ,
prompt_attention_mask = prompt_attention_mask ,
negative_prompt_attention_mask = negative_prompt_attention_mask ,
decode_timestep = decode_timestep ,
decode_noise_scale = decode_noise_scale ,
output_type = "latent" ,
return_dict = return_dict ,
attention_kwargs = attention_kwargs ,
max_sequence_length = max_sequence_length ,
) . output
upscaled_video_latent = self . upsample_pipe (
latents = video_latent ,
output_type = "latent" ,
return_dict = False ,
)[ 0 ]
if not self . distilled :
# Load Stage 2 distilled LoRA
lora_path = f " { self . model_path } /ltx-2-19b-distilled-lora-384.safetensors"
lora_request = LoRARequest (
lora_name = "stage_2_distilled" ,
lora_int_id = 1 ,
lora_path = lora_path ,
)
self . lora_manager . set_active_adapter ( lora_request , lora_scale = 1.0 )
# Change scheduler to use Stage 2 distilled sigmas as is
new_scheduler = FlowMatchEulerDiscreteScheduler . from_config (
self . pipe . scheduler . config ,
use_dynamic_shifting = False ,
shift_terminal = None ,
)
self . pipe . scheduler = new_scheduler
stage_2_req = copy . copy ( req )
stage_2_req . sampling_params = req . sampling_params . clone ()
stage_2_req . sampling_params . num_inference_steps = 3
video , audio = self . pipe (
req = stage_2_req ,
latents = upscaled_video_latent ,
audio_latents = audio_latent ,
prompt = prompt ,
negative_prompt = negative_prompt ,
noise_scale = STAGE_2_DISTILLED_SIGMA_VALUES [ 0 ],
sigmas = STAGE_2_DISTILLED_SIGMA_VALUES ,
guidance_scale = 1.0 ,
generator = generator ,
output_type = "np" ,
return_dict = False ,
) . output
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 )
2026-04-20 06:21:43 -07:00
class LTX2I2VDMD2Pipeline ( DMD2PipelineMixin , LTX2ImageToVideoPipeline ):
"""LTX-2 I2V pipeline for FastGen DMD2-distilled models."""
def __init__ ( self , * , od_config : OmniDiffusionConfig , prefix : str = "" ):
super () . __init__ ( od_config = od_config , prefix = prefix )
self . __init_dmd2__ ()
