From 99c02d3871a2d4529c0fe477322dae4c4525dbc0 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Mon, 19 Jan 2026 04:31:55 +0000 Subject: [PATCH] Add LTX-2 model support with pipeline, registry, and example Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com> --- .../offline_inference/ltx2/text_to_video.md | 77 +++++++ .../offline_inference/ltx2/text_to_video.py | 123 +++++++++++ vllm_omni/diffusion/models/ltx2/__init__.py | 13 ++ .../diffusion/models/ltx2/pipeline_ltx2.py | 196 ++++++++++++++++++ vllm_omni/diffusion/registry.py | 24 ++- 5 files changed, 429 insertions(+), 4 deletions(-) create mode 100644 examples/offline_inference/ltx2/text_to_video.md create mode 100644 examples/offline_inference/ltx2/text_to_video.py create mode 100644 vllm_omni/diffusion/models/ltx2/__init__.py create mode 100644 vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py diff --git a/examples/offline_inference/ltx2/text_to_video.md b/examples/offline_inference/ltx2/text_to_video.md new file mode 100644 index 000000000..82308f7cb --- /dev/null +++ b/examples/offline_inference/ltx2/text_to_video.md @@ -0,0 +1,77 @@ +# LTX-2 Text-to-Video Generation + +This example demonstrates how to use Lightricks LTX-2 model for text-to-video generation with vLLM-Omni. + +## Model Information + +LTX-2 is a state-of-the-art text-to-video model by Lightricks that can generate high-quality videos from text prompts. It supports: +- Text-to-video generation +- Image-to-video generation (conditioned on an input image) +- High-resolution video generation (up to 4K) +- Synchronized audio-video generation + +For more information, see: +- [Lightricks/LTX-2 on HuggingFace](https://huggingface.co/Lightricks/LTX-2) +- [LTX-2 GitHub Repository](https://github.com/Lightricks/LTX-2) + +## Requirements + +- Python >= 3.10 +- diffusers >= 0.36.0 (already included in vllm-omni dependencies) +- GPU with sufficient VRAM (recommended: >= 24GB for high-quality generation) + +## Usage + +### Basic Text-to-Video Generation + +```bash +python text_to_video.py \ + --model Lightricks/LTX-2 \ + --prompt "A panda riding a bicycle through a forest, cinematic lighting" \ + --height 512 \ + --width 768 \ + --num_frames 121 \ + --num_inference_steps 40 \ + --guidance_scale 4.0 \ + --output output.mp4 +``` + +### Parameters + +- `--model`: Model ID or local path (default: `Lightricks/LTX-2`) +- `--prompt`: Text description of the video to generate +- `--negative_prompt`: Text describing what to avoid in the video +- `--height`: Video height in pixels, must be divisible by 32 (default: 512) +- `--width`: Video width in pixels, must be divisible by 32 (default: 768) +- `--num_frames`: Number of frames, should be 8*n+1 (e.g., 25, 81, 121) (default: 121) +- `--num_inference_steps`: Number of denoising steps (default: 40) +- `--guidance_scale`: Classifier-free guidance scale (default: 4.0) +- `--seed`: Random seed for reproducibility (default: 42) +- `--fps`: Output video frames per second (default: 24) +- `--output`: Output video file path (default: ltx2_output.mp4) + +### Example Prompts + +```bash +# Cinematic scene +python text_to_video.py --prompt "A serene lakeside sunrise with mist over the water, cinematic" + +# Action scene +python text_to_video.py --prompt "A futuristic cityscape with flying cars, neon lights, cyberpunk style" + +# Nature scene +python text_to_video.py --prompt "A waterfall in a lush rainforest, birds flying overhead" +``` + +## Performance Tips + +1. **Resolution**: Start with lower resolutions (512x768) for faster generation. Higher resolutions require more VRAM. +2. **Frames**: More frames take longer to generate. Start with 25 or 81 frames for testing. +3. **Inference Steps**: 40 steps provide good quality. Fewer steps are faster but may reduce quality. +4. **Guidance Scale**: Values between 3.0-5.0 typically work well. Higher values follow the prompt more strictly. + +## Notes + +- The first run will download the model from HuggingFace (approximately 20-30GB). +- Generation time depends on your GPU and the parameters chosen. +- For best results, use descriptive prompts with style keywords (e.g., "cinematic", "high quality", "detailed"). diff --git a/examples/offline_inference/ltx2/text_to_video.py b/examples/offline_inference/ltx2/text_to_video.py new file mode 100644 index 000000000..47bfb476b --- /dev/null +++ b/examples/offline_inference/ltx2/text_to_video.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import argparse +from pathlib import Path + +import numpy as np +import torch + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.utils.platform_utils import detect_device_type + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Generate a video with LTX-2 T2V.") + parser.add_argument( + "--model", + default="Lightricks/LTX-2", + help="LTX-2 model ID or local path.", + ) + parser.add_argument( + "--prompt", + default="A serene lakeside sunrise with mist over the water.", + help="Text prompt for video generation.", + ) + parser.add_argument("--negative_prompt", default="", help="Negative prompt.") + parser.add_argument("--seed", type=int, default=42, help="Random seed.") + parser.add_argument("--guidance_scale", type=float, default=4.0, help="CFG scale.") + parser.add_argument("--height", type=int, default=512, help="Video height (must be divisible by 32).") + parser.add_argument("--width", type=int, default=768, help="Video width (must be divisible by 32).") + parser.add_argument( + "--num_frames", type=int, default=121, help="Number of frames (should be 8*n+1, e.g., 25, 81, 121)." + ) + parser.add_argument("--num_inference_steps", type=int, default=40, help="Sampling steps.") + parser.add_argument("--output", type=str, default="ltx2_output.mp4", help="Path to save the video (mp4).") + parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.") + return parser.parse_args() + + +def main(): + args = parse_args() + device = detect_device_type() + generator = torch.Generator(device=device).manual_seed(args.seed) + + omni = Omni( + model=args.model, + ) + + frames = omni.generate( + args.prompt, + negative_prompt=args.negative_prompt, + height=args.height, + width=args.width, + generator=generator, + guidance_scale=args.guidance_scale, + num_inference_steps=args.num_inference_steps, + num_frames=args.num_frames, + ) + + # Extract video frames from OmniRequestOutput + if isinstance(frames, list) and len(frames) > 0: + first_item = frames[0] + + # Check if it's an OmniRequestOutput + if hasattr(first_item, "final_output_type"): + if first_item.final_output_type != "image": + raise ValueError( + f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation." + ) + + # Pipeline mode: extract from nested request_output + if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output: + if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0: + inner_output = first_item.request_output[0] + if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"): + frames = inner_output.images[0] if inner_output.images else None + if frames is None: + raise ValueError("No video frames found in output.") + # Diffusion mode: use direct images field + elif hasattr(first_item, "images") and first_item.images: + frames = first_item.images + else: + raise ValueError("No video frames found in OmniRequestOutput.") + + output_path = Path(args.output) + output_path.parent.mkdir(parents=True, exist_ok=True) + try: + from diffusers.utils import export_to_video + except ImportError: + raise ImportError("diffusers is required for export_to_video.") + + # frames may be np.ndarray (preferred) or torch.Tensor + # export_to_video expects a list of frames with values in [0, 1] + if isinstance(frames, torch.Tensor): + video_tensor = frames.detach().cpu() + if video_tensor.dim() == 5: + # [B, C, F, H, W] or [B, F, H, W, C] + if video_tensor.shape[1] in (3, 4): + video_tensor = video_tensor[0].permute(1, 2, 3, 0) + else: + video_tensor = video_tensor[0] + elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4): + video_tensor = video_tensor.permute(1, 2, 3, 0) + # If float, assume [-1,1] and normalize to [0,1] + if video_tensor.is_floating_point(): + video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5 + video_array = video_tensor.float().numpy() + else: + video_array = frames + if hasattr(video_array, "shape") and video_array.ndim == 5: + video_array = video_array[0] + + # Convert 4D array (frames, H, W, C) to list of frames for export_to_video + if isinstance(video_array, np.ndarray) and video_array.ndim == 4: + video_array = list(video_array) + + export_to_video(video_array, str(output_path), fps=args.fps) + print(f"Saved generated video to {output_path}") + + +if __name__ == "__main__": + main() diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py new file mode 100644 index 000000000..08e4975e0 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -0,0 +1,13 @@ +"""LTX-2 model components for text-to-video and image-to-video generation.""" + +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( + LTX2Pipeline, + get_ltx2_post_process_func, + get_ltx2_pre_process_func, +) + +__all__ = [ + "LTX2Pipeline", + "get_ltx2_post_process_func", + "get_ltx2_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py new file mode 100644 index 000000000..ba95db235 --- /dev/null +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -0,0 +1,196 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import os +from typing import Optional + +import torch +from diffusers import LTXVideoPipeline +from torch import nn + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.request import OmniDiffusionRequest + + +def get_ltx2_post_process_func( + od_config: OmniDiffusionConfig, +): + """Post-process function for LTX-2: convert video output to numpy array.""" + from diffusers.video_processor import VideoProcessor + + video_processor = VideoProcessor(vae_scale_factor=8) + + def post_process_func( + video: torch.Tensor, + output_type: str = "np", + ): + if output_type == "latent": + return video + return video_processor.postprocess_video(video, output_type=output_type) + + return post_process_func + + +def get_ltx2_pre_process_func( + od_config: OmniDiffusionConfig, +): + """Pre-process function for LTX-2: validate and prepare image inputs for I2V mode.""" + import PIL.Image + + def pre_process_func(requests: list[OmniDiffusionRequest]) -> list[OmniDiffusionRequest]: + for req in requests: + # Load image if path is provided + if req.image_path is not None and req.pil_image is None: + req.pil_image = PIL.Image.open(req.image_path).convert("RGB") + + # For image-to-video, store the image for later use + if req.pil_image is not None: + # LTX-2 will handle image resizing internally + pass + + return requests + + return pre_process_func + + +class LTX2Pipeline(nn.Module): + """ + Wrapper for Lightricks LTX-2 text-to-video pipeline. + + This model uses the diffusers LTXVideoPipeline for efficient video generation. + Supports text-to-video and image-to-video generation. + """ + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.device = get_local_device() + + model = od_config.model + local_files_only = os.path.exists(model) + + # Load the LTX-2 pipeline from diffusers + dtype = getattr(od_config, "dtype", torch.bfloat16) + + self.pipe = LTXVideoPipeline.from_pretrained( + model, + torch_dtype=dtype, + local_files_only=local_files_only, + ).to(self.device) + + # Store configuration + self.output_type = self.od_config.output_type + + # No weights_sources needed since we use the diffusers pipeline directly + self.weights_sources = [] + + @torch.no_grad() + def forward( + self, + req: OmniDiffusionRequest, + prompt: Optional[str] = None, + negative_prompt: Optional[str] = None, + height: Optional[int] = None, + width: Optional[int] = None, + num_inference_steps: Optional[int] = None, + guidance_scale: float = 4.0, + frame_num: Optional[int] = None, + output_type: Optional[str] = "np", + generator: Optional[torch.Generator] = None, + **kwargs, + ) -> DiffusionOutput: + """ + Generate video from text prompt. + + Args: + req: OmniDiffusionRequest containing generation parameters + prompt: Text prompt for video generation + negative_prompt: Negative text prompt + height: Video height (must be divisible by 32) + width: Video width (must be divisible by 32) + num_inference_steps: Number of diffusion steps + guidance_scale: Guidance scale for classifier-free guidance + frame_num: Number of frames to generate + output_type: Output format ("np", "pt", or "pil") + generator: Random generator for reproducibility + **kwargs: Additional arguments passed to the pipeline + + Returns: + DiffusionOutput containing the generated video frames + """ + # Extract parameters from request + prompt = req.prompt if req.prompt is not None else prompt + negative_prompt = req.negative_prompt if req.negative_prompt is not None else negative_prompt + + if prompt is None: + raise ValueError("Prompt is required for LTX-2 video generation.") + + # Set default values for video generation + height = req.height or height or 512 + width = req.width or width or 768 + num_frames = req.num_frames if req.num_frames else frame_num or 121 + num_steps = req.num_inference_steps or num_inference_steps or 40 + + # Ensure dimensions are divisible by 32 (LTX-2 requirement) + height = (height // 32) * 32 + width = (width // 32) * 32 + + # Ensure num_frames follows LTX-2 convention (divisible by 8 + 1) + # e.g., 25, 41, 81, 121, etc. + if (num_frames - 1) % 8 != 0: + num_frames = ((num_frames - 1) // 8) * 8 + 1 + num_frames = max(num_frames, 25) # Minimum 25 frames + + # Prepare generator + if generator is None: + seed = req.seed if hasattr(req, 'seed') and req.seed is not None else 42 + generator = torch.Generator(device=self.device).manual_seed(seed) + + # Prepare image input for I2V mode if available + image = None + if hasattr(req, 'pil_image') and req.pil_image is not None: + image = req.pil_image + + # Call the diffusers pipeline + pipeline_kwargs = { + "prompt": prompt, + "negative_prompt": negative_prompt, + "height": height, + "width": width, + "num_frames": num_frames, + "num_inference_steps": num_steps, + "guidance_scale": guidance_scale, + "generator": generator, + "output_type": output_type, + } + + # Add image if provided (for I2V mode) + if image is not None: + pipeline_kwargs["image"] = image + + # Add any extra kwargs + pipeline_kwargs.update(kwargs) + + # Generate video + output = self.pipe(**pipeline_kwargs) + + # Extract frames from the pipeline output + if hasattr(output, "frames"): + frames = output.frames + elif hasattr(output, "videos"): + frames = output.videos + else: + # Fallback: assume the output itself is the frames + frames = output + + # Return as DiffusionOutput + return DiffusionOutput( + images=frames, + final_output_type="image", + ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 5edd87a82..b8fa5cde3 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -79,6 +79,11 @@ _DIFFUSION_MODELS = { "pipeline_flux2_klein", "Flux2KleinPipeline", ), + "LTX2Pipeline": ( + "ltx2", + "pipeline_ltx2", + "LTX2Pipeline", + ), } @@ -100,10 +105,19 @@ def initialize_model( if model_class is not None: model = model_class(od_config=od_config) # Configure VAE memory optimization settings from config - if hasattr(model.vae, "use_slicing"): - model.vae.use_slicing = od_config.vae_use_slicing - if hasattr(model.vae, "use_tiling"): - model.vae.use_tiling = od_config.vae_use_tiling + # For pipelines that wrap diffusers pipelines (like LTX2Pipeline), + # VAE is accessed via model.pipe.vae + vae = None + if hasattr(model, "vae"): + vae = model.vae + elif hasattr(model, "pipe") and hasattr(model.pipe, "vae"): + vae = model.pipe.vae + + if vae is not None: + if hasattr(vae, "use_slicing"): + vae.use_slicing = od_config.vae_use_slicing + if hasattr(vae, "use_tiling"): + vae.use_tiling = od_config.vae_use_tiling return model else: @@ -127,6 +141,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", "Flux2KleinPipeline": "get_flux2_klein_post_process_func", + "LTX2Pipeline": "get_ltx2_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { @@ -139,6 +154,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { "QwenImageLayeredPipeline": "get_qwen_image_layered_pre_process_func", "WanPipeline": "get_wan22_pre_process_func", "WanImageToVideoPipeline": "get_wan22_i2v_pre_process_func", + "LTX2Pipeline": "get_ltx2_pre_process_func", }