Fix griffe warning: add type annotation for **kwargs parameter

Co-authored-by: hsliuustc0106 <222337142+hsliuustc0106@users.noreply.github.com>

docs/.nav.yml

@@ -12,6 +12,7 @@ nav:
     - Offline Inference:
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - LTX-2: user_guide/examples/offline_inference/ltx2.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md

docs/models/supported_models.md

@@ -34,6 +34,7 @@ th {
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
 |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
+|`LTX2Pipeline` | LTX-2 | `Lightricks/LTX-2` |
 
 
 ## List of Supported Models for NPU

docs/user_guide/examples/offline_inference/ltx2.md

@@ -0,0 +1,42 @@
+# LTX-2 Text-To-Video
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/ltx2>.
+
+The `Lightricks/LTX-2` pipeline generates high-quality videos from text prompts using a 19B parameter DiT architecture. It supports up to 4K video generation with synchronized audio.
+
+## Local CLI Usage
+
+```bash
+python text_to_video.py \
+  --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 ltx2_output.mp4
+```
+
+Key arguments:
+
+- `--prompt`: Text description of the video to generate (string).
+- `--height/--width`: Output resolution (defaults 512x768). Dimensions must be divisible by 32.
+- `--num_frames`: Number of frames (should be 8*n+1, e.g., 25, 81, 121). Default is 121.
+- `--guidance_scale`: Classifier-free guidance scale (default: 4.0). Range 3.0-5.0 recommended.
+- `--negative_prompt`: Optional text describing what to avoid in the video.
+- `--num_inference_steps`: Number of denoising steps (default: 40).
+- `--seed`: Random seed for reproducibility (default: 42).
+- `--fps`: Frames per second for the saved MP4 (default: 24).
+- `--output`: Path to save the generated video.
+
+## Example materials
+
+??? abstract "text_to_video.py"
+    ``````py
+    --8<-- "examples/offline_inference/ltx2/text_to_video.py"
+    ``````
+
+??? abstract "text_to_video.md"
+    ``````md
+    --8<-- "examples/offline_inference/ltx2/text_to_video.md"
+    ``````

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").

examples/offline_inference/ltx2/text_to_video.py

@@ -0,0 +1,124 @@
+# 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}', "
+                    f"expected 'image' for video frames output."
+                )
+
+            # 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()

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",
+]

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -0,0 +1,186 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import os
+from typing import Any
+
+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")
+
+        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: str | None = None,
+        negative_prompt: str | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_inference_steps: int | None = None,
+        guidance_scale: float = 4.0,
+        frame_num: int | None = None,
+        output_type: str | None = "np",
+        generator: torch.Generator | None = None,
+        **kwargs: Any,
+    ) -> 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(output=frames)

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"):
+        # For pipelines that wrap diffusers pipelines (like LTX2Pipeline),
-            model.vae.use_slicing = od_config.vae_use_slicing
+        # VAE is accessed via model.pipe.vae
-        if hasattr(model.vae, "use_tiling"):
+        vae = None
-            model.vae.use_tiling = od_config.vae_use_tiling
+        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",
 }