# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """LTX-2.3 image-to-video pipeline.""" from __future__ import annotations from typing import Any import PIL.Image import torch 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 from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from .pipeline_ltx2_3 import ( LTX23Pipeline, _LTX23DenoiseContext, _LTX23ForwardContext, _LTX23PromptContext, _LTX23RequestInputs, get_ltx2_post_process_func, ) from .pipeline_ltx2_image2video import LTX2ImageToVideoPipeline, _I2VVideoAudioScheduler class LTX23ImageToVideoPipeline(LTX23Pipeline): """LTX-2.3 image-to-video pipeline. This keeps the LTX-2.3 prompt connector, x0-space CFG, sigma prompt modulation, and audio branch semantics from ``LTX23Pipeline`` while reusing the existing LTX image-conditioning contract: the first video latent frame is encoded from the input image and remains fixed during denoising. """ 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") support_image_input = True _normalize_latents = staticmethod(LTX2ImageToVideoPipeline._normalize_latents) _create_noised_state = staticmethod(LTX2ImageToVideoPipeline._create_noised_state) @staticmethod def _resolve_single_prompt_image(raw_image: Any) -> Any: if isinstance(raw_image, list): if len(raw_image) != 1: raise ValueError( "LTX-2.3 I2V prompt dictionaries support exactly one image per prompt. " "Pass one image per prompt for batched I2V requests." ) return raw_image[0] return raw_image @staticmethod def _resolve_additional_image(additional: dict[str, Any]) -> Any: raw_image = additional.get("preprocessed_image") if raw_image is None: raw_image = additional.get("pixel_values") if raw_image is None: raw_image = additional.get("image") return raw_image 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, 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, torch.Tensor]: """Prepare I2V latents and the first-frame conditioning mask. If caller-provided latents are used without an image, the latents must already represent the full video state including the conditioning first frame. Packed 3D latents are assumed to be in transformer token layout. """ 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: if latents.ndim == 5: 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 = 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 conditioning_mask = self._pack_latents( conditioning_mask, self.transformer_spatial_patch_size, self.transformer_temporal_patch_size, ).squeeze(-1) 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 image is None and latents is None: raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.") super().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, ) _step_video_latents_i2v = LTX2ImageToVideoPipeline._step_video_latents_i2v def _resolve_request_image( self, req: DiffusionRequestBatch, image: PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None, request_inputs: _LTX23RequestInputs, ) -> PIL.Image.Image | torch.Tensor | list[PIL.Image.Image | torch.Tensor] | None: if image is not None or not req.prompts: return image 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 = self._resolve_additional_image(additional) raw_image = self._resolve_single_prompt_image(raw_image) 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) and request_inputs.latents is None: raise ValueError("Image is required for LTX-2.3 I2V generation.") if len(raw_images) == 1: return raw_images[0] if raw_images: return raw_images return image def _check_forward_inputs( self, request_inputs: _LTX23RequestInputs, image: Any | None = None, ) -> None: self.check_inputs( image=image, height=request_inputs.height, width=request_inputs.width, prompt=request_inputs.prompt, latents=request_inputs.latents, prompt_embeds=request_inputs.prompt_embeds, negative_prompt_embeds=request_inputs.negative_prompt_embeds, prompt_attention_mask=request_inputs.prompt_attention_mask, negative_prompt_attention_mask=request_inputs.negative_prompt_attention_mask, ) def _resolve_video_latent_dimensions( self, request_inputs: _LTX23RequestInputs, ) -> tuple[int, int, int]: latent_num_frames = (request_inputs.num_frames - 1) // self.vae_temporal_compression_ratio + 1 latent_height = request_inputs.height // self.vae_spatial_compression_ratio latent_width = request_inputs.width // self.vae_spatial_compression_ratio if request_inputs.latents is not None: if request_inputs.latents.ndim == 5: _, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape elif request_inputs.latents.ndim != 3: raise ValueError( f"Provided `latents` tensor has shape {request_inputs.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]." ) return latent_num_frames, latent_height, latent_width def _prepare_video_latents_stage( self, request_inputs: _LTX23RequestInputs, prompt_context: _LTX23PromptContext, *, device: torch.device, noise_scale: float, image: Any | None = None, ) -> tuple[torch.Tensor, torch.Tensor | None]: if request_inputs.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=request_inputs.height, width=request_inputs.width, ) image = image.to(device=device, dtype=prompt_context.positive_connector_prompt_embeds.dtype) num_channels_latents = self.transformer.config.in_channels return self.prepare_latents( image, prompt_context.batch_size * request_inputs.num_videos_per_prompt, num_channels_latents, request_inputs.height, request_inputs.width, request_inputs.num_frames, noise_scale, torch.float32, device, request_inputs.generator, request_inputs.latents, ) def _make_video_audio_scheduler( self, audio_scheduler: Any, latent_num_frames: int, latent_height: int, latent_width: int, ) -> Any: return _I2VVideoAudioScheduler( self, audio_scheduler, latent_num_frames, latent_height, latent_width, ) def _prepare_denoise_context_for_cfg( self, forward_ctx: _LTX23ForwardContext, denoise_ctx: _LTX23DenoiseContext, ) -> _LTX23DenoiseContext: denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) if denoise_ctx.conditioning_mask is None: raise ValueError("I2V denoising requires a conditioning mask.") if self.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready: denoise_ctx.conditioning_mask_for_model = torch.cat( [denoise_ctx.conditioning_mask, denoise_ctx.conditioning_mask] ) else: denoise_ctx.conditioning_mask_for_model = denoise_ctx.conditioning_mask return denoise_ctx def _denoise_timestep_kwargs( self, ts: torch.Tensor, forward_ctx: _LTX23ForwardContext, denoise_ctx: _LTX23DenoiseContext, ) -> dict[str, torch.Tensor]: conditioning_mask = ( denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model ) if conditioning_mask is None: raise ValueError("I2V denoising requires a conditioning mask.") return { "timestep": ts.unsqueeze(-1) * (1 - conditioning_mask), "audio_timestep": ts, "sigma": ts, } def _step_denoised_latents( self, forward_ctx: _LTX23ForwardContext, denoise_ctx: _LTX23DenoiseContext, noise_pred_video: torch.Tensor, noise_pred_audio: torch.Tensor, t: torch.Tensor, ) -> tuple[torch.Tensor, torch.Tensor]: latents, audio_latents = self.scheduler_step_maybe_with_cfg( (noise_pred_video, noise_pred_audio), (t, t), (denoise_ctx.latents, denoise_ctx.audio_latents), do_true_cfg=self.do_classifier_free_guidance, per_request_scheduler=forward_ctx.video_audio_scheduler, ) return self._synchronize_cfg_parallel_step_output( (latents, audio_latents), do_true_cfg=self.do_classifier_free_guidance, ) @torch.no_grad() def forward( self, req: DiffusionRequestBatch, image: PIL.Image.Image | torch.Tensor | list[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, 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, ) -> list[DiffusionOutput]: request_inputs = self._resolve_request_inputs( req, prompt=prompt, negative_prompt=negative_prompt, height=height, width=width, num_frames=num_frames, frame_rate=frame_rate, num_inference_steps=num_inference_steps, timesteps=timesteps, guidance_scale=guidance_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=output_type, max_sequence_length=max_sequence_length, ) image = self._resolve_request_image(req, image, request_inputs) return self._forward_impl( req, request_inputs, noise_scale=noise_scale, sigmas=sigmas, timesteps=timesteps, attention_kwargs=attention_kwargs, image=image, ) __all__ = [ "LTX23ImageToVideoPipeline", "get_ltx2_post_process_func", ]