diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md index 6e105741a..dc1d3d6b2 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -3,7 +3,9 @@ Source . -This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API. +This example demonstrates how to generate videos from images using Wan2.2, +LTX-2/LTX-2.3, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's +offline inference API. ## Local CLI Usage @@ -51,13 +53,41 @@ python image_to_video.py \ --output i2v_output.mp4 ``` +### LTX-2.3 + +```bash +python image_to_video.py \ + --model dg845/LTX-2.3-Diffusers \ + --model-class-name LTX23ImageToVideoPipeline \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --guidance-scale 4.0 \ + --num-inference-steps 20 \ + --frame-rate 24 \ + --fps 24 \ + --output ltx23_i2v_output.mp4 +``` + +Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select +the LTX-2.3 image-to-video pipeline. + Key arguments: -- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V). +- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or + LTX-2/LTX-2.3). - `--image`: Path to input image (required). - `--prompt`: Text description of desired motion/animation. -- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16. -- `--num-frames`: Number of frames (default 81). +- `--height/--width`: Output resolution (auto-calculated from image if not set). + Wan dimensions should be multiples of 16; LTX dimensions should be multiples + of 32. +- `--num-frames`: Number of frames (model-specific default; LTX-style models + work best with `8k + 1`). - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE). - `--negative-prompt`: Optional list of artifacts to suppress. - `--boundary-ratio`: Boundary split ratio for two-stage MoE models. @@ -65,6 +95,7 @@ Key arguments: - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints. - `--num-inference-steps`: Number of denoising steps (default 50). - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video). +- `--audio-sample-rate`: fallback audio sample rate for embedded audio. - `--output`: Path to save the generated video. - `--vae-use-slicing`: Enable VAE slicing for memory optimization. - `--vae-use-tiling`: Enable VAE tiling for memory optimization. diff --git a/docs/user_guide/examples/offline_inference/text_to_video.md b/docs/user_guide/examples/offline_inference/text_to_video.md index a09dbfc97..5d89eac34 100644 --- a/docs/user_guide/examples/offline_inference/text_to_video.md +++ b/docs/user_guide/examples/offline_inference/text_to_video.md @@ -12,6 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) | |---|---|---|---|---|---| | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | @@ -50,6 +51,30 @@ python text_to_video.py \ --output ltx2_out.mp4 ``` +### LTX-2.3 + +```bash +python text_to_video.py \ + --model dg845/LTX-2.3-Diffusers \ + --model-class-name LTX23Pipeline \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --num-inference-steps 20 \ + --guidance-scale 4.0 \ + --frame-rate 24 \ + --fps 24 \ + --audio-sample-rate 48000 \ + --output ltx23_t2v_output.mp4 +``` + +Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3 +text-to-video pipeline explicitly. + ### HunyuanVideo-1.5 (480p) ```bash @@ -110,6 +135,8 @@ python text_to_video.py \ ### Common - `--model`: Diffusers model ID or local path. +- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline` + for LTX-2.3 text-to-video. - `--prompt`: text description (string). - `--height/--width`: output resolution. Default depends on model. - `--num-frames`: number of frames. Default depends on model. @@ -124,7 +151,8 @@ python text_to_video.py \ - `--enable-cpu-offload`: enable CPU offloading for diffusion models. - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules. - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). -- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio). +- `--audio-sample-rate`: audio sample rate for embedded audio (when the + pipeline returns audio; LTX-2.3 outputs 48kHz audio). - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). - `--flow-shift`: scheduler flow_shift parameter. diff --git a/examples/offline_inference/image_to_video/README.md b/examples/offline_inference/image_to_video/README.md index 68bddca06..17977bb7a 100644 --- a/examples/offline_inference/image_to_video/README.md +++ b/examples/offline_inference/image_to_video/README.md @@ -1,7 +1,7 @@ # Image-To-Video -This shared example generates videos from images with VACE, Wan2.2, LTX2, -HunyuanVideo-1.5, Cosmos3, and other compatible pipelines. +This shared example generates videos from images with VACE, Wan2.2, +LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines. ## Local CLI Usage @@ -49,6 +49,30 @@ python image_to_video.py \ --output i2v_output.mp4 ``` +### LTX-2.3 + +```bash +python image_to_video.py \ + --model dg845/LTX-2.3-Diffusers \ + --model-class-name LTX23ImageToVideoPipeline \ + --image cherry_blossom.jpg \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --guidance-scale 4.0 \ + --num-inference-steps 20 \ + --frame-rate 24 \ + --fps 24 \ + --output ltx23_i2v_output.mp4 +``` + +Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select +the LTX-2.3 image-to-video pipeline. + ### Cosmos3 ```bash @@ -66,15 +90,19 @@ python image_to_video.py \ Key arguments: -- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V). +- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, + LTX-2/LTX-2.3, Cosmos3, or VACE). - `--image`: Path to the first-frame or source image. - `--last-image`: Optional last-frame condition for models such as VACE. - `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved. - `--reference-image`: Optional reference image; repeat it to provide multiple references. - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3. - `--prompt`: Text description of desired motion/animation. -- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16. -- `--num-frames`: Number of frames (default: model-specific — Wan 81, LTX2 121, Cosmos3 189). +- `--height/--width`: Output resolution (auto-calculated from image if not set). + Wan dimensions should be multiples of 16; LTX dimensions should be multiples + of 32. +- `--num-frames`: Number of frames (model-specific default; LTX-style models + work best with `8k + 1`; Cosmos3 defaults to 189). - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE). - `--negative-prompt`: Optional list of artifacts to suppress. - `--boundary-ratio`: Boundary split ratio for two-stage MoE models. @@ -82,6 +110,7 @@ Key arguments: - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints. - `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35). - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video). +- `--audio-sample-rate`: fallback audio sample rate for embedded audio. - `--output`: Path to save the generated video. - `--vae-use-slicing`: Enable VAE slicing for memory optimization. - `--vae-use-tiling`: Enable VAE tiling for memory optimization. diff --git a/examples/offline_inference/image_to_video/image_to_video.py b/examples/offline_inference/image_to_video/image_to_video.py index cc8f1ff2a..1a0691ebb 100644 --- a/examples/offline_inference/image_to_video/image_to_video.py +++ b/examples/offline_inference/image_to_video/image_to_video.py @@ -2,7 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """ -Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE. +Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2/LTX-2.3, +HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE. Supports: - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder @@ -27,6 +28,13 @@ Usage: --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \ --frame-rate 24 --fps 24 --output ltx2_i2v.mp4 + # LTX-2.3 image-to-video + python image_to_video.py --model dg845/LTX-2.3-Diffusers \ + --model-class-name LTX23ImageToVideoPipeline \ + --image input.jpg --prompt "A cinematic dolly shot of a boat" \ + --height 384 --width 512 --num-frames 25 --num-inference-steps 20 \ + --frame-rate 24 --fps 24 --output ltx23_i2v.mp4 + # HunyuanVideo-1.5 I2V (480p) python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \ --image input.jpg --prompt "A cat playing with yarn" \ @@ -79,17 +87,20 @@ parse_profiler_config = functools.partial(parse_json_object, flag_name="--profil def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( - description="Generate a video from one or more images (Wan2.2, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE)." + description=( + "Generate a video from one or more images " + "(Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE)." + ) ) parser.add_argument( "--model", default="Wan-AI/Wan2.2-I2V-A14B-Diffusers", - help="Diffusers I2V model ID or local path (Wan2.2, HunyuanVideo-1.5, or Wan2.1 VACE).", + help="Diffusers I2V model ID or local path (Wan2.2, LTX2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, or Wan2.1 VACE).", ) parser.add_argument( "--model-class-name", default=None, - help="Override model class name (e.g., LTX2ImageToVideoPipeline).", + help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).", ) parser.add_argument("--image", help="Path to the first-frame or source image.") parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).") @@ -113,7 +124,10 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).") parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.") parser.add_argument( - "--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific." + "--num-inference-steps", + type=int, + default=None, + help="Sampling steps. Default: model-specific.", ) parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.") parser.add_argument( @@ -312,10 +326,8 @@ def main(): generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed) model_name = str(args.model).lower() if args.model is not None else "" model_class_name = args.model_class_name - is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower()) + is_ltx2 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"} is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower()) - if model_class_name is None and is_ltx2: - model_class_name = "LTX2ImageToVideoPipeline" image = PIL.Image.open(args.image).convert("RGB") if args.image else None last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None @@ -340,7 +352,15 @@ def main(): 16, ) elif is_ltx2: - d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 24, 4.0, 121, 40, 5.0, 512 * 768, 32 + d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = ( + 24, + 4.0, + 121, + 40, + 5.0, + 512 * 768, + 32, + ) else: # Wan2.2 / HunyuanVideo-1.5 d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16 @@ -462,7 +482,7 @@ def main(): f" tensor_parallel_size={args.tensor_parallel_size}, vae_patch_parallel_size={args.vae_patch_parallel_size}," f" pipeline_parallel_size={args.pipeline_parallel_size}" ) - print(f" Video size: {args.width}x{args.height}") + print(f" Video size: {width}x{height}") print(f"{'=' * 60}\n") prompt_dict = build_image_to_video_prompt( @@ -509,6 +529,7 @@ def main(): print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)") audio = None + audio_sample_rate = args.audio_sample_rate if isinstance(frames, list): frames = frames[0] if frames else None @@ -519,11 +540,13 @@ def main(): ) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] + audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) if frames.is_pipeline_output and frames.request_output is not None: inner_output = frames.request_output if isinstance(inner_output, OmniRequestOutput): if inner_output.multimodal_output and "audio" in inner_output.multimodal_output: audio = inner_output.multimodal_output["audio"] + audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate) frames = inner_output if isinstance(frames, OmniRequestOutput): if frames.images: @@ -531,6 +554,7 @@ def main(): frames, audio = frames.images[0] elif len(frames.images) == 1 and isinstance(frames.images[0], dict): audio = frames.images[0].get("audio") + audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate) frames = frames.images[0].get("frames") or frames.images[0].get("video") else: frames = frames.images @@ -543,6 +567,7 @@ def main(): frames, audio = first_item elif isinstance(first_item, dict): audio = first_item.get("audio") + audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate) frames = first_item.get("frames") or first_item.get("video") elif isinstance(first_item, list): frames = first_item @@ -551,6 +576,7 @@ def main(): frames, audio = frames elif isinstance(frames, dict): audio = frames.get("audio") + audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate) frames = frames.get("frames") or frames.get("video") if frames is None: @@ -669,7 +695,7 @@ def main(): frames_u8, audio_np, fps=float(fps), - audio_sample_rate=args.audio_sample_rate, + audio_sample_rate=audio_sample_rate, ) with open(str(output_path), "wb") as f: f.write(video_bytes) diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index a2634d5ae..4c2b5a83a 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -8,6 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod |---|---|---|---|---|---| | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) | | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB | +| `dg845/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB | | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required | | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) | @@ -70,6 +71,30 @@ python text_to_video.py \ --output ltx2_out.mp4 ``` +### LTX-2.3 + +```bash +python text_to_video.py \ + --model dg845/LTX-2.3-Diffusers \ + --model-class-name LTX23Pipeline \ + --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \ + --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \ + --height 384 \ + --width 512 \ + --num-frames 25 \ + --num-inference-steps 20 \ + --guidance-scale 4.0 \ + --frame-rate 24 \ + --fps 24 \ + --audio-sample-rate 48000 \ + --output ltx23_t2v_output.mp4 +``` + +Use a Diffusers-format checkpoint such as `dg845/LTX-2.3-Diffusers`; the +upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by +this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3 +text-to-video pipeline explicitly. + ### HunyuanVideo-1.5 (480p) ```bash @@ -187,6 +212,8 @@ python text_to_video.py \ ### Common - `--model`: Diffusers model ID or local path. +- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline` + for LTX-2.3 text-to-video. - `--prompt`: text description (string). - `--height/--width`: output resolution. Default depends on model. - `--num-frames`: number of frames. Default depends on model. @@ -202,7 +229,8 @@ python text_to_video.py \ - `--enable-cpu-offload`: enable CPU offloading for diffusion models. - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules. - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2). -- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio). +- `--audio-sample-rate`: audio sample rate for embedded audio (when the + pipeline returns audio; LTX-2.3 outputs 48kHz audio). - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF). - `--flow-shift`: scheduler flow_shift parameter. - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above). diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index 1323985e5..200515a83 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -455,6 +455,7 @@ def main(): print(f"Worker peak GPU memory (reserved): {peak_mb:.2f} MiB ({peak_mb / 1024:.2f} GiB)") audio = None + audio_sample_rate = args.audio_sample_rate if isinstance(frames, list): frames = frames[0] if frames else None @@ -465,11 +466,13 @@ def main(): ) if frames.multimodal_output and "audio" in frames.multimodal_output: audio = frames.multimodal_output["audio"] + audio_sample_rate = frames.multimodal_output.get("audio_sample_rate", audio_sample_rate) if frames.is_pipeline_output and frames.request_output is not None: inner_output = frames.request_output if isinstance(inner_output, OmniRequestOutput): if inner_output.multimodal_output and "audio" in inner_output.multimodal_output: audio = inner_output.multimodal_output["audio"] + audio_sample_rate = inner_output.multimodal_output.get("audio_sample_rate", audio_sample_rate) frames = inner_output if isinstance(frames, OmniRequestOutput): if frames.images: @@ -477,6 +480,7 @@ def main(): frames, audio = frames.images[0] elif len(frames.images) == 1 and isinstance(frames.images[0], dict): audio = frames.images[0].get("audio") + audio_sample_rate = frames.images[0].get("audio_sample_rate", audio_sample_rate) frames = frames.images[0].get("frames") or frames.images[0].get("video") else: frames = frames.images @@ -489,6 +493,7 @@ def main(): frames, audio = first_item elif isinstance(first_item, dict): audio = first_item.get("audio") + audio_sample_rate = first_item.get("audio_sample_rate", audio_sample_rate) frames = first_item.get("frames") or first_item.get("video") elif isinstance(first_item, list): frames = first_item @@ -497,6 +502,7 @@ def main(): frames, audio = frames elif isinstance(frames, dict): audio = frames.get("audio") + audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate) frames = frames.get("frames") or frames.get("video") if frames is None: @@ -611,7 +617,7 @@ def main(): frames_u8, audio_np, fps=float(args.fps), - audio_sample_rate=args.audio_sample_rate, + audio_sample_rate=audio_sample_rate, ) with open(str(output_path), "wb") as f: f.write(video_bytes) diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md index 269ea6425..c83a14883 100644 --- a/recipes/LTX/LTX-2.3.md +++ b/recipes/LTX/LTX-2.3.md @@ -6,21 +6,23 @@ - Vendor: Lightricks - Model: `dg845/LTX-2.3-Diffusers` -- Task: Text-to-video with synchronized audio generation +- Task: Text-to-video and image-to-video with synchronized audio generation - Mode: Online serving (pure diffusion) - Maintainer: @oglok ## When to use this recipe -Use this recipe when you want to serve LTX-2.3 for text-to-video generation -with audio. The model generates videos up to 20+ seconds at 768x512 resolution -with 48kHz audio, all from a single text prompt. Requires a GPU with at least -96GB VRAM due to the 22B parameter transformer (~44GB weights) plus text +Use this recipe when you want to serve LTX-2.3 for text-to-video or +image-to-video generation with audio. The model generates videos up to 20+ +seconds at 768x512 resolution with 48kHz audio from text prompts, and can +condition video generation on an initial image. Start with a 96GB-class GPU for +validation because the model combines a 22B parameter transformer with text encoder, VAE, and vocoder components. ## References -- Model: +- Upstream raw checkpoints: +- Diffusers-format checkpoint: - Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`) ## Serving @@ -34,7 +36,16 @@ vllm serve dg845/LTX-2.3-Diffusers \ --stage-init-timeout 600 ``` -### Verification +For image-to-video: + +```bash +vllm serve dg845/LTX-2.3-Diffusers \ + --omni \ + --model-class-name LTX23ImageToVideoPipeline \ + --stage-init-timeout 600 +``` + +### T2V Verification ```bash # Health check @@ -73,19 +84,48 @@ curl -X POST http://localhost:8000/v1/videos \ -F "guidance_scale=4.0" ``` +### I2V Verification + +Start the server with `--model-class-name LTX23ImageToVideoPipeline`, then pass +an image reference. The I2V pipeline requires `input_reference` or +`image_reference` because the first frame is the conditioning input. + +```bash +curl http://localhost:8000/health + +curl -X POST http://localhost:8000/v1/videos/sync \ + -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \ + -F "input_reference=@/absolute/path/to/reference.png" \ + -F "model=dg845/LTX-2.3-Diffusers" \ + -F "num_frames=81" \ + -F "fps=24" \ + -F "size=768x512" \ + -F "num_inference_steps=30" \ + -F "guidance_scale=4.0" \ + -F "seed=42" \ + -o ltx23_i2v.mp4 +``` + ### Notes -- Memory usage: Model loads at ~36 GiB, peaks at ~62 GiB during inference +- Memory usage: run a latest-head serving smoke or DFX benchmark on the target + GPU before publishing load/peak VRAM numbers. +- Checkpoint format: use a Diffusers-format checkpoint such as + `dg845/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships + raw safetensors and does not contain the subfolder configs required by this + pipeline loader. - Key flags: - - `--stage-init-timeout 600`: Required for the initial `torch.compile` warmup (~90-140 seconds on first request) - - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 pipeline (not LTX-2) + - `--stage-init-timeout 600`: Gives the initial `torch.compile` warmup enough time on large checkpoints + - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 text-to-video pipeline (not LTX-2) + - `--model-class-name LTX23ImageToVideoPipeline`: Selects the LTX-2.3 image-to-video pipeline + - `input_reference=@...`: Uploads the I2V reference image directly - Audio: 48kHz AAC via BWE vocoder, automatically synced with video - CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed - Supported resolutions: 768x512, 512x384 (must be divisible by 32) - Frame rate: 24 fps - Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1) - Known limitations: - - No image-to-video support yet (LTX23ImageToVideoPipeline is a placeholder) + - LTX-2.3 I2V currently supports first-frame image conditioning, matching the LTX image-conditioning contract. - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git) ## Hardware Support @@ -99,13 +139,16 @@ curl -X POST http://localhost:8000/v1/videos \ - OS: Ubuntu 22.04 - Python: 3.10+ - Driver / runtime: CUDA 13.0, Driver 580.126.09 -- vLLM version: 0.19.x -- vLLM-Omni version: 0.19.x +- vLLM version: Match the PR or release checkout being validated +- vLLM-Omni version or commit: Use the commit you are deploying from -### Validated configurations +### Benchmarking -| Duration | Frames | Resolution | Steps | Guidance | Inference Time | Peak VRAM | -|----------|--------|------------|-------|----------|----------------|-----------| -| 3s | 81 | 768x512 | 30 | 4.0 | ~110s | ~62 GB | -| 10s | 241 | 768x512 | 30 | 4.0 | ~130s | ~62 GB | -| 20s | 481 | 768x512 | 30 | 4.0 | ~420s | ~62 GB | +Do not copy latency or VRAM numbers between commits or machines. Record them +only after running a latest-head sweep on the target hardware. + +For formal PR benchmarking, reuse +`tests/dfx/perf/tests/test_ltx2_3_vllm_omni.json` with +`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. That config captures the +single-device eager baseline and CFG-parallel=2 case for a small 384x512, +25-frame, 20-step workload. diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py new file mode 100644 index 000000000..3e592f14b --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_3_forward_parallel.py @@ -0,0 +1,436 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for LTX-2.3 forward-parallel behavior.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestCFGParallelIntegration: + def test_ltx23_pipeline_has_cfg_parallel_mixin(self): + """LTX23Pipeline must use the shared CFG parallel implementation.""" + from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + assert issubclass(LTX23Pipeline, CFGParallelMixin) + + def test_ltx2_transformer_has_dit_cache_config(self): + """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG.""" + from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig + from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel + + adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config") + assert isinstance(adapter_config, CacheDiTAdapterConfig) + assert adapter_config.has_separate_cfg + + +class TestCFGParallelHelpers: + """Test LTX-2.3 CFG helper math without loading model weights.""" + + def test_combine_cfg_noise_matches_x0_space_formula(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + pipe = object.__new__(LTX23Pipeline) + video_sample = torch.tensor([[[1.0, -2.0]]]) + audio_sample = torch.tensor([[[0.5, 3.0]]]) + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2]]]) + audio_neg = torch.tensor([[[0.1, 0.4]]]) + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + scale = 4.0 + + video_combined, audio_combined = pipe.combine_cfg_noise( + (video_pos, audio_pos), + (video_neg, audio_neg), + scale, + video_latents=video_sample, + audio_latents=audio_sample, + video_sigma=video_sigma, + audio_sigma=audio_sigma, + ) + + x0_video_cond = video_sample - video_pos * video_sigma + x0_video_uncond = video_sample - video_neg * video_sigma + x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond) + expected_video = (video_sample - x0_video_guided) / video_sigma + + x0_audio_cond = audio_sample - audio_pos * audio_sigma + x0_audio_uncond = audio_sample - audio_neg * audio_sigma + x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond) + expected_audio = (audio_sample - x0_audio_guided) / audio_sigma + assert torch.allclose(video_combined, expected_video) + assert torch.allclose(audio_combined, expected_audio) + + def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 + + pipe = object.__new__(ltx23.LTX23Pipeline) + video_sample = torch.tensor([[[1.0, -2.0]]]) + audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]]) + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]]) + audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]]) + video_sigma = torch.tensor(0.25) + audio_sigma = torch.tensor(0.5) + scale = 4.0 + + class FakeCfgGroup: + def all_gather(self, tensor, separate_tensors=True): + assert separate_tensors + if tensor.shape == video_pos.shape: + return [video_pos, video_neg] + return [audio_pos, audio_neg] + + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) + + expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg( + video_sample, + video_pos, + video_neg, + video_sigma, + scale, + ) + expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg( + audio_sample, + audio_pos, + audio_neg, + audio_sigma, + scale, + ) + + for rank, expected_branch in ((0, "positive"), (1, "negative")): + calls = [] + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank) + + def fake_predict_noise(**kwargs): + calls.append(kwargs["branch"]) + if kwargs["branch"] == "positive": + return video_pos, audio_pos + return video_neg, audio_neg + + object.__setattr__(pipe, "predict_noise", fake_predict_noise) + video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg( + true_cfg_scale=scale, + positive_kwargs={"branch": "positive"}, + negative_kwargs={"branch": "negative"}, + cfg_normalize=False, + video_latents=video_sample, + audio_latents=audio_sample, + video_sigma=video_sigma, + audio_sigma=audio_sigma, + ) + + assert calls == [expected_branch] + torch.testing.assert_close(video_combined, expected_video) + torch.testing.assert_close(audio_combined, expected_audio) + + assert "_cfg_video_latents" not in pipe.__dict__ + assert "_cfg_audio_latents" not in pipe.__dict__ + + +class TestCFGParallelForwardPath: + """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" + + def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = object.__new__(ltx23.LTX23Pipeline) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 4 + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1) + + class StopAtEncodePromptError(Exception): + pass + + captured = {} + + def fake_encode_prompt(**kwargs): + captured.update(kwargs) + raise StopAtEncodePromptError + + object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) + + prompt_embeds_a = torch.zeros(2, 3) + prompt_embeds_b = torch.ones(2, 3) + negative_prompt_embeds_a = torch.full((2, 3), 2.0) + negative_prompt_embeds_b = torch.full((2, 3), 3.0) + prompt_attention_mask_a = torch.tensor([True, True]) + prompt_attention_mask_b = torch.tensor([True, False]) + negative_attention_mask_a = torch.tensor([False, True]) + negative_attention_mask_b = torch.tensor([False, False]) + + requests = [ + OmniDiffusionRequest( + prompt={ + "prompt": "prompt-a", + "negative_prompt": "negative-a", + "prompt_embeds": prompt_embeds_a, + "negative_prompt_embeds": negative_prompt_embeds_a, + "prompt_attention_mask": prompt_attention_mask_a, + "negative_prompt_attention_mask": negative_attention_mask_a, + }, + sampling_params=OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=1, + frame_rate=1.0, + num_inference_steps=2, + ), + request_id="ltx23-prompt-local-a", + ), + OmniDiffusionRequest( + prompt={ + "prompt": "prompt-b", + "negative_prompt": "negative-b", + "prompt_embeds": prompt_embeds_b, + "negative_prompt_embeds": negative_prompt_embeds_b, + "attention_mask": prompt_attention_mask_b, + "negative_attention_mask": negative_attention_mask_b, + }, + sampling_params=OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=1, + frame_rate=1.0, + num_inference_steps=2, + ), + request_id="ltx23-prompt-local-b", + ), + ] + + with pytest.raises(StopAtEncodePromptError): + pipe.forward(DiffusionRequestBatch(requests=requests)) + + assert captured["prompt"] is None + assert captured["negative_prompt"] is None + torch.testing.assert_close( + captured["prompt_embeds"], + torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + captured["negative_prompt_embeds"], + torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0), + ) + torch.testing.assert_close( + captured["prompt_attention_mask"], + torch.stack([prompt_attention_mask_a, prompt_attention_mask_b], dim=0), + ) + torch.testing.assert_close( + captured["negative_prompt_attention_mask"], + torch.stack([negative_attention_mask_a, negative_attention_mask_b], dim=0), + ) + + @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)]) + @pytest.mark.parametrize( + ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"), + [(1.0, 1, 1.0), (None, 24, 24.0)], + ) + def test_forward_cfg_parallel_steps_video_and_audio_scheduler( + self, + monkeypatch, + cfg_rank, + expected_prompt_value, + frame_rate_input, + audio_sampling_rate, + expected_frame_rate, + ): + from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = object.__new__(ltx23.LTX23Pipeline) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 1 + pipe.vae_spatial_compression_ratio = 32 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.audio_sampling_rate = audio_sampling_rate + pipe.audio_hop_length = 1 + pipe.audio_vae_temporal_compression_ratio = 1 + pipe.audio_vae_mel_compression_ratio = 1 + pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)) + pipe.tokenizer = SimpleNamespace(padding_side="left") + pipe.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + pipe.audio_vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(mel_bins=2, latent_channels=1), + ) + + video_pos = torch.tensor([[[0.2, -0.3]]]) + video_neg = torch.tensor([[[-0.4, 0.1]]]) + audio_pos = torch.tensor([[[0.7, -0.2]]]) + audio_neg = torch.tensor([[[0.1, 0.4]]]) + + class FakeCfgGroup: + def all_gather(self, tensor, separate_tensors=True): + assert separate_tensors + if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg): + return [video_pos, video_neg] + if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg): + return [audio_pos, audio_neg] + raise AssertionError(f"Unexpected gathered tensor: {tensor}") + + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) + monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank) + monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) + + def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None): + scheduler.sigmas = torch.tensor([0.25, 0.25], device=device) + return torch.tensor([1.0, 0.5], device=device), 2 + + monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps) + + class FakeScheduler: + def __init__(self, name="video", calls=None): + self.name = name + self.calls = [] if calls is None else calls + self.config = { + "max_image_seq_len": 4096, + "base_image_seq_len": 1024, + "base_shift": 0.95, + "max_shift": 2.05, + } + self.sigmas = torch.tensor([0.25, 0.25]) + + def __deepcopy__(self, memo): + return FakeScheduler("audio", self.calls) + + def step(self, noise_pred, t, latents, return_dict=False, generator=None): + self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone())) + return (latents - noise_pred,) + + class FakeConnectors: + def to(self, device): + return self + + def __call__(self, prompt_embeds, prompt_attention_mask, padding_side): + assert padding_side == "left" + assert prompt_embeds.shape[0] == 2 + return prompt_embeds, prompt_embeds, prompt_attention_mask + + rope_video_fps: list[float] = [] + + class FakeRope: + def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps): + rope_video_fps.append(fps) + return torch.zeros(batch_size, num_frames * height * width, 3, device=device) + + def prepare_audio_coords(self, batch_size, num_frames, device): + return torch.zeros(batch_size, num_frames, 1, device=device) + + class FakeTransformer: + def __init__(self): + self.config = SimpleNamespace(in_channels=2) + self.rope = FakeRope() + self.audio_rope = FakeRope() + self.calls = [] + + def __call__(self, **kwargs): + self.calls.append(kwargs) + expected_prompt = torch.full((1, 1, 1), expected_prompt_value) + torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt) + torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt) + assert kwargs["hidden_states"].shape == (1, 1, 2) + assert kwargs["audio_hidden_states"].shape == (1, 1, 2) + if cfg_rank == 0: + return video_pos, audio_pos + return video_neg, audio_neg + + class DummyProgress: + def __enter__(self): + return self + + def __exit__(self, exc_type, exc, tb): + return False + + def update(self): + pass + + pipe.scheduler = FakeScheduler() + pipe.connectors = FakeConnectors() + pipe.transformer = FakeTransformer() + object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress()) + + def fake_encode_prompt(**kwargs): + return ( + torch.ones(1, 1, 1), + torch.ones(1, 1, dtype=torch.bool), + torch.zeros(1, 1, 1), + torch.ones(1, 1, dtype=torch.bool), + ) + + object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) + + video_latents = torch.tensor([[[1.0, -2.0]]]) + audio_latents = torch.tensor([[[0.5, 3.0]]]) + req = OmniDiffusionRequest( + prompt={"prompt": "prompt", "negative_prompt": "negative"}, + sampling_params=OmniDiffusionSamplingParams( + height=32, + width=32, + num_frames=1, + frame_rate=frame_rate_input, + num_inference_steps=2, + guidance_scale=4.0, + latents=video_latents, + audio_latents=audio_latents, + output_type="latent", + ), + request_id="ltx23-cfg-parallel-forward-test", + ) + + output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0] + + expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( + video_latents, + video_pos, + video_neg, + pipe.scheduler.sigmas[0], + 4.0, + ) + expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( + audio_latents, + audio_pos, + audio_neg, + pipe.scheduler.sigmas[0], + 4.0, + ) + scheduler_call_names = [call[0] for call in pipe.scheduler.calls] + assert scheduler_call_names == ["video", "audio", "video", "audio"] + assert len(pipe.transformer.calls) == 2 + torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise) + torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise) + torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise) + torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise) + + video_out, audio_out = output.output + torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1)) + torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2)) + + # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve + # to the model's own 24.0 default, not crash on None; a provided rate is passed through. + assert rope_video_fps + assert all(fps == expected_frame_rate for fps in rope_video_fps) diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py new file mode 100644 index 000000000..9b883ff34 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_3_image2video_pipeline.py @@ -0,0 +1,244 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for the LTX-2.3 image-to-video pipeline.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_ltx23_request_pipe(cls): + pipe = object.__new__(cls) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 99 + return pipe + + +class TestLTX23ImageToVideoForwardStages: + def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline) + image = torch.zeros(3, 8, 8) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "make the image move", + "negative_prompt": "jitter", + "multi_modal_data": {"image": image}, + }, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + num_inference_steps=2, + ), + request_id="ltx23-i2v-forward-stage-delegation", + ) + ] + ) + seen = {} + + def fake_forward_impl(req_arg, request_inputs, **kwargs): + seen["req"] = req_arg + seen["request_inputs"] = request_inputs + seen["kwargs"] = kwargs + return ["i2v-delegated"] + + object.__setattr__(pipe, "_forward_impl", fake_forward_impl) + + output = pipe.forward(req, noise_scale=0.5) + + assert output == ["i2v-delegated"] + assert seen["req"] is req + assert seen["request_inputs"].prompt == ["make the image move"] + assert seen["request_inputs"].negative_prompt == ["jitter"] + assert seen["kwargs"]["image"] is image + assert seen["kwargs"]["noise_scale"] == 0.5 + + def test_denoise_timestep_kwargs_masks_video_only(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + ts = torch.tensor([2.0, 4.0]) + denoise_ctx = SimpleNamespace( + conditioning_mask=torch.tensor([[1.0, 0.0]]), + conditioning_mask_for_model=torch.tensor([[1.0, 0.0], [0.0, 1.0]]), + ) + + kwargs = pipe._denoise_timestep_kwargs( + ts, + SimpleNamespace(cfg_parallel_ready=False), + denoise_ctx, + ) + + torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]])) + torch.testing.assert_close(kwargs["audio_timestep"], ts) + torch.testing.assert_close(kwargs["sigma"], ts) + + +class TestLTX23ImageToVideoPipeline: + def test_ltx23_i2v_pipeline_reuses_ltx23_semantics(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + assert issubclass(LTX23ImageToVideoPipeline, LTX23Pipeline) + assert LTX23ImageToVideoPipeline.support_image_input is True + + def test_ltx23_i2v_rejects_multi_image_prompt_list(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + image = object() + + assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image + with pytest.raises(ValueError, match="exactly one image per prompt"): + LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()]) + + def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + image = torch.zeros(1, 3, 4, 4) + additional = { + "preprocessed_image": None, + "pixel_values": image, + "image": torch.ones_like(image), + } + + assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image + + def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + torch.nn.Module.__init__(pipe) + pipe.vae_spatial_compression_ratio = 1 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + + def fake_randn_tensor(shape, generator=None, device=None, dtype=None): + raise AssertionError("packed I2V latents should not be noised") + + monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor) + + latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]]) + + out, conditioning_mask = pipe.prepare_latents( + image=None, + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=3, + noise_scale=1.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) + torch.testing.assert_close(out, latents) + + def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx2 + import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video as ltx23_i2v + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + torch.nn.Module.__init__(pipe) + pipe.vae_spatial_compression_ratio = 1 + pipe.vae_temporal_compression_ratio = 1 + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + pipe.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + + def fake_randn_tensor(shape, generator=None, device=None, dtype=None): + return torch.ones(shape, device=device, dtype=dtype) + + monkeypatch.setattr(ltx23_i2v, "randn_tensor", fake_randn_tensor) + monkeypatch.setattr(ltx2, "randn_tensor", fake_randn_tensor) + + latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]]) + + out, conditioning_mask = pipe.prepare_latents( + image=None, + batch_size=1, + num_channels_latents=2, + height=1, + width=1, + num_frames=3, + noise_scale=1.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=latents, + ) + + torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]])) + torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]])) + + def test_ltx23_i2v_video_step_preserves_conditioning_frame(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + pipe = object.__new__(LTX23ImageToVideoPipeline) + torch.nn.Module.__init__(pipe) + pipe.transformer_spatial_patch_size = 1 + pipe.transformer_temporal_patch_size = 1 + + class FakeScheduler: + def step(self, noise_pred, t, latents, return_dict=False): + return (latents + noise_pred + t,) + + pipe.scheduler = FakeScheduler() + latents = torch.tensor([[[1.0], [2.0], [3.0]]]) + noise_pred = torch.full_like(latents, 10.0) + + out = pipe._step_video_latents_i2v( + noise_pred, + latents, + torch.tensor(0.5), + latent_num_frames=3, + latent_height=1, + latent_width=1, + ) + + torch.testing.assert_close(out[:, :1], latents[:, :1]) + torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5) + + +class TestLTX23ImageToVideoModule: + """Test that pipeline_ltx2_3_image2video.py exposes I2V entry points.""" + + def test_i2v_classes_importable(self): + """I2V classes must be importable from the implementation module.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + + assert LTX23ImageToVideoPipeline is not None + + def test_post_process_func_importable(self): + """get_ltx2_post_process_func must be importable for registry lookup.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func + + assert callable(get_ltx2_post_process_func) + + def test_i2v_class_matches_package_export(self): + """Package export must point at the I2V implementation module.""" + from vllm_omni.diffusion.models.ltx2 import LTX23ImageToVideoPipeline as PackageExported + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import ( + LTX23ImageToVideoPipeline as ModuleExported, + ) + + assert PackageExported is ModuleExported diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py index 96051416b..25aae4510 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_3_pipeline.py @@ -38,6 +38,279 @@ def _make_ltx23_pipeline(sequence_parallel_size: int = 1): return pipeline +def _make_ltx23_request_pipe(cls): + pipe = object.__new__(cls) + torch.nn.Module.__init__(pipe) + pipe.device = torch.device("cpu") + pipe.tokenizer_max_length = 99 + return pipe + + +def _resolve_request_inputs_for_test(pipe, req): + return pipe._resolve_request_inputs( + req, + prompt=None, + negative_prompt=None, + height=None, + width=None, + num_frames=None, + frame_rate=None, + num_inference_steps=None, + timesteps=None, + guidance_scale=4.0, + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=None, + ) + + +class TestLTX23RequestParsing: + def test_t2v_and_i2v_share_request_input_resolution(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + prompt_embeds = torch.tensor([[1.0, 2.0]]) + negative_prompt_embeds = torch.tensor([[3.0, 4.0]]) + prompt_attention_mask = torch.tensor([True, False]) + negative_attention_mask = torch.tensor([False, True]) + video_latents = torch.ones(1, 2, 3) + audio_latents = torch.zeros(1, 2, 3) + generator = torch.Generator().manual_seed(123) + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "shared prompt", + "negative_prompt": "shared negative", + "additional_information": { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "attention_mask": prompt_attention_mask, + "negative_attention_mask": negative_attention_mask, + }, + }, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + frame_rate=12.5, + num_inference_steps=1, + num_outputs_per_prompt=2, + guidance_scale=6.0, + generator=generator, + latents=video_latents, + extra_args={"audio_latents": audio_latents}, + decode_timestep=[0.1], + decode_noise_scale=[0.2], + output_type="latent", + max_sequence_length=17, + ), + request_id="ltx23-shared-request-inputs", + ) + ] + ) + + resolved_t2v = _resolve_request_inputs_for_test( + _make_ltx23_request_pipe(LTX23Pipeline), + req, + ) + resolved_i2v = _resolve_request_inputs_for_test( + _make_ltx23_request_pipe(LTX23ImageToVideoPipeline), + req, + ) + + assert resolved_i2v.prompt == resolved_t2v.prompt is None + assert resolved_i2v.negative_prompt == resolved_t2v.negative_prompt is None + assert resolved_i2v.height == resolved_t2v.height == 384 + assert resolved_i2v.width == resolved_t2v.width == 512 + assert resolved_i2v.num_frames == resolved_t2v.num_frames == 25 + assert resolved_i2v.frame_rate == resolved_t2v.frame_rate == 12.5 + assert resolved_i2v.num_inference_steps == resolved_t2v.num_inference_steps == 2 + assert resolved_i2v.guidance_scale == resolved_t2v.guidance_scale == 6.0 + assert resolved_i2v.num_videos_per_prompt == resolved_t2v.num_videos_per_prompt == 2 + assert isinstance(resolved_i2v.generator, list) + assert isinstance(resolved_t2v.generator, list) + assert [gen.initial_seed() for gen in resolved_i2v.generator] == [123, 123] + assert [gen.initial_seed() for gen in resolved_t2v.generator] == [123, 123] + assert resolved_i2v.decode_timestep == resolved_t2v.decode_timestep == [0.1] + assert resolved_i2v.decode_noise_scale == resolved_t2v.decode_noise_scale == [0.2] + assert resolved_i2v.output_type == resolved_t2v.output_type == "latent" + assert resolved_i2v.max_sequence_length == resolved_t2v.max_sequence_length == 17 + torch.testing.assert_close(resolved_i2v.latents, video_latents) + torch.testing.assert_close(resolved_t2v.latents, video_latents) + torch.testing.assert_close(resolved_i2v.audio_latents, audio_latents) + torch.testing.assert_close(resolved_t2v.audio_latents, audio_latents) + torch.testing.assert_close(resolved_i2v.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved_t2v.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved_i2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) + torch.testing.assert_close(resolved_t2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) + torch.testing.assert_close(resolved_i2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) + torch.testing.assert_close(resolved_t2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) + torch.testing.assert_close( + resolved_i2v.negative_prompt_attention_mask, + torch.stack([negative_attention_mask]), + ) + torch.testing.assert_close( + resolved_t2v.negative_prompt_attention_mask, + torch.stack([negative_attention_mask]), + ) + + def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = _make_ltx23_request_pipe(LTX23Pipeline) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "with embeds", + "additional_information": {"prompt_embeds": torch.tensor([[1.0]])}, + }, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + num_inference_steps=2, + ), + request_id="ltx23-mixed-precomputed-fields-0", + ), + OmniDiffusionRequest( + prompt={"prompt": "without embeds"}, + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + num_inference_steps=2, + ), + request_id="ltx23-mixed-precomputed-fields-1", + ), + ] + ) + + with pytest.raises(ValueError, match="mix of provided and missing prompt_embeds"): + _resolve_request_inputs_for_test(pipe, req) + + +class TestLTX23ForwardStages: + def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + + pipe = _make_ltx23_request_pipe(LTX23Pipeline) + prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]]) + negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]]) + prompt_attention_mask = torch.tensor([[True, False], [False, True]]) + negative_prompt_attention_mask = torch.tensor([[False, False], [True, True]]) + + ( + repeated_prompt_embeds, + repeated_prompt_attention_mask, + repeated_negative_prompt_embeds, + repeated_negative_prompt_attention_mask, + ) = pipe.encode_prompt( + prompt=None, + do_classifier_free_guidance=True, + num_videos_per_prompt=2, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + prompt_attention_mask=prompt_attention_mask, + negative_prompt_attention_mask=negative_prompt_attention_mask, + device=torch.device("cpu"), + ) + + torch.testing.assert_close( + repeated_prompt_embeds, + torch.tensor([[[1.0], [2.0]], [[1.0], [2.0]], [[3.0], [4.0]], [[3.0], [4.0]]]), + ) + torch.testing.assert_close( + repeated_negative_prompt_embeds, + torch.tensor([[[-1.0], [-2.0]], [[-1.0], [-2.0]], [[-3.0], [-4.0]], [[-3.0], [-4.0]]]), + ) + torch.testing.assert_close( + repeated_prompt_attention_mask, + torch.tensor([[True, False], [True, False], [False, True], [False, True]]), + ) + torch.testing.assert_close( + repeated_negative_prompt_attention_mask, + torch.tensor([[False, False], [False, False], [True, True], [True, True]]), + ) + + def test_t2v_forward_delegates_to_shared_forward_impl(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + pipe = _make_ltx23_request_pipe(LTX23Pipeline) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="stage refactor prompt", + sampling_params=OmniDiffusionSamplingParams( + height=384, + width=512, + num_frames=25, + frame_rate=16, + num_inference_steps=3, + guidance_scale=4.5, + output_type="latent", + ), + request_id="ltx23-t2v-forward-stage-delegation", + ) + ] + ) + sigmas = [1.0, 0.5] + timesteps = [1000, 500] + attention_kwargs = {"scale": 1.0} + seen = {} + + def fake_forward_impl(req_arg, request_inputs, **kwargs): + seen["req"] = req_arg + seen["request_inputs"] = request_inputs + seen["kwargs"] = kwargs + return ["delegated"] + + object.__setattr__(pipe, "_forward_impl", fake_forward_impl) + + output = pipe.forward( + req, + sigmas=sigmas, + timesteps=timesteps, + noise_scale=0.25, + attention_kwargs=attention_kwargs, + ) + + assert output == ["delegated"] + assert seen["req"] is req + assert seen["request_inputs"].height == 384 + assert seen["request_inputs"].width == 512 + assert seen["request_inputs"].num_frames == 25 + assert seen["request_inputs"].frame_rate == 16.0 + assert seen["request_inputs"].guidance_scale == 4.5 + assert seen["request_inputs"].output_type == "latent" + assert seen["kwargs"] == { + "noise_scale": 0.25, + "sigmas": sigmas, + "timesteps": timesteps, + "attention_kwargs": attention_kwargs, + } + + class TestPipelineIndependence: """Verify LTX23Pipeline is fully independent from LTX2Pipeline.""" @@ -65,13 +338,6 @@ class TestPipelineIndependence: assert issubclass(LTX23Pipeline, ProgressBarMixin) - def test_ltx23_pipeline_has_cfg_parallel_mixin(self): - """LTX23Pipeline must use the shared CFG parallel implementation.""" - from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - - assert issubclass(LTX23Pipeline, CFGParallelMixin) - def test_ltx23_pipeline_declares_offload_components(self): """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery.""" from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline @@ -105,661 +371,40 @@ class TestPipelineIndependence: assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin) -class TestLTX23VaeDecodeParallel: - """Test LTX-2.3 video VAE tiled parallel helpers without loading weights.""" +class TestLTX23DecodeConditioning: + def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value - def test_ltx23_video_vae_is_distributed_tile_only_class(self): - from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( - DistributedAutoencoderKLLTX2Video, - ) - from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin + assert _expand_per_prompt_decode_value( + [0.1, 0.2], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", + ) == [0.1, 0.1, 0.2, 0.2] + assert _expand_per_prompt_decode_value( + [0.3], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", + ) == [0.3, 0.3, 0.3, 0.3] + assert _expand_per_prompt_decode_value( + [0.1, 0.2, 0.3, 0.4], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", + ) == [0.1, 0.2, 0.3, 0.4] - assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin) - assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split") + def test_decode_conditioning_rejects_ambiguous_lengths(self): + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import _expand_per_prompt_decode_value - def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self): - from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( - DistributedAutoencoderKLLTX2Video, - ) - - vae = SimpleNamespace( - spatial_compression_ratio=32, - tile_sample_min_height=512, - tile_sample_min_width=512, - tile_sample_stride_height=448, - tile_sample_stride_width=448, - temporal_compression_ratio=8, - dtype=torch.float32, - ) - - z = torch.zeros(1, 2, 5, 16, 24) - tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z) - - assert grid_spec.grid_shape == (2, 2) - assert grid_spec.split_dims == (3, 4) - assert grid_spec.tile_spec["sample_height"] == 512 - assert grid_spec.tile_spec["sample_width"] == 768 - assert grid_spec.tile_spec["blend_height"] == 64 - assert grid_spec.tile_spec["blend_width"] == 64 - assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512) - assert grid_spec.tile_spec["tile_output_shapes"] == { - 0: (1, 3, 33, 512, 512), - 1: (1, 3, 33, 512, 320), - 2: (1, 3, 33, 64, 512), - 3: (1, 3, 33, 64, 320), - } - assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)] - assert [tuple(task.tensor.shape) for task in tasks] == [ - (1, 2, 5, 16, 16), - (1, 2, 5, 16, 10), - (1, 2, 5, 2, 16), - (1, 2, 5, 2, 10), - ] - assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10] - - def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self): - from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( - DistributedAutoencoderKLLTX2Video, - ) - from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec - - class FakeVae: - def __init__(self): - self.blend_calls = [] - - def clear_cache(self): - pass - - def blend_v(self, _previous, current, blend_height): - self.blend_calls.append(("v", blend_height)) - return current - - def blend_h(self, _previous, current, blend_width): - self.blend_calls.append(("h", blend_width)) - return current - - fake_vae = FakeVae() - grid_spec = GridSpec( - split_dims=(3, 4), - grid_shape=(2, 2), - tile_spec={ - "sample_height": 10, - "sample_width": 10, - "blend_height": 1, - "blend_width": 2, - "tile_sample_stride_height": 5, - "tile_sample_stride_width": 5, - }, - ) - tiles = { - (0, 0): torch.full((1, 3, 2, 6, 6), 1.0), - (0, 1): torch.full((1, 3, 2, 6, 6), 2.0), - (1, 0): torch.full((1, 3, 2, 6, 6), 3.0), - (1, 1): torch.full((1, 3, 2, 6, 6), 4.0), - } - - merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec) - - assert merged.shape == (1, 3, 2, 10, 10) - assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)] - torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5)) - torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0)) - torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0)) - torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0)) - - def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self): - from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2 - - z = torch.zeros(1, 2, 1, 16, 24) - expected = torch.ones(1, 3, 1, 512, 768) - seen = {} - - class FakeExecutor: - def execute(self, tensor, operator, broadcast_result=True): - seen["tensor"] = tensor - seen["operator"] = operator - seen["broadcast_result"] = broadcast_result - return expected - - vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True) - vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae) - vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae) - vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae) - - output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode( - vae, - z, - temb=torch.tensor(0.5), - return_dict=False, - ) - - assert len(output) == 1 - assert output[0] is expected - assert seen["tensor"] is z - assert seen["broadcast_result"] is False - assert seen["operator"].split.__name__ == "tile_split" - assert seen["operator"].merge.__name__ == "tile_merge" - - def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self): - from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor - from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( - DistributedOperator, - GridSpec, - TileTask, - ) - - z = torch.zeros(1, 1, 1, 1, 1) - tile_output_shapes = { - 0: (1, 1, 1, 2, 2), - 1: (1, 1, 1, 2, 1), - 2: (1, 1, 1, 1, 2), - 3: (1, 1, 1, 1, 1), - } - tasks = [ - TileTask(0, (0, 0), z, workload=4), - TileTask(1, (0, 1), z, workload=2), - TileTask(2, (1, 0), z, workload=2), - TileTask(3, (1, 1), z, workload=1), - ] - grid_spec = GridSpec( - split_dims=(3, 4), - grid_shape=(2, 2), - tile_spec={ - "max_tile_output_shape": (1, 1, 1, 2, 2), - "tile_output_shapes": tile_output_shapes, - }, - output_dtype=torch.float32, - ) - seen = {} - - def exec_tile(task): - return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1)) - - def merge_tiles(coord_tensor_map, passed_grid_spec): - seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()} - assert passed_grid_spec is grid_spec - return torch.stack( - [ - coord_tensor_map[(0, 0)].flatten()[0], - coord_tensor_map[(0, 1)].flatten()[0], - coord_tensor_map[(1, 0)].flatten()[0], - coord_tensor_map[(1, 1)].flatten()[0], - ] + with pytest.raises(ValueError, match="decode_timestep"): + _expand_per_prompt_decode_value( + [0.1, 0.2, 0.3], + prompt_batch_size=2, + effective_batch_size=4, + field_name="decode_timestep", ) - operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles) - - rank0_executor = object.__new__(LTX2VaeExecutor) - rank0_executor.parallel_size = 2 - rank0_executor.world_size = 2 - rank0_executor.rank = 0 - - def gather_rank0(local_tile_tensor): - assigned = rank0_executor._balance_tasks(tasks, 2) - rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]] - rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta( - rank1_results, - list(local_tile_tensor.shape), - z.device, - torch.float32, - ) - seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape) - return [local_tile_tensor, rank1_tile_tensor] - - def fail_final_sync(*_args, **_kwargs): - raise AssertionError("broadcast_result=False should not sync the final result") - - rank0_executor.gather_tensors = gather_rank0 - rank0_executor._sync_final_result = fail_final_sync - - rank0_result = rank0_executor.execute(z, operator, broadcast_result=False) - - torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0])) - assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2) - assert seen["merged_shapes"] == { - (0, 0): (1, 1, 1, 2, 2), - (0, 1): (1, 1, 1, 2, 1), - (1, 0): (1, 1, 1, 1, 2), - (1, 1): (1, 1, 1, 1, 1), - } - - non_rank0_executor = object.__new__(LTX2VaeExecutor) - non_rank0_executor.parallel_size = 2 - non_rank0_executor.world_size = 2 - non_rank0_executor.rank = 1 - - def gather_rank1(local_tile_tensor): - seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape) - return None - - def fail_non_rank0_merge(*_args, **_kwargs): - raise AssertionError("non-rank0 should not merge gathered tiles") - - non_rank0_executor.gather_tensors = gather_rank1 - non_rank0_executor._sync_final_result = fail_final_sync - - empty_result = non_rank0_executor.execute( - z, - DistributedOperator( - split=lambda _z: (tasks, grid_spec), - exec=exec_tile, - merge=fail_non_rank0_merge, - ), - broadcast_result=False, - ) - - assert tuple(empty_result.shape) == (0,) - assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2) - - -class TestCFGParallelHelpers: - """Test LTX-2.3 CFG helper math without loading model weights.""" - - def test_combine_cfg_noise_matches_x0_space_formula(self): - import torch - - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline - - pipe = object.__new__(LTX23Pipeline) - video_sample = torch.tensor([[[1.0, -2.0]]]) - audio_sample = torch.tensor([[[0.5, 3.0]]]) - video_pos = torch.tensor([[[0.2, -0.3]]]) - video_neg = torch.tensor([[[-0.4, 0.1]]]) - audio_pos = torch.tensor([[[0.7, -0.2]]]) - audio_neg = torch.tensor([[[0.1, 0.4]]]) - video_sigma = torch.tensor(0.25) - audio_sigma = torch.tensor(0.5) - scale = 4.0 - - video_combined, audio_combined = pipe.combine_cfg_noise( - (video_pos, audio_pos), - (video_neg, audio_neg), - scale, - video_latents=video_sample, - audio_latents=audio_sample, - video_sigma=video_sigma, - audio_sigma=audio_sigma, - ) - - x0_video_cond = video_sample - video_pos * video_sigma - x0_video_uncond = video_sample - video_neg * video_sigma - x0_video_guided = x0_video_cond + (scale - 1) * (x0_video_cond - x0_video_uncond) - expected_video = (video_sample - x0_video_guided) / video_sigma - - x0_audio_cond = audio_sample - audio_pos * audio_sigma - x0_audio_uncond = audio_sample - audio_neg * audio_sigma - x0_audio_guided = x0_audio_cond + (scale - 1) * (x0_audio_cond - x0_audio_uncond) - expected_audio = (audio_sample - x0_audio_guided) / audio_sigma - assert torch.allclose(video_combined, expected_video) - assert torch.allclose(audio_combined, expected_audio) - - def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch): - from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 - - pipe = object.__new__(ltx23.LTX23Pipeline) - video_sample = torch.tensor([[[1.0, -2.0]]]) - audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]]) - video_pos = torch.tensor([[[0.2, -0.3]]]) - video_neg = torch.tensor([[[-0.4, 0.1]]]) - audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]]) - audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]]) - video_sigma = torch.tensor(0.25) - audio_sigma = torch.tensor(0.5) - scale = 4.0 - - class FakeCfgGroup: - def all_gather(self, tensor, separate_tensors=True): - assert separate_tensors - if tensor.shape == video_pos.shape: - return [video_pos, video_neg] - return [audio_pos, audio_neg] - - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) - - expected_video = ltx23.LTX23Pipeline._combine_x0_space_cfg( - video_sample, - video_pos, - video_neg, - video_sigma, - scale, - ) - expected_audio = ltx23.LTX23Pipeline._combine_x0_space_cfg( - audio_sample, - audio_pos, - audio_neg, - audio_sigma, - scale, - ) - - for rank, expected_branch in ((0, "positive"), (1, "negative")): - calls = [] - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda rank=rank: rank) - - def fake_predict_noise(**kwargs): - calls.append(kwargs["branch"]) - if kwargs["branch"] == "positive": - return video_pos, audio_pos - return video_neg, audio_neg - - object.__setattr__(pipe, "predict_noise", fake_predict_noise) - video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg( - true_cfg_scale=scale, - positive_kwargs={"branch": "positive"}, - negative_kwargs={"branch": "negative"}, - cfg_normalize=False, - video_latents=video_sample, - audio_latents=audio_sample, - video_sigma=video_sigma, - audio_sigma=audio_sigma, - ) - - assert calls == [expected_branch] - torch.testing.assert_close(video_combined, expected_video) - torch.testing.assert_close(audio_combined, expected_audio) - - assert "_cfg_video_latents" not in pipe.__dict__ - assert "_cfg_audio_latents" not in pipe.__dict__ - - -class TestCFGParallelForwardPath: - """Test the LTX-2.3 CFG-parallel denoising path without loading model weights.""" - - def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch): - from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 - from vllm_omni.diffusion.request import OmniDiffusionRequest - from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - pipe = object.__new__(ltx23.LTX23Pipeline) - torch.nn.Module.__init__(pipe) - pipe.device = torch.device("cpu") - pipe.tokenizer_max_length = 4 - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 1) - - class StopAtEncodePromptError(Exception): - pass - - captured = {} - - def fake_encode_prompt(**kwargs): - captured.update(kwargs) - raise StopAtEncodePromptError - - object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) - - prompt_embeds_a = torch.zeros(2, 3) - prompt_embeds_b = torch.ones(2, 3) - negative_prompt_embeds_a = torch.full((2, 3), 2.0) - negative_prompt_embeds_b = torch.full((2, 3), 3.0) - prompt_attention_mask_a = torch.tensor([True, True]) - prompt_attention_mask_b = torch.tensor([True, False]) - negative_attention_mask_a = torch.tensor([False, True]) - negative_attention_mask_b = torch.tensor([False, False]) - - requests = [ - OmniDiffusionRequest( - prompt={ - "prompt": "prompt-a", - "negative_prompt": "negative-a", - "prompt_embeds": prompt_embeds_a, - "negative_prompt_embeds": negative_prompt_embeds_a, - "prompt_attention_mask": prompt_attention_mask_a, - "negative_prompt_attention_mask": negative_attention_mask_a, - }, - sampling_params=OmniDiffusionSamplingParams( - height=32, - width=32, - num_frames=1, - frame_rate=1.0, - num_inference_steps=2, - ), - request_id="ltx23-prompt-local-a", - ), - OmniDiffusionRequest( - prompt={ - "prompt": "prompt-b", - "negative_prompt": "negative-b", - "prompt_embeds": prompt_embeds_b, - "negative_prompt_embeds": negative_prompt_embeds_b, - "attention_mask": prompt_attention_mask_b, - "negative_attention_mask": negative_attention_mask_b, - }, - sampling_params=OmniDiffusionSamplingParams( - height=32, - width=32, - num_frames=1, - frame_rate=1.0, - num_inference_steps=2, - ), - request_id="ltx23-prompt-local-b", - ), - ] - - with pytest.raises(StopAtEncodePromptError): - pipe.forward(DiffusionRequestBatch(requests=requests)) - - assert captured["prompt"] is None - assert captured["negative_prompt"] is None - torch.testing.assert_close( - captured["prompt_embeds"], - torch.stack([prompt_embeds_a, prompt_embeds_b], dim=0), - ) - torch.testing.assert_close( - captured["negative_prompt_embeds"], - torch.stack([negative_prompt_embeds_a, negative_prompt_embeds_b], dim=0), - ) - torch.testing.assert_close( - captured["prompt_attention_mask"], - torch.stack([prompt_attention_mask_a, prompt_attention_mask_b], dim=0), - ) - torch.testing.assert_close( - captured["negative_prompt_attention_mask"], - torch.stack([negative_attention_mask_a, negative_attention_mask_b], dim=0), - ) - - @pytest.mark.parametrize(("cfg_rank", "expected_prompt_value"), [(0, 1.0), (1, 0.0)]) - @pytest.mark.parametrize( - ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"), - [(1.0, 1, 1.0), (None, 24, 24.0)], - ) - def test_forward_cfg_parallel_steps_video_and_audio_scheduler( - self, - monkeypatch, - cfg_rank, - expected_prompt_value, - frame_rate_input, - audio_sampling_rate, - expected_frame_rate, - ): - from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2_3 as ltx23 - from vllm_omni.diffusion.request import OmniDiffusionRequest - from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - pipe = object.__new__(ltx23.LTX23Pipeline) - torch.nn.Module.__init__(pipe) - pipe.device = torch.device("cpu") - pipe.tokenizer_max_length = 1 - pipe.vae_spatial_compression_ratio = 32 - pipe.vae_temporal_compression_ratio = 1 - pipe.transformer_spatial_patch_size = 1 - pipe.transformer_temporal_patch_size = 1 - pipe.audio_sampling_rate = audio_sampling_rate - pipe.audio_hop_length = 1 - pipe.audio_vae_temporal_compression_ratio = 1 - pipe.audio_vae_mel_compression_ratio = 1 - pipe.od_config = SimpleNamespace(parallel_config=SimpleNamespace(sequence_parallel_size=1)) - pipe.tokenizer = SimpleNamespace(padding_side="left") - pipe.vae = SimpleNamespace( - latents_mean=torch.zeros(2), - latents_std=torch.ones(2), - config=SimpleNamespace(scaling_factor=1.0), - ) - pipe.audio_vae = SimpleNamespace( - latents_mean=torch.zeros(2), - latents_std=torch.ones(2), - config=SimpleNamespace(mel_bins=2, latent_channels=1), - ) - - video_pos = torch.tensor([[[0.2, -0.3]]]) - video_neg = torch.tensor([[[-0.4, 0.1]]]) - audio_pos = torch.tensor([[[0.7, -0.2]]]) - audio_neg = torch.tensor([[[0.1, 0.4]]]) - - class FakeCfgGroup: - def all_gather(self, tensor, separate_tensors=True): - assert separate_tensors - if torch.equal(tensor, video_pos) or torch.equal(tensor, video_neg): - return [video_pos, video_neg] - if torch.equal(tensor, audio_pos) or torch.equal(tensor, audio_neg): - return [audio_pos, audio_neg] - raise AssertionError(f"Unexpected gathered tensor: {tensor}") - - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_world_size", lambda: 2) - monkeypatch.setattr(ltx23, "get_classifier_free_guidance_rank", lambda: cfg_rank) - monkeypatch.setattr(ltx23, "get_cfg_group", lambda: FakeCfgGroup()) - - def fake_retrieve_timesteps(scheduler, num_inference_steps, device, timesteps, sigmas=None, mu=None): - scheduler.sigmas = torch.tensor([0.25, 0.25], device=device) - return torch.tensor([1.0, 0.5], device=device), 2 - - monkeypatch.setattr(ltx23, "retrieve_timesteps", fake_retrieve_timesteps) - - class FakeScheduler: - def __init__(self, name="video", calls=None): - self.name = name - self.calls = [] if calls is None else calls - self.config = { - "max_image_seq_len": 4096, - "base_image_seq_len": 1024, - "base_shift": 0.95, - "max_shift": 2.05, - } - self.sigmas = torch.tensor([0.25, 0.25]) - - def __deepcopy__(self, memo): - return FakeScheduler("audio", self.calls) - - def step(self, noise_pred, t, latents, return_dict=False, generator=None): - self.calls.append((self.name, noise_pred.clone(), t.clone(), latents.clone())) - return (latents - noise_pred,) - - class FakeConnectors: - def to(self, device): - return self - - def __call__(self, prompt_embeds, prompt_attention_mask, padding_side): - assert padding_side == "left" - assert prompt_embeds.shape[0] == 2 - return prompt_embeds, prompt_embeds, prompt_attention_mask - - rope_video_fps: list[float] = [] - - class FakeRope: - def prepare_video_coords(self, batch_size, num_frames, height, width, device, fps): - rope_video_fps.append(fps) - return torch.zeros(batch_size, num_frames * height * width, 3, device=device) - - def prepare_audio_coords(self, batch_size, num_frames, device): - return torch.zeros(batch_size, num_frames, 1, device=device) - - class FakeTransformer: - def __init__(self): - self.config = SimpleNamespace(in_channels=2) - self.rope = FakeRope() - self.audio_rope = FakeRope() - self.calls = [] - - def __call__(self, **kwargs): - self.calls.append(kwargs) - expected_prompt = torch.full((1, 1, 1), expected_prompt_value) - torch.testing.assert_close(kwargs["encoder_hidden_states"], expected_prompt) - torch.testing.assert_close(kwargs["audio_encoder_hidden_states"], expected_prompt) - assert kwargs["hidden_states"].shape == (1, 1, 2) - assert kwargs["audio_hidden_states"].shape == (1, 1, 2) - if cfg_rank == 0: - return video_pos, audio_pos - return video_neg, audio_neg - - class DummyProgress: - def __enter__(self): - return self - - def __exit__(self, exc_type, exc, tb): - return False - - def update(self): - pass - - pipe.scheduler = FakeScheduler() - pipe.connectors = FakeConnectors() - pipe.transformer = FakeTransformer() - object.__setattr__(pipe, "progress_bar", lambda total: DummyProgress()) - - def fake_encode_prompt(**kwargs): - return ( - torch.ones(1, 1, 1), - torch.ones(1, 1, dtype=torch.bool), - torch.zeros(1, 1, 1), - torch.ones(1, 1, dtype=torch.bool), - ) - - object.__setattr__(pipe, "encode_prompt", fake_encode_prompt) - - video_latents = torch.tensor([[[1.0, -2.0]]]) - audio_latents = torch.tensor([[[0.5, 3.0]]]) - req = OmniDiffusionRequest( - prompt={"prompt": "prompt", "negative_prompt": "negative"}, - sampling_params=OmniDiffusionSamplingParams( - height=32, - width=32, - num_frames=1, - frame_rate=frame_rate_input, - num_inference_steps=2, - guidance_scale=4.0, - latents=video_latents, - audio_latents=audio_latents, - output_type="latent", - ), - request_id="ltx23-cfg-parallel-forward-test", - ) - - output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0] - - expected_video_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( - video_latents, - video_pos, - video_neg, - pipe.scheduler.sigmas[0], - 4.0, - ) - expected_audio_noise = ltx23.LTX23Pipeline._combine_x0_space_cfg( - audio_latents, - audio_pos, - audio_neg, - pipe.scheduler.sigmas[0], - 4.0, - ) - scheduler_call_names = [call[0] for call in pipe.scheduler.calls] - assert scheduler_call_names == ["video", "audio", "video", "audio"] - assert len(pipe.transformer.calls) == 2 - torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise) - torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise) - torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise) - torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise) - - video_out, audio_out = output.output - torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1)) - torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2)) - - # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve - # to the model's own 24.0 default, not crash on None; a provided rate is passed through. - assert rope_video_fps - assert all(fps == expected_frame_rate for fps in rope_video_fps) - class TestRegistryIntegration: """Verify all LTX-2.3 pipeline variants are registered.""" @@ -814,15 +459,6 @@ class TestRegistryIntegration: for name in expected: assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS" - def test_ltx2_transformer_has_dit_cache_config(self): - """Ensure LTX2 has a Cache DiT adapter config and that it uses separate CFG.""" - from vllm_omni.diffusion.cache.cache_dit_backend import CacheDiTAdapterConfig - from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel - - adapter_config = getattr(LTX2VideoTransformer3DModel, "_cache_dit_adapter_config") - assert isinstance(adapter_config, CacheDiTAdapterConfig) - assert adapter_config.has_separate_cfg - class TestVocoderSampleRateDetection: """Test _detect_vocoder_output_sample_rate logic.""" @@ -912,31 +548,6 @@ class TestPostProcessFunction: assert "audio_sample_rate" not in result -class TestReExportModule: - """Test that pipeline_ltx2_3_image2video.py correctly re-exports.""" - - def test_i2v_classes_importable(self): - """I2V classes must be importable from the re-export module.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline - - assert LTX23ImageToVideoPipeline is not None - - def test_post_process_func_importable(self): - """get_ltx2_post_process_func must be importable from re-export module.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import get_ltx2_post_process_func - - assert callable(get_ltx2_post_process_func) - - def test_i2v_classes_are_same_as_direct_import(self): - """Re-exported classes must be the same objects as direct imports.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23ImageToVideoPipeline as Direct - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import ( - LTX23ImageToVideoPipeline as ReExported, - ) - - assert Direct is ReExported - - class TestInitExports: """Test that __init__.py exports all LTX-2.3 classes.""" diff --git a/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py b/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py new file mode 100644 index 000000000..9b53bb4f8 --- /dev/null +++ b/tests/diffusion/models/ltx2/test_ltx2_3_vae_tiling.py @@ -0,0 +1,282 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Unit tests for LTX-2.3 VAE tiling and distributed decode behavior.""" + +from types import SimpleNamespace + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class TestLTX23OutputRank: + def test_single_process_rank_is_output_rank(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23 + + monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: False) + + assert ltx23._is_output_rank() is True + assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is True + + def test_non_output_rank_skips_decode_unless_vae_decode_is_distributed(self, monkeypatch): + import vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 as ltx23 + + monkeypatch.setattr(ltx23.torch.distributed, "is_initialized", lambda: True) + monkeypatch.setattr(ltx23.torch.distributed, "get_rank", lambda: 1) + + assert ltx23._is_output_rank() is False + assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: False)) is False + assert ltx23._should_decode_video_on_rank(SimpleNamespace(is_distributed_enabled=lambda: True)) is True + + +class TestLTX23VaeDistributedDecode: + """Test LTX-2.3 distributed VAE helpers without loading weights.""" + + def test_ltx23_video_vae_is_distributed_tile_only_class(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin + + assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin) + assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split") + + def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor + from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import ( + DistributedOperator, + GridSpec, + TileTask, + ) + + z = torch.zeros(1, 1, 1, 1, 1) + tile_output_shapes = { + 0: (1, 1, 1, 2, 2), + 1: (1, 1, 1, 2, 1), + 2: (1, 1, 1, 1, 2), + 3: (1, 1, 1, 1, 1), + } + tasks = [ + TileTask(0, (0, 0), z, workload=4), + TileTask(1, (0, 1), z, workload=2), + TileTask(2, (1, 0), z, workload=2), + TileTask(3, (1, 1), z, workload=1), + ] + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(2, 2), + tile_spec={ + "max_tile_output_shape": (1, 1, 1, 2, 2), + "tile_output_shapes": tile_output_shapes, + }, + output_dtype=torch.float32, + ) + seen = {} + + def exec_tile(task): + return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1)) + + def merge_tiles(coord_tensor_map, passed_grid_spec): + seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()} + assert passed_grid_spec is grid_spec + return torch.stack( + [ + coord_tensor_map[(0, 0)].flatten()[0], + coord_tensor_map[(0, 1)].flatten()[0], + coord_tensor_map[(1, 0)].flatten()[0], + coord_tensor_map[(1, 1)].flatten()[0], + ] + ) + + operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles) + + rank0_executor = object.__new__(LTX2VaeExecutor) + rank0_executor.parallel_size = 2 + rank0_executor.world_size = 2 + rank0_executor.rank = 0 + + def gather_rank0(local_tile_tensor): + assigned = rank0_executor._balance_tasks(tasks, 2) + rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]] + rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta( + rank1_results, + list(local_tile_tensor.shape), + z.device, + torch.float32, + ) + seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape) + return [local_tile_tensor, rank1_tile_tensor] + + def fail_final_sync(*_args, **_kwargs): + raise AssertionError("broadcast_result=False should not sync the final result") + + rank0_executor.gather_tensors = gather_rank0 + rank0_executor._sync_final_result = fail_final_sync + + rank0_result = rank0_executor.execute(z, operator, broadcast_result=False) + + torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0])) + assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2) + assert seen["merged_shapes"] == { + (0, 0): (1, 1, 1, 2, 2), + (0, 1): (1, 1, 1, 2, 1), + (1, 0): (1, 1, 1, 1, 2), + (1, 1): (1, 1, 1, 1, 1), + } + + non_rank0_executor = object.__new__(LTX2VaeExecutor) + non_rank0_executor.parallel_size = 2 + non_rank0_executor.world_size = 2 + non_rank0_executor.rank = 1 + + def gather_rank1(local_tile_tensor): + seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape) + return None + + def fail_non_rank0_merge(*_args, **_kwargs): + raise AssertionError("non-rank0 should not merge gathered tiles") + + non_rank0_executor.gather_tensors = gather_rank1 + non_rank0_executor._sync_final_result = fail_final_sync + + empty_result = non_rank0_executor.execute( + z, + DistributedOperator( + split=lambda _z: (tasks, grid_spec), + exec=exec_tile, + merge=fail_non_rank0_merge, + ), + broadcast_result=False, + ) + + assert tuple(empty_result.shape) == (0,) + assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2) + + +class TestLTX23VaeTiling: + """Test LTX-2.3 video VAE tile helpers without loading weights.""" + + def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + + vae = SimpleNamespace( + spatial_compression_ratio=32, + tile_sample_min_height=512, + tile_sample_min_width=512, + tile_sample_stride_height=448, + tile_sample_stride_width=448, + temporal_compression_ratio=8, + dtype=torch.float32, + ) + + z = torch.zeros(1, 2, 5, 16, 24) + tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z) + + assert grid_spec.grid_shape == (2, 2) + assert grid_spec.split_dims == (3, 4) + assert grid_spec.tile_spec["sample_height"] == 512 + assert grid_spec.tile_spec["sample_width"] == 768 + assert grid_spec.tile_spec["blend_height"] == 64 + assert grid_spec.tile_spec["blend_width"] == 64 + assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512) + assert grid_spec.tile_spec["tile_output_shapes"] == { + 0: (1, 3, 33, 512, 512), + 1: (1, 3, 33, 512, 320), + 2: (1, 3, 33, 64, 512), + 3: (1, 3, 33, 64, 320), + } + assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)] + assert [tuple(task.tensor.shape) for task in tasks] == [ + (1, 2, 5, 16, 16), + (1, 2, 5, 16, 10), + (1, 2, 5, 2, 16), + (1, 2, 5, 2, 10), + ] + assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10] + + def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self): + from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( + DistributedAutoencoderKLLTX2Video, + ) + from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec + + class FakeVae: + def __init__(self): + self.blend_calls = [] + + def clear_cache(self): + pass + + def blend_v(self, _previous, current, blend_height): + self.blend_calls.append(("v", blend_height)) + return current + + def blend_h(self, _previous, current, blend_width): + self.blend_calls.append(("h", blend_width)) + return current + + fake_vae = FakeVae() + grid_spec = GridSpec( + split_dims=(3, 4), + grid_shape=(2, 2), + tile_spec={ + "sample_height": 10, + "sample_width": 10, + "blend_height": 1, + "blend_width": 2, + "tile_sample_stride_height": 5, + "tile_sample_stride_width": 5, + }, + ) + tiles = { + (0, 0): torch.full((1, 3, 2, 6, 6), 1.0), + (0, 1): torch.full((1, 3, 2, 6, 6), 2.0), + (1, 0): torch.full((1, 3, 2, 6, 6), 3.0), + (1, 1): torch.full((1, 3, 2, 6, 6), 4.0), + } + + merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec) + + assert merged.shape == (1, 3, 2, 10, 10) + assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)] + torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5)) + torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0)) + torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0)) + torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0)) + + def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self): + from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2 + + z = torch.zeros(1, 2, 1, 16, 24) + expected = torch.ones(1, 3, 1, 512, 768) + seen = {} + + class FakeExecutor: + def execute(self, tensor, operator, broadcast_result=True): + seen["tensor"] = tensor + seen["operator"] = operator + seen["broadcast_result"] = broadcast_result + return expected + + vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True) + vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae) + vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae) + vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae) + + output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode( + vae, + z, + temb=torch.tensor(0.5), + return_dict=False, + ) + + assert len(output) == 1 + assert output[0] is expected + assert seen["tensor"] is z + assert seen["broadcast_result"] is False + assert seen["operator"].split.__name__ == "tile_split" + assert seen["operator"].merge.__name__ == "tile_merge" diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index bf57d3f9e..76a26b5b0 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -10,10 +10,8 @@ from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( get_ltx2_post_process_func, load_transformer_config, ) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import ( - LTX23ImageToVideoPipeline, - LTX23Pipeline, -) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3_image2video import LTX23ImageToVideoPipeline from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_image2video import ( LTX2I2VDMD2Pipeline, LTX2ImageToVideoPipeline, diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py index 9183065dc..45be23296 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3.py @@ -19,6 +19,7 @@ import json import os from collections.abc import Iterable from contextlib import nullcontext +from dataclasses import dataclass from typing import Any, ClassVar import numpy as np @@ -68,6 +69,26 @@ def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None: return sampling.extra_args.get("audio_latents") +def _is_output_rank() -> bool: + return not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0 + + +def _vae_decode_needs_all_ranks(vae: Any) -> bool: + if not torch.distributed.is_initialized(): + return False + is_distributed_enabled = getattr(vae, "is_distributed_enabled", None) + if not callable(is_distributed_enabled): + return False + try: + return bool(is_distributed_enabled()) + except Exception: + return False + + +def _should_decode_video_on_rank(vae: Any) -> bool: + return _is_output_rank() or _vae_decode_needs_all_ranks(vae) + + # Try to import LTX2VocoderWithBWE (diffusers >= 0.38.0) try: from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE @@ -75,6 +96,87 @@ except ImportError: LTX2VocoderWithBWE = None +@dataclass +class _LTX23RequestInputs: + prompt: str | list[str] | None + negative_prompt: str | list[str] | None + height: int + width: int + num_frames: int + frame_rate: float + num_inference_steps: int + guidance_scale: float + num_videos_per_prompt: int + generator: torch.Generator | list[torch.Generator] | None + latents: torch.Tensor | None + audio_latents: torch.Tensor | None + prompt_embeds: torch.Tensor | None + negative_prompt_embeds: torch.Tensor | None + prompt_attention_mask: torch.Tensor | None + negative_prompt_attention_mask: torch.Tensor | None + decode_timestep: float | list[float] + decode_noise_scale: float | list[float] | None + output_type: str + max_sequence_length: int + + +@dataclass +class _LTX23PromptContext: + batch_size: int + connector_prompt_embeds: torch.Tensor + connector_audio_prompt_embeds: torch.Tensor + connector_attention_mask: torch.Tensor + positive_connector_prompt_embeds: torch.Tensor + positive_connector_audio_prompt_embeds: torch.Tensor + positive_connector_attention_mask: torch.Tensor + negative_connector_prompt_embeds: torch.Tensor | None + negative_connector_audio_prompt_embeds: torch.Tensor | None + negative_connector_attention_mask: torch.Tensor | None + + +@dataclass +class _LTX23ForwardContext: + req: DiffusionRequestBatch + request_inputs: _LTX23RequestInputs + prompt_context: _LTX23PromptContext + device: torch.device + cfg_parallel_ready: bool + attention_kwargs: dict[str, Any] | None + latent_num_frames: int + latent_height: int + latent_width: int + latent_mel_bins: int + original_audio_num_frames: int + padded_audio_num_frames: int + timesteps: torch.Tensor + audio_scheduler: Any + video_audio_scheduler: Any + + @property + def batch_size(self) -> int: + return self.prompt_context.batch_size + + @property + def num_videos_per_prompt(self) -> int: + return self.request_inputs.num_videos_per_prompt + + +@dataclass +class _LTX23DenoiseContext: + latents: torch.Tensor + audio_latents: torch.Tensor + video_coords: torch.Tensor + audio_coords: torch.Tensor + conditioning_mask: torch.Tensor | None = None + conditioning_mask_for_model: torch.Tensor | None = None + + +def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_videos_per_prompt: int) -> torch.Tensor: + if num_videos_per_prompt == 1: + return tensor + return tensor.repeat_interleave(num_videos_per_prompt, dim=0) + + def _detect_vocoder_output_sample_rate(model: str) -> int | None: """Detect the vocoder output sample rate from vocoder/config.json. @@ -121,6 +223,58 @@ def get_ltx2_post_process_func(od_config: OmniDiffusionConfig): return post_process_func +def _expand_per_prompt_decode_value( + value: float | list[float], + *, + prompt_batch_size: int, + effective_batch_size: int, + field_name: str, +) -> list[float]: + if not isinstance(value, list): + return [value] * effective_batch_size + if len(value) == 1: + return value * effective_batch_size + if len(value) == effective_batch_size: + return value + if prompt_batch_size > 0 and len(value) == prompt_batch_size and effective_batch_size % prompt_batch_size == 0: + repeats = effective_batch_size // prompt_batch_size + return [item for item in value for _ in range(repeats)] + raise ValueError( + f"`{field_name}` must have length 1, prompt batch size ({prompt_batch_size}), or effective batch size" + f" ({effective_batch_size}); got {len(value)}." + ) + + +def _prepare_decode_timestep_conditioning( + *, + decode_timestep: float | list[float], + decode_noise_scale: float | list[float] | None, + prompt_batch_size: int, + effective_batch_size: int, + device: torch.device, + dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + decode_timestep_values = _expand_per_prompt_decode_value( + decode_timestep, + prompt_batch_size=prompt_batch_size, + effective_batch_size=effective_batch_size, + field_name="decode_timestep", + ) + if decode_noise_scale is None: + decode_noise_scale_values = decode_timestep_values + else: + decode_noise_scale_values = _expand_per_prompt_decode_value( + decode_noise_scale, + prompt_batch_size=prompt_batch_size, + effective_batch_size=effective_batch_size, + field_name="decode_noise_scale", + ) + return ( + torch.tensor(decode_timestep_values, device=device, dtype=dtype), + torch.tensor(decode_noise_scale_values, device=device, dtype=dtype)[:, None, None, None, None], + ) + + class LTX23Pipeline( nn.Module, CFGParallelMixin, @@ -353,12 +507,9 @@ class LTX23Pipeline( # [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) + prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) + prompt_attention_mask = _repeat_prompt_tensor_for_outputs(prompt_attention_mask, num_videos_per_prompt) return prompt_embeds, prompt_attention_mask @@ -384,6 +535,8 @@ class LTX23Pipeline( else: batch_size = prompt_embeds.shape[0] + negative_prompt_embeds_provided = negative_prompt_embeds is not None + if prompt_embeds is None: prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds( prompt=prompt, @@ -392,6 +545,9 @@ class LTX23Pipeline( device=device, dtype=dtype, ) + elif num_videos_per_prompt > 1: + prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt) + prompt_attention_mask = _repeat_prompt_tensor_for_outputs(prompt_attention_mask, num_videos_per_prompt) if do_classifier_free_guidance and negative_prompt_embeds is None: negative_prompt = negative_prompt or "" @@ -416,6 +572,12 @@ class LTX23Pipeline( device=device, dtype=dtype, ) + elif do_classifier_free_guidance and negative_prompt_embeds_provided and num_videos_per_prompt > 1: + negative_prompt_embeds = _repeat_prompt_tensor_for_outputs(negative_prompt_embeds, num_videos_per_prompt) + negative_prompt_attention_mask = _repeat_prompt_tensor_for_outputs( + negative_prompt_attention_mask, + num_videos_per_prompt, + ) return prompt_embeds, prompt_attention_mask, negative_prompt_embeds, negative_prompt_attention_mask @@ -542,6 +704,67 @@ class LTX23Pipeline( return audio_latent_length return provided_latent_length + def _decode_output( + self, + *, + latents: torch.Tensor, + audio_latents: torch.Tensor, + output_type: str, + connector_prompt_embeds: torch.Tensor, + generator: torch.Generator | list[torch.Generator] | None, + device: torch.device, + decode_timestep: float | list[float], + decode_noise_scale: float | list[float] | None, + prompt_batch_size: int, + ) -> DiffusionOutput: + if output_type == "latent": + return DiffusionOutput( + output=(latents, audio_latents), + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + 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) + timestep_decode, decode_noise_scale_t = _prepare_decode_timestep_conditioning( + decode_timestep=decode_timestep, + decode_noise_scale=decode_noise_scale, + prompt_batch_size=prompt_batch_size, + effective_batch_size=latents.shape[0], + device=device, + dtype=latents.dtype, + ) + latents = (1 - decode_noise_scale_t) * latents + decode_noise_scale_t * noise + + if _should_decode_video_on_rank(self.vae): + latents = latents.to(self.vae.dtype) + video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] + else: + video = torch.empty(0, device=latents.device, dtype=latents.dtype) + + if not _is_output_rank(): + return DiffusionOutput( + output=( + torch.empty(0, device=video.device, dtype=video.dtype), + torch.empty(0, device=audio_latents.device, dtype=audio_latents.dtype), + ), + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + + if video.numel() > 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) + + return DiffusionOutput( + output=(video, audio), + stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, + ) + # ------------------------------------------------------------------ # Latent preparation # ------------------------------------------------------------------ @@ -830,6 +1053,713 @@ class LTX23Pipeline( torch.cuda.current_stream(device).synchronize() return latents + def _resolve_request_inputs( + self, + req: DiffusionRequestBatch, + *, + prompt: str | list[str] | None, + negative_prompt: str | list[str] | None, + height: int | None, + width: int | None, + num_frames: int | None, + frame_rate: float | None, + num_inference_steps: int | None, + timesteps: list[int] | None, + guidance_scale: float, + num_videos_per_prompt: int | None, + generator: torch.Generator | list[torch.Generator] | None, + latents: torch.Tensor | None, + audio_latents: torch.Tensor | None, + prompt_embeds: torch.Tensor | None, + negative_prompt_embeds: torch.Tensor | None, + prompt_attention_mask: torch.Tensor | None, + negative_prompt_attention_mask: torch.Tensor | None, + decode_timestep: float | list[float], + decode_noise_scale: float | list[float] | None, + output_type: str, + max_sequence_length: int | None, + ) -> _LTX23RequestInputs: + sampling_params_list = req.sampling_params_list + common_sampling_params = sampling_params_list[0] + 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 = common_sampling_params.height or height or 512 + width = common_sampling_params.width or width or 768 + num_frames = common_sampling_params.num_frames or num_frames or 121 + frame_rate = common_sampling_params.resolved_frame_rate or frame_rate or 24.0 + num_inference_steps = common_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 = ( + common_sampling_params.num_outputs_per_prompt + if common_sampling_params.num_outputs_per_prompt > 0 + else num_videos_per_prompt or 1 + ) + max_sequence_length = ( + common_sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length + ) + + if common_sampling_params.guidance_scale_provided: + guidance_scale = common_sampling_params.guidance_scale + + if generator is None: + generator = req.collate_request_generators(num_videos_per_prompt, generator) + + latents = req.collate_request_tensors("latents", latents) + audio_latents = DiffusionRequestBatch.collate_tensors( + [_get_audio_latents_from_sampling(sampling) for sampling in sampling_params_list], + "audio_latents", + audio_latents, + ) + + prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( + req.prompts, + { + "prompt_embeds": prompt_embeds, + "negative_prompt_embeds": negative_prompt_embeds, + "prompt_attention_mask": prompt_attention_mask, + "negative_prompt_attention_mask": negative_prompt_attention_mask, + }, + field_aliases={ + "prompt_attention_mask": ("prompt_attention_mask", "attention_mask"), + "negative_prompt_attention_mask": ( + "negative_prompt_attention_mask", + "negative_attention_mask", + ), + }, + ) + prompt_embeds = prompt_fields["prompt_embeds"] + negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] + prompt_attention_mask = prompt_fields["prompt_attention_mask"] + negative_prompt_attention_mask = prompt_fields["negative_prompt_attention_mask"] + if prompt_embeds is not None: + prompt = None + if negative_prompt_embeds is not None: + negative_prompt = None + + if common_sampling_params.decode_timestep is not None: + decode_timestep = common_sampling_params.decode_timestep + if common_sampling_params.decode_noise_scale is not None: + decode_noise_scale = common_sampling_params.decode_noise_scale + if common_sampling_params.output_type is not None: + output_type = common_sampling_params.output_type + + return _LTX23RequestInputs( + prompt=prompt, + negative_prompt=negative_prompt, + height=int(height), + width=int(width), + num_frames=int(num_frames), + frame_rate=float(frame_rate), + num_inference_steps=int(num_inference_steps), + guidance_scale=guidance_scale, + num_videos_per_prompt=int(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=int(max_sequence_length), + ) + + def _prepare_prompt_context( + self, + *, + prompt: str | list[str] | None, + negative_prompt: str | list[str] | None, + prompt_embeds: torch.Tensor | None, + negative_prompt_embeds: torch.Tensor | None, + prompt_attention_mask: torch.Tensor | None, + negative_prompt_attention_mask: torch.Tensor | None, + num_videos_per_prompt: int, + max_sequence_length: int, + ) -> _LTX23PromptContext: + 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] + + 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=self.device, + ) + ) + + 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) + + connector_prompt_embeds, connector_audio_prompt_embeds, connector_attention_mask = self.connectors( + prompt_embeds, + prompt_attention_mask, + padding_side=getattr(self.tokenizer, "padding_side", "left"), + ) + + positive_connector_prompt_embeds = connector_prompt_embeds + positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds + positive_connector_attention_mask = connector_attention_mask + negative_connector_prompt_embeds = None + negative_connector_audio_prompt_embeds = None + negative_connector_attention_mask = None + if self.do_classifier_free_guidance: + split_batch = batch_size * num_videos_per_prompt + negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch] + positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:] + negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch] + positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:] + negative_connector_attention_mask = connector_attention_mask[:split_batch] + positive_connector_attention_mask = connector_attention_mask[split_batch:] + + return _LTX23PromptContext( + batch_size=batch_size, + connector_prompt_embeds=connector_prompt_embeds, + connector_audio_prompt_embeds=connector_audio_prompt_embeds, + connector_attention_mask=connector_attention_mask, + positive_connector_prompt_embeds=positive_connector_prompt_embeds, + positive_connector_audio_prompt_embeds=positive_connector_audio_prompt_embeds, + positive_connector_attention_mask=positive_connector_attention_mask, + negative_connector_prompt_embeds=negative_connector_prompt_embeds, + negative_connector_audio_prompt_embeds=negative_connector_audio_prompt_embeds, + negative_connector_attention_mask=negative_connector_attention_mask, + ) + + def _setup_forward_runtime( + self, + request_inputs: _LTX23RequestInputs, + attention_kwargs: dict[str, Any] | None, + ) -> bool: + self._guidance_scale = request_inputs.guidance_scale + self._attention_kwargs = attention_kwargs + self._interrupt = False + self._current_timestep = None + cfg_world_size = get_classifier_free_guidance_world_size() + if self.do_classifier_free_guidance and cfg_world_size not in (1, 2): + raise ValueError( + f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}." + ) + return self.do_classifier_free_guidance and cfg_world_size > 1 + + def _check_forward_inputs( + self, + request_inputs: _LTX23RequestInputs, + image: Any | None = None, + ) -> None: + self.check_inputs( + prompt=request_inputs.prompt, + height=request_inputs.height, + width=request_inputs.width, + 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 and request_inputs.latents.ndim == 5: + _, _, latent_num_frames, latent_height, latent_width = request_inputs.latents.shape + 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]: + num_channels_latents = self.transformer.config.in_channels + latents = self.prepare_latents( + 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, + ) + return latents, None + + def _prepare_audio_latents_stage( + self, + request_inputs: _LTX23RequestInputs, + prompt_context: _LTX23PromptContext, + *, + device: torch.device, + noise_scale: float, + ) -> tuple[torch.Tensor, int, int, int]: + duration_s = request_inputs.num_frames / request_inputs.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) + audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, request_inputs.audio_latents) + + 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( + prompt_context.batch_size * request_inputs.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=request_inputs.generator, + latents=request_inputs.audio_latents, + ) + return audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins + + def _make_video_audio_scheduler( + self, + audio_scheduler: Any, + latent_num_frames: int, + latent_height: int, + latent_width: int, + ) -> Any: + return _VideoAudioScheduler(self.scheduler, audio_scheduler) + + def _prepare_scheduler_stage( + self, + request_inputs: _LTX23RequestInputs, + *, + device: torch.device, + sigmas: list[float] | None, + timesteps: list[int] | None, + latent_num_frames: int, + latent_height: int, + latent_width: int, + ) -> tuple[Any, Any, torch.Tensor]: + sigmas = ( + np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.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) + video_audio_scheduler = self._make_video_audio_scheduler( + audio_scheduler, + latent_num_frames, + latent_height, + latent_width, + ) + _ = retrieve_timesteps( + audio_scheduler, + request_inputs.num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + timesteps, _ = retrieve_timesteps( + self.scheduler, + request_inputs.num_inference_steps, + device, + timesteps, + sigmas=sigmas, + mu=mu, + ) + self._num_timesteps = len(timesteps) + return audio_scheduler, video_audio_scheduler, timesteps + + def _prepare_rope_coords_stage( + self, + forward_ctx: _LTX23ForwardContext, + latents: torch.Tensor, + audio_latents: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + request_inputs = forward_ctx.request_inputs + video_coords = self.transformer.rope.prepare_video_coords( + latents.shape[0], + forward_ctx.latent_num_frames, + forward_ctx.latent_height, + forward_ctx.latent_width, + latents.device, + fps=request_inputs.frame_rate, + ) + audio_coords = self.transformer.audio_rope.prepare_audio_coords( + audio_latents.shape[0], + forward_ctx.padded_audio_num_frames, + audio_latents.device, + ) + return video_coords, audio_coords + + def _prepare_denoise_context_for_cfg( + self, + forward_ctx: _LTX23ForwardContext, + denoise_ctx: _LTX23DenoiseContext, + ) -> _LTX23DenoiseContext: + if self.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready: + denoise_ctx.video_coords = denoise_ctx.video_coords.repeat( + (2,) + (1,) * (denoise_ctx.video_coords.ndim - 1) + ) + denoise_ctx.audio_coords = denoise_ctx.audio_coords.repeat( + (2,) + (1,) * (denoise_ctx.audio_coords.ndim - 1) + ) + return denoise_ctx + + def _denoise_timestep_kwargs( + self, + ts: torch.Tensor, + forward_ctx: _LTX23ForwardContext, + denoise_ctx: _LTX23DenoiseContext, + ) -> dict[str, torch.Tensor]: + return {"timestep": ts, "sigma": ts} + + def _build_transformer_kwargs( + self, + forward_ctx: _LTX23ForwardContext, + denoise_ctx: _LTX23DenoiseContext, + *, + hidden_states: torch.Tensor, + audio_hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + audio_encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor, + audio_encoder_attention_mask: torch.Tensor, + ts: torch.Tensor, + ) -> dict[str, Any]: + request_inputs = forward_ctx.request_inputs + return { + "hidden_states": hidden_states, + "audio_hidden_states": audio_hidden_states, + "encoder_hidden_states": encoder_hidden_states, + "audio_encoder_hidden_states": audio_encoder_hidden_states, + **self._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx), + "encoder_attention_mask": encoder_attention_mask, + "audio_encoder_attention_mask": audio_encoder_attention_mask, + "num_frames": forward_ctx.latent_num_frames, + "height": forward_ctx.latent_height, + "width": forward_ctx.latent_width, + "fps": request_inputs.frame_rate, + "audio_num_frames": forward_ctx.padded_audio_num_frames, + "video_coords": denoise_ctx.video_coords, + "audio_coords": denoise_ctx.audio_coords, + "attention_kwargs": forward_ctx.attention_kwargs, + "return_dict": False, + } + + 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]: + if forward_ctx.cfg_parallel_ready: + 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, + ) + + latents = self.scheduler.step(noise_pred_video, t, denoise_ctx.latents, return_dict=False)[0] + audio_latents = forward_ctx.audio_scheduler.step( + noise_pred_audio, + t, + denoise_ctx.audio_latents, + return_dict=False, + )[0] + return latents, audio_latents + + def _denoise_loop( + self, + forward_ctx: _LTX23ForwardContext, + denoise_ctx: _LTX23DenoiseContext, + ) -> tuple[torch.Tensor, torch.Tensor]: + request_inputs = forward_ctx.request_inputs + prompt_context = forward_ctx.prompt_context + guidance_scale = request_inputs.guidance_scale + audio_scheduler = forward_ctx.audio_scheduler + + with self.progress_bar(total=len(forward_ctx.timesteps)) as pbar: + for i, t in enumerate(forward_ctx.timesteps): + if self.interrupt: + continue + + self._current_timestep = t + + if forward_ctx.cfg_parallel_ready: + latent_model_input = denoise_ctx.latents.to(prompt_context.positive_connector_prompt_embeds.dtype) + audio_latent_model_input = denoise_ctx.audio_latents.to( + prompt_context.positive_connector_prompt_embeds.dtype + ) + ts = t.expand(latent_model_input.shape[0]) + positive_kwargs = self._build_transformer_kwargs( + forward_ctx, + denoise_ctx, + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=prompt_context.positive_connector_prompt_embeds, + audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds, + encoder_attention_mask=prompt_context.positive_connector_attention_mask, + audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask, + ts=ts, + ) + negative_kwargs = { + **positive_kwargs, + "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds, + "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds, + "encoder_attention_mask": prompt_context.negative_connector_attention_mask, + "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask, + } + noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg( + true_cfg_scale=guidance_scale, + positive_kwargs=positive_kwargs, + negative_kwargs=negative_kwargs, + cfg_normalize=False, + video_latents=denoise_ctx.latents, + audio_latents=denoise_ctx.audio_latents, + video_sigma=self.scheduler.sigmas[i], + audio_sigma=audio_scheduler.sigmas[i], + ) + else: + latent_model_input = ( + torch.cat([denoise_ctx.latents] * 2) + if self.do_classifier_free_guidance + else denoise_ctx.latents + ) + latent_model_input = latent_model_input.to(prompt_context.connector_prompt_embeds.dtype) + audio_latent_model_input = ( + torch.cat([denoise_ctx.audio_latents] * 2) + if self.do_classifier_free_guidance + else denoise_ctx.audio_latents + ) + audio_latent_model_input = audio_latent_model_input.to(prompt_context.connector_prompt_embeds.dtype) + ts = t.expand(latent_model_input.shape[0]) + + transformer_kwargs = self._build_transformer_kwargs( + forward_ctx, + denoise_ctx, + hidden_states=latent_model_input, + audio_hidden_states=audio_latent_model_input, + encoder_hidden_states=prompt_context.connector_prompt_embeds, + audio_encoder_hidden_states=prompt_context.connector_audio_prompt_embeds, + encoder_attention_mask=prompt_context.connector_attention_mask, + audio_encoder_attention_mask=prompt_context.connector_attention_mask, + ts=ts, + ) + with self._transformer_cache_context("cond_uncond"): + noise_pred_video, noise_pred_audio = self.transformer(**transformer_kwargs) + + noise_pred_video = noise_pred_video.float() + noise_pred_audio = noise_pred_audio.float() + + if self.do_classifier_free_guidance: + noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) + noise_pred_video = self._combine_x0_space_cfg( + denoise_ctx.latents, + noise_pred_video_cond, + noise_pred_video_uncond, + self.scheduler.sigmas[i], + guidance_scale, + ) + + noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) + noise_pred_audio = self._combine_x0_space_cfg( + denoise_ctx.audio_latents, + noise_pred_audio_cond, + noise_pred_audio_uncond, + audio_scheduler.sigmas[i], + guidance_scale, + ) + + denoise_ctx.latents, denoise_ctx.audio_latents = self._step_denoised_latents( + forward_ctx, + denoise_ctx, + noise_pred_video, + noise_pred_audio, + t, + ) + pbar.update() + + return denoise_ctx.latents, denoise_ctx.audio_latents + + def _unpack_and_denormalize_stage( + self, + forward_ctx: _LTX23ForwardContext, + latents: torch.Tensor, + audio_latents: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + latents = self._unpack_latents( + latents, + forward_ctx.latent_num_frames, + forward_ctx.latent_height, + forward_ctx.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, forward_ctx.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, + forward_ctx.original_audio_num_frames, + num_mel_bins=forward_ctx.latent_mel_bins, + ) + return latents, audio_latents + + def _decode_and_split( + self, + forward_ctx: _LTX23ForwardContext, + latents: torch.Tensor, + audio_latents: torch.Tensor, + ) -> list[DiffusionOutput]: + request_inputs = forward_ctx.request_inputs + return split_diffusion_output_by_request( + self._decode_output( + latents=latents, + audio_latents=audio_latents, + output_type=request_inputs.output_type, + connector_prompt_embeds=forward_ctx.prompt_context.connector_prompt_embeds, + generator=request_inputs.generator, + device=forward_ctx.device, + decode_timestep=request_inputs.decode_timestep, + decode_noise_scale=request_inputs.decode_noise_scale, + prompt_batch_size=forward_ctx.batch_size, + ), + forward_ctx.req, + num_outputs_per_prompt=forward_ctx.num_videos_per_prompt, + ) + + def _forward_impl( + self, + req: DiffusionRequestBatch, + request_inputs: _LTX23RequestInputs, + *, + noise_scale: float, + sigmas: list[float] | None, + timesteps: list[int] | None, + attention_kwargs: dict[str, Any] | None, + image: Any | None = None, + ) -> list[DiffusionOutput]: + self._check_forward_inputs(request_inputs, image=image) + cfg_parallel_ready = self._setup_forward_runtime(request_inputs, attention_kwargs) + device = self.device + prompt_context = self._prepare_prompt_context( + prompt=request_inputs.prompt, + negative_prompt=request_inputs.negative_prompt, + 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, + num_videos_per_prompt=request_inputs.num_videos_per_prompt, + max_sequence_length=request_inputs.max_sequence_length, + ) + + latent_num_frames, latent_height, latent_width = self._resolve_video_latent_dimensions(request_inputs) + latents, conditioning_mask = self._prepare_video_latents_stage( + request_inputs, + prompt_context, + device=device, + noise_scale=noise_scale, + image=image, + ) + audio_latents, original_audio_num_frames, padded_audio_num_frames, latent_mel_bins = ( + self._prepare_audio_latents_stage( + request_inputs, + prompt_context, + device=device, + noise_scale=noise_scale, + ) + ) + audio_scheduler, video_audio_scheduler, timesteps_tensor = self._prepare_scheduler_stage( + request_inputs, + device=device, + sigmas=sigmas, + timesteps=timesteps, + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + ) + forward_ctx = _LTX23ForwardContext( + req=req, + request_inputs=request_inputs, + prompt_context=prompt_context, + device=device, + cfg_parallel_ready=cfg_parallel_ready, + attention_kwargs=attention_kwargs, + latent_num_frames=latent_num_frames, + latent_height=latent_height, + latent_width=latent_width, + latent_mel_bins=latent_mel_bins, + original_audio_num_frames=original_audio_num_frames, + padded_audio_num_frames=padded_audio_num_frames, + timesteps=timesteps_tensor, + audio_scheduler=audio_scheduler, + video_audio_scheduler=video_audio_scheduler, + ) + video_coords, audio_coords = self._prepare_rope_coords_stage(forward_ctx, latents, audio_latents) + denoise_ctx = _LTX23DenoiseContext( + latents=latents, + audio_latents=audio_latents, + video_coords=video_coords, + audio_coords=audio_coords, + conditioning_mask=conditioning_mask, + ) + denoise_ctx = self._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx) + latents, audio_latents = self._denoise_loop(forward_ctx, denoise_ctx) + latents, audio_latents = self._unpack_and_denormalize_stage(forward_ctx, latents, audio_latents) + return self._decode_and_split(forward_ctx, latents, audio_latents) + # ------------------------------------------------------------------ # Forward pass # ------------------------------------------------------------------ @@ -864,446 +1794,39 @@ class LTX23Pipeline( attention_kwargs: dict[str, Any] | None = None, max_sequence_length: int | None = None, ) -> list[DiffusionOutput]: - # ---- Extract from request ---- - sampling_params_list = req.sampling_params_list - common_sampling_params = sampling_params_list[0] - 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 = common_sampling_params.height or height or 512 - width = common_sampling_params.width or width or 768 - num_frames = common_sampling_params.num_frames or num_frames or 121 - frame_rate = common_sampling_params.resolved_frame_rate or frame_rate or 24.0 - num_inference_steps = common_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 = ( - common_sampling_params.num_outputs_per_prompt - if common_sampling_params.num_outputs_per_prompt > 0 - else num_videos_per_prompt or 1 - ) - max_sequence_length = ( - common_sampling_params.max_sequence_length or max_sequence_length or self.tokenizer_max_length - ) - - if common_sampling_params.guidance_scale_provided: - guidance_scale = common_sampling_params.guidance_scale - - if generator is None: - generator = req.collate_request_generators(num_videos_per_prompt, generator) - - latents = req.collate_request_tensors("latents", latents) - audio_latents = DiffusionRequestBatch.collate_tensors( - [_get_audio_latents_from_sampling(sampling) for sampling in sampling_params_list], - "audio_latents", - audio_latents, - ) - - # Override with pre-computed embeddings if provided in request - prompt_fields = DiffusionRequestBatch.collate_prompt_field_map( - req.prompts, - { - "prompt_embeds": prompt_embeds, - "negative_prompt_embeds": negative_prompt_embeds, - "prompt_attention_mask": prompt_attention_mask, - "negative_prompt_attention_mask": negative_prompt_attention_mask, - }, - field_aliases={ - "prompt_attention_mask": ("prompt_attention_mask", "attention_mask"), - "negative_prompt_attention_mask": ( - "negative_prompt_attention_mask", - "negative_attention_mask", - ), - }, - ) - prompt_embeds = prompt_fields["prompt_embeds"] - negative_prompt_embeds = prompt_fields["negative_prompt_embeds"] - prompt_attention_mask = prompt_fields["prompt_attention_mask"] - negative_prompt_attention_mask = prompt_fields["negative_prompt_attention_mask"] - if prompt_embeds is not None: - prompt = None - if negative_prompt_embeds is not None: - negative_prompt = None - - if common_sampling_params.decode_timestep is not None: - decode_timestep = common_sampling_params.decode_timestep - if common_sampling_params.decode_noise_scale is not None: - decode_noise_scale = common_sampling_params.decode_noise_scale - if common_sampling_params.output_type is not None: - output_type = common_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 - cfg_world_size = get_classifier_free_guidance_world_size() - if self.do_classifier_free_guidance and cfg_world_size not in (1, 2): - raise ValueError( - f"LTX23Pipeline supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}." - ) - cfg_parallel_ready = self.do_classifier_free_guidance and cfg_world_size > 1 - - 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( + request_inputs = self._resolve_request_inputs( + req, prompt=prompt, negative_prompt=negative_prompt, - do_classifier_free_guidance=self.do_classifier_free_guidance, + 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, - 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) - - 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 - ) - - positive_connector_prompt_embeds = connector_prompt_embeds - positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds - positive_connector_attention_mask = connector_attention_mask - negative_connector_prompt_embeds = None - negative_connector_audio_prompt_embeds = None - negative_connector_attention_mask = None - if self.do_classifier_free_guidance: - split_batch = batch_size * num_videos_per_prompt - negative_connector_prompt_embeds = connector_prompt_embeds[:split_batch] - positive_connector_prompt_embeds = connector_prompt_embeds[split_batch:] - negative_connector_audio_prompt_embeds = connector_audio_prompt_embeds[:split_batch] - positive_connector_audio_prompt_embeds = connector_audio_prompt_embeds[split_batch:] - negative_connector_attention_mask = connector_attention_mask[:split_batch] - positive_connector_attention_mask = connector_attention_mask[split_batch:] - - # ---- 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) - audio_num_frames = self._resolve_audio_latent_length(audio_num_frames, audio_latents) - - 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) - video_audio_scheduler = _VideoAudioScheduler(self.scheduler, audio_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 single-rank batch=2 CFG ---- - # Connector outputs are already batch=2 (neg+pos concatenated before connector call) - if self.do_classifier_free_guidance and not cfg_parallel_ready: - 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 - - if cfg_parallel_ready: - latent_model_input = latents.to(positive_connector_prompt_embeds.dtype) - audio_latent_model_input = audio_latents.to(positive_connector_prompt_embeds.dtype) - ts = t.expand(latent_model_input.shape[0]) - positive_kwargs = { - "hidden_states": latent_model_input, - "audio_hidden_states": audio_latent_model_input, - "encoder_hidden_states": positive_connector_prompt_embeds, - "audio_encoder_hidden_states": positive_connector_audio_prompt_embeds, - "timestep": ts, - "sigma": ts, - "encoder_attention_mask": positive_connector_attention_mask, - "audio_encoder_attention_mask": positive_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, - "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, - } - noise_pred_video, noise_pred_audio = self.predict_noise_with_parallel_cfg( - true_cfg_scale=guidance_scale, - positive_kwargs=positive_kwargs, - negative_kwargs=negative_kwargs, - cfg_normalize=False, - video_latents=latents, - audio_latents=audio_latents, - video_sigma=self.scheduler.sigmas[i], - audio_sigma=audio_scheduler.sigmas[i], - ) - - latents, audio_latents = self.scheduler_step_maybe_with_cfg( - (noise_pred_video, noise_pred_audio), - (t, t), - (latents, audio_latents), - do_true_cfg=self.do_classifier_free_guidance, - per_request_scheduler=video_audio_scheduler, - ) - latents, audio_latents = self._synchronize_cfg_parallel_step_output( - (latents, audio_latents), - do_true_cfg=self.do_classifier_free_guidance, - ) - else: - 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() - - if self.do_classifier_free_guidance: - noise_pred_video_uncond, noise_pred_video_cond = noise_pred_video.chunk(2) - noise_pred_video = self._combine_x0_space_cfg( - latents, - noise_pred_video_cond, - noise_pred_video_uncond, - self.scheduler.sigmas[i], - guidance_scale, - ) - - noise_pred_audio_uncond, noise_pred_audio_cond = noise_pred_audio.chunk(2) - noise_pred_audio = self._combine_x0_space_cfg( - audio_latents, - noise_pred_audio_cond, - noise_pred_audio_uncond, - audio_scheduler.sigmas[i], - guidance_scale, - ) - - 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 - - latents = latents.to(self.vae.dtype) - video = self.vae.decode(latents, timestep_decode, return_dict=False)[0] - if video.numel() > 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) - - return split_diffusion_output_by_request( - DiffusionOutput( - output=(video, audio), - stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, - ), + return self._forward_impl( req, - num_outputs_per_prompt=num_videos_per_prompt, + request_inputs, + noise_scale=noise_scale, + sigmas=sigmas, + timesteps=timesteps, + attention_kwargs=attention_kwargs, ) 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." - ) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py index d30ef5464..bcccd7596 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_3_image2video.py @@ -1,16 +1,452 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Re-exports for LTX-2.3 I2V pipeline variants. +"""LTX-2.3 image-to-video pipeline.""" -The registry loads pipeline classes by (mod_folder, mod_relname, cls_name). -This module exposes the I2V class names so the registry can find them. -""" +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 ( - LTX23ImageToVideoPipeline, - get_ltx2_post_process_func, # noqa: F401 - loaded by registry via getattr + 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",