anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Feature] Add official LTX multimodal guidance and precision guard (#5148)

Browse Source 
Signed-off-by: mglyn <1203789601@qq.com>
This commit is contained in:
Mu GuanLin authored and GitHub committed 2026-07-28 10:44:25 +00:00
1 parent 66a9b0c84f
commit 12d75cf63a
51 files changed
+4547 -3062
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.buildkite/cuda/test-nightly.yml
+1 -1
View File
@@ -426,7 +426,7 @@ steps:
commands:
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_t2v/test_hunyuanvideo15_t2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
- pytest -s -v 'tests/e2e/accuracy/ltx/test_ltx_official_similarity.py::test_ltx_one_stage_matches_official[ltx2_3_i2v]' -m full_model --run-level full_model
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
mirror_hardwares: h100_2
docs/models/supported_models.md
+2 -6
View File
@@ -41,12 +41,8 @@ th {
| `Wan22S2VPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX23Pipeline` | LTX-2.3-T2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX2Pipeline` | LTX-2 / LTX-2.3 one-stage T2V and I2V | `Lightricks/LTX-2`, `diffusers/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
| `LTX2DistilledPipeline` | LTX-2 distilled two-stage T2V and I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LingBotVideoPipeline` | LingBot-Video dense and MoE T2V | `robbyant/lingbot-video-dense-1.3b`, `robbyant/lingbot-video-moe-30b-a3b` | ✅︎ | | | |
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
docs/serving/videos_api.md
+1 -1
View File
@@ -73,7 +73,7 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
| `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
| `width` | integer | model default | Output video width |
| `height` | integer | model default | Output video height |
| `num_frames` | integer | model default | Number of generated frames |
| `num_frames` | integer | 1 | Number of generated frames |
| `fps` | integer | model default | Output frames per second |
| `num_inference_steps` | integer | model default | Number of diffusion steps |
| `guidance_scale` | number | null | CFG guidance scale for the low-noise stage |
docs/user_guide/examples/offline_inference/image_to_video.md
+7 -19
View File
@@ -4,7 +4,7 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
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
LTX-2, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
offline inference API.
## Local CLI Usage
@@ -53,34 +53,22 @@ python image_to_video.py \
--output i2v_output.mp4
```
### LTX-2.3
### LTX-2
```bash
python image_to_video.py \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23ImageToVideoPipeline \
--model Lightricks/LTX-2 \
--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
--output ltx2_i2v_output.mp4
```
Use the Diffusers-format checkpoint `diffusers/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.
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, T2V, defaults, and advanced options.
Key arguments:
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or
LTX-2/LTX-2.3).
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or LTX-2).
- `--image`: Path to input image (required).
- `--prompt`: Text description of desired motion/animation.
- `--height/--width`: Output resolution (auto-calculated from image if not set).
docs/user_guide/examples/offline_inference/text_to_video.md
+8 -38
View File
@@ -12,7 +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 |
| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `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 |
@@ -35,45 +35,17 @@ python text_to_video.py \
--output t2v_out.mp4
```
LTX2 example:
### LTX-2
```bash
python text_to_video.py \
--model "Lightricks/LTX-2" \
--prompt "A cinematic close-up of ocean waves at golden hour." \
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
--height 512 \
--width 768 \
--num-frames 121 \
--num-inference-steps 40 \
--guidance-scale 4.0 \
--frame-rate 24 \
--output ltx2_out.mp4
```
### LTX-2.3
```bash
python text_to_video.py \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23Pipeline \
--model Lightricks/LTX-2 \
--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
--output ltx2_output.mp4
```
Use the Diffusers-format checkpoint `diffusers/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.
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, I2V, defaults, and advanced options.
### HunyuanVideo-1.5 (480p)
@@ -135,8 +107,7 @@ 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.
- `--model-class-name`: Optional explicit pipeline override.
- `--prompt`: text description (string).
- `--height/--width`: output resolution. Default depends on model.
- `--num-frames`: number of frames. Default depends on model.
@@ -151,8 +122,7 @@ 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; LTX-2.3 outputs 48kHz audio).
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
- `--flow-shift`: scheduler flow_shift parameter.
docs/user_guide/examples/online_serving/text_to_video.md
+4 -94
View File
@@ -13,7 +13,6 @@ This example demonstrates how to deploy text-to-video models for online video ge
| Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` |
| Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
| LTX-2 | `Lightricks/LTX-2` |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` |
## Wan2.2 T2V
@@ -192,7 +191,7 @@ curl -X POST http://localhost:8091/v1/videos \
| `negative_prompt` | str | None | Negative prompt |
| `width` | int | None | Video width in pixels |
| `height` | int | None | Video height in pixels |
| `num_frames` | int | None | Number of frames to generate |
| `num_frames` | int | 1 | Number of frames to generate; set explicitly for video generation |
| `fps` | int | None | Frames per second for output video |
| `num_inference_steps` | int | None | Number of denoising steps |
| `guidance_scale` | float | None | CFG guidance scale (low-noise stage) |
@@ -290,101 +289,12 @@ done
## LTX-2
### Start Server
#### Basic Start
```bash
vllm serve Lightricks/LTX-2 --omni --port 8098 \
--enforce-eager --boundary-ratio 1.0
vllm serve Lightricks/LTX-2 --omni --port 8098
```
For multi-GPU memory reduction, you can enable HSDP:
```bash
vllm serve Lightricks/LTX-2 --omni --port 8098 \
--enforce-eager --boundary-ratio 1.0 \
--use-hsdp --hsdp-shard-size 2
```
#### Start with Optimization Presets
Use the LTX-2 startup script with built-in optimization presets:
```bash
# Baseline (1 GPU, eager)
bash run_server_ltx2.sh baseline
# 4-GPU Ulysses sequence parallelism (lossless)
bash run_server_ltx2.sh ulysses4
# Cache-DiT lossy acceleration (1 GPU, ~1.4× speedup)
bash run_server_ltx2.sh cache-dit
# Best combo: 4-GPU Ulysses SP + Cache-DiT (~2.2× speedup)
bash run_server_ltx2.sh best-combo
```
#### Optimization Benchmarks
Benchmarked on H800, online serving (480×768, 41 frames, 20 steps, `seed=42`).
"Inference" is the server-reported inference time; excludes HTTP/poll overhead.
| Preset | Server Command | Inference (s) | Speedup | Type |
|--------|---------------|---------------|---------|------|
| `baseline` | `--enforce-eager` | 10.3 | 1.00× | — |
| `compile` | *(default, no --enforce-eager)* | ~10.3 (warm) | ~1.00× | Lossless |
| `ulysses4` | `--enforce-eager --usp 4` | ~10.3 | ~1.00× | Lossless |
| `cache-dit` | `--enforce-eager --cache-backend cache_dit` | 7.4 avg | ~1.4× | Lossy |
| `best-combo` | `--enforce-eager --usp 4 --cache-backend cache_dit` | 4.7 avg | **~2.2×** | Lossless + Lossy |
**Observations**:
- **torch.compile**: On H800, warm-request inference time matches the eager baseline (~10.3s).
The first request pays ~6s compilation overhead. Benefit depends on model architecture and GPU.
- **Ulysses SP (4 GPU)**: No measurable speedup alone for 41-frame generation at this resolution.
Communication overhead outweighs gains at this sequence length.
- **Cache-DiT**: Inference varies per request (6–10s) due to dynamic caching decisions.
Average is ~7.4s (~1.4× speedup) with slight quality tradeoff.
- **Best combo**: 4-GPU Ulysses SP + Cache-DiT synergize well — Cache-DiT reduces per-step
computation, making the communication overhead of Ulysses SP worthwhile. Average ~4.7s
(~2.2× speedup).
- **FP8 quantization**: Reduces VRAM but does not speed up LTX-2 on H800 (compute-bound).
**Deployment Recommendations**:
- For **production with quality priority**: use `baseline` with `--enforce-eager`
- For **maximum throughput** (4 GPUs, quality tradeoff): use `best-combo` (~2.2× speedup)
- For **single-GPU throughput**: use `cache-dit` (~1.4× speedup)
- `--enforce-eager` is recommended to avoid torch.compile warmup latency on first request
### Send Requests (curl)
```bash
# Using the provided script
bash run_curl_ltx2.sh
# Or directly
curl -sS -X POST http://localhost:8098/v1/videos \
-H "Accept: application/json" \
-F "prompt=A serene lakeside sunrise with mist over the water." \
-F "width=768" \
-F "height=480" \
-F "num_frames=41" \
-F "fps=24" \
-F "num_inference_steps=20" \
-F "guidance_scale=3.0" \
-F "seed=42"
```
## LTX-2.3
LTX-2.3 uses a Diffusers-format checkpoint and a dedicated pipeline class:
```bash
vllm serve diffusers/LTX-2.3-Diffusers --omni --port 8098 \
--model-class-name LTX23Pipeline
```
Send requests through the same `/v1/videos` API shown in the LTX-2 example.
See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, requests, defaults, and advanced options.
## Example materials
examples/offline_inference/image_to_video/README.md
+12 -41
View File
@@ -1,7 +1,7 @@
# Image-To-Video
This shared example generates videos from images with VACE, Wan2.2,
LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
This shared example generates videos from images with VACE, Wan2.2, LTX-2,
HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
- `image_to_video.py`: command-line script for single video generation with advanced options.
@@ -26,7 +26,7 @@ This folder provides a unified CLI script for image-to-video generation using vL
| `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
| `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
| `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism |
!!! info
Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
@@ -105,7 +105,7 @@ python image_to_video.py \
| Argument | Type | Default | Description |
| -------- | ---- | ------- | ----------- |
| `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
| `--model-class-name` | str | `None` | Optional pipeline override |
| `--image` | str | (required) | Path to input image |
| `--prompt` | str | `""` | Text description of desired motion/animation |
| `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
@@ -195,24 +195,19 @@ python image_to_video.py \
--output hunyuan_i2v.mp4
```
### LTX2 Image-to-Video
### LTX-2
```bash
python image_to_video.py \
--model /path/to/LTX-2 \
--model-class-name LTX2ImageToVideoPipeline \
--model Lightricks/LTX-2 \
--image cherry_blossom.jpg \
--prompt "A cinematic dolly shot of cherry blossoms" \
--height 512 \
--width 768 \
--num-frames 121 \
--num-inference-steps 40 \
--guidance-scale 4.0 \
--frame-rate 24 \
--fps 24 \
--output ltx2_i2v.mp4
```
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, T2V, defaults, and advanced options.
## Advanced Features
### CFG Parallel
@@ -233,30 +228,6 @@ python image_to_video.py \
--output i2v_cached.mp4
```
### LTX-2.3
```bash
python image_to_video.py \
--model diffusers/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 the Diffusers-format checkpoint `diffusers/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
@@ -274,8 +245,8 @@ python image_to_video.py \
Key arguments:
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V,
LTX-2/LTX-2.3, Cosmos3, or VACE).
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, LTX-2,
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.
@@ -290,7 +261,7 @@ Key arguments:
- `--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.
- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
- `--flow-shift`: Scheduler flow shift (default: Wan 5.0, Cosmos3 10.0). LTX does not consume this option.
- `--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).
examples/offline_inference/image_to_video/image_to_video.py
+31 -16
View File
@@ -23,16 +23,14 @@ Usage:
# LTX2 image-to-video
python image_to_video.py --model /path/to/LTX-2 \
--model-class-name LTX2ImageToVideoPipeline \
--image input.jpg --prompt "A cinematic dolly shot of a boat" \
--num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
--num-frames 121 --num-inference-steps 40 \
--frame-rate 24 --fps 24 --output ltx2_i2v.mp4
# LTX-2.3 image-to-video
python image_to_video.py --model diffusers/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 \
--height 512 --width 768 --num-frames 121 --num-inference-steps 30 \
--frame-rate 24 --fps 24 --output ltx23_i2v.mp4
# HunyuanVideo-1.5 I2V (480p)
@@ -100,7 +98,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument(
"--model-class-name",
default=None,
help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).",
help="Override model class name (LTX checkpoints default to LTX2Pipeline).",
)
parser.add_argument(
"--deploy-config",
@@ -117,7 +115,7 @@ def parse_args() -> argparse.Namespace:
help="Path to a reference image. Repeat to provide multiple references.",
)
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
parser.add_argument(
@@ -331,7 +329,10 @@ 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 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"}
model_class_name_lower = (model_class_name or "").lower()
is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name
is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
image = PIL.Image.open(args.image).convert("RGB") if args.image else None
@@ -356,13 +357,23 @@ def main():
1280 * 720,
16,
)
elif is_ltx2_distilled:
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
24,
None,
121,
8,
None,
1024 * 1536,
64,
)
elif is_ltx2:
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
24,
4.0,
None,
121,
40,
5.0,
30 if is_ltx23 else 40,
None,
512 * 768,
32,
)
@@ -445,7 +456,6 @@ def main():
vae_use_slicing=args.vae_use_slicing,
vae_use_tiling=args.vae_use_tiling,
boundary_ratio=args.boundary_ratio,
flow_shift=flow_shift,
diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
@@ -460,6 +470,8 @@ def main():
)
if args.deploy_config:
omni_kwargs["deploy_config"] = args.deploy_config
if flow_shift is not None:
omni_kwargs["flow_shift"] = flow_shift
if args.quantization is not None:
omni_kwargs["quantization"] = args.quantization
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
@@ -492,10 +504,14 @@ def main():
print(f" Video size: {width}x{height}")
print(f"{'=' * 60}\n")
negative_prompt = args.negative_prompt
if negative_prompt is None and not is_ltx2:
# Preserve the historical empty-prompt behavior for non-LTX examples.
negative_prompt = ""
prompt_dict = build_image_to_video_prompt(
model_class_name=model_class_name,
prompt=args.prompt,
negative_prompt=args.negative_prompt,
negative_prompt=negative_prompt,
media_inputs=media_inputs,
height=height,
width=width,
@@ -511,11 +527,10 @@ def main():
num_inference_steps=num_inference_steps,
num_frames=num_frames,
frame_rate=frame_rate,
extra_args={
"sample_solver": args.sample_solver,
"flow_shift": flow_shift,
},
extra_args={"sample_solver": args.sample_solver},
)
if flow_shift is not None:
sampling_params.extra_args["flow_shift"] = flow_shift
# Route model-specific knobs through extra_body, filtered against the model's
# declared extra_body_params. Models without a declaration only forward explicit
examples/offline_inference/text_to_video/text_to_video.md
+8 -38
View File
@@ -8,7 +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 |
| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
| `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) |
@@ -55,45 +55,17 @@ python text_to_video.py \
--output vace_t2v_output.mp4
```
LTX2 example:
### LTX-2
```bash
python text_to_video.py \
--model "Lightricks/LTX-2" \
--prompt "A cinematic close-up of ocean waves at golden hour." \
--negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
--height 512 \
--width 768 \
--num-frames 121 \
--num-inference-steps 40 \
--guidance-scale 4.0 \
--frame-rate 24 \
--output ltx2_out.mp4
```
### LTX-2.3
```bash
python text_to_video.py \
--model diffusers/LTX-2.3-Diffusers \
--model-class-name LTX23Pipeline \
--model Lightricks/LTX-2 \
--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
--output ltx2_output.mp4
```
Use the Diffusers-format checkpoint `diffusers/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.
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, I2V, defaults, and advanced options.
### HunyuanVideo-1.5 (480p)
@@ -212,8 +184,7 @@ 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.
- `--model-class-name`: Optional explicit pipeline override.
- `--prompt`: text description (string).
- `--height/--width`: output resolution. Default depends on model.
- `--num-frames`: number of frames. Default depends on model.
@@ -229,8 +200,7 @@ 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; LTX-2.3 outputs 48kHz audio).
- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
- `--quantization`: quantization method (such as `fp8` for FP8).
- `--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).
examples/offline_inference/text_to_video/text_to_video.py
+39 -5
View File
@@ -65,11 +65,42 @@ _MODEL_PRESETS = {
"fps": 16,
"output": "helios_output.mp4",
},
"ltx2": {
"height": 512,
"width": 768,
"num_frames": 121,
"num_inference_steps": 40,
"fps": 24,
"output": "ltx2_output.mp4",
},
"ltx2_distilled": {
"height": 1024,
"width": 1536,
"num_frames": 121,
"num_inference_steps": 8,
"fps": 24,
"output": "ltx2_distilled_output.mp4",
},
"ltx23": {
"height": 512,
"width": 768,
"num_frames": 121,
"num_inference_steps": 30,
"fps": 24,
"output": "ltx23_output.mp4",
},
}
def _detect_preset(model: str) -> dict:
def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
model_lower = model.lower()
class_lower = (model_class_name or "").lower()
if "distilled" in class_lower or "distilled" in model_lower:
return _MODEL_PRESETS["ltx2_distilled"]
if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
return _MODEL_PRESETS["ltx23"]
if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
return _MODEL_PRESETS["ltx2"]
if "vace" in model_lower:
return _MODEL_PRESETS["vace"]
if "cosmos" in model_lower:
@@ -122,7 +153,7 @@ def parse_args() -> argparse.Namespace:
parser.add_argument(
"--model-class-name",
default=None,
help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
help="Override model class name (e.g., LTX2Pipeline).",
)
parser.add_argument(
"--deploy-config",
@@ -130,7 +161,7 @@ def parse_args() -> argparse.Namespace:
help="Optional deploy config YAML to use for pipeline-backed runs.",
)
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.")
parser.add_argument(
"--extra-body",
type=parse_extra_body,
@@ -331,7 +362,7 @@ def main():
args = parse_args()
model_class_name = args.model_class_name
preset = _detect_preset(args.model)
preset = _detect_preset(args.model, model_class_name)
for key, default_val in preset.items():
if getattr(args, key.replace("-", "_"), None) is None:
setattr(args, key.replace("-", "_"), default_val)
@@ -422,8 +453,11 @@ def main():
print(f"{'=' * 60}\n")
prompt_dict = {"prompt": args.prompt}
if args.negative_prompt:
if args.negative_prompt is not None:
prompt_dict["negative_prompt"] = args.negative_prompt
elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
# Preserve the historical empty-prompt behavior for non-LTX examples.
prompt_dict["negative_prompt"] = ""
sampling_kwargs = dict(
height=args.height,
examples/online_serving/text_to_video/README.md
+3 -73
View File
@@ -248,82 +248,12 @@ done
## LTX-2
### Start Server
#### Basic Start
```bash
vllm serve Lightricks/LTX-2 --omni --port 8098 \
--enforce-eager --boundary-ratio 1.0
vllm serve Lightricks/LTX-2 --omni --port 8098
```
#### Start with Optimization Presets
Use the LTX-2 startup script with built-in optimization presets:
```bash
# Baseline (1 GPU, eager)
bash run_server_ltx2.sh baseline
# 4-GPU Ulysses sequence parallelism (lossless)
bash run_server_ltx2.sh ulysses4
# Cache-DiT lossy acceleration (1 GPU, ~1.4× speedup)
bash run_server_ltx2.sh cache-dit
# Best combo: 4-GPU Ulysses SP + Cache-DiT (~2.2× speedup)
bash run_server_ltx2.sh best-combo
```
#### Optimization Benchmarks
Benchmarked on H800, online serving (480×768, 41 frames, 20 steps, `seed=42`).
"Inference" is the server-reported inference time; excludes HTTP/poll overhead.
| Preset | Server Command | Inference (s) | Speedup | Type |
|--------|---------------|---------------|---------|------|
| `baseline` | `--enforce-eager` | 10.3 | 1.00× | — |
| `compile` | *(default, no --enforce-eager)* | ~10.3 (warm) | ~1.00× | Lossless |
| `ulysses4` | `--enforce-eager --usp 4` | ~10.3 | ~1.00× | Lossless |
| `cache-dit` | `--enforce-eager --cache-backend cache_dit` | 7.4 avg | ~1.4× | Lossy |
| `best-combo` | `--enforce-eager --usp 4 --cache-backend cache_dit` | 4.7 avg | **~2.2×** | Lossless + Lossy |
**Observations**:
- **torch.compile**: On H800, warm-request inference time matches the eager baseline (~10.3s).
The first request pays ~6s compilation overhead. Benefit depends on model architecture and GPU.
- **Ulysses SP (4 GPU)**: No measurable speedup alone for 41-frame generation at this resolution.
Communication overhead outweighs gains at this sequence length.
- **Cache-DiT**: Inference varies per request (6–10s) due to dynamic caching decisions.
Average is ~7.4s (~1.4× speedup) with slight quality tradeoff.
- **Best combo**: 4-GPU Ulysses SP + Cache-DiT synergize well — Cache-DiT reduces per-step
computation, making the communication overhead of Ulysses SP worthwhile. Average ~4.7s
(~2.2× speedup).
- **FP8 quantization**: Reduces VRAM but does not speed up LTX-2 on H800 (compute-bound).
**Deployment Recommendations**:
- For **production with quality priority**: use `baseline` with `--enforce-eager`
- For **maximum throughput** (4 GPUs, quality tradeoff): use `best-combo` (~2.2× speedup)
- For **single-GPU throughput**: use `cache-dit` (~1.4× speedup)
- `--enforce-eager` is recommended to avoid torch.compile warmup latency on first request
### Send Requests (curl)
```bash
# Using the provided script
bash run_curl_ltx2.sh
# Or directly
curl -sS -X POST http://localhost:8098/v1/videos \
-H "Accept: application/json" \
-F "prompt=A serene lakeside sunrise with mist over the water." \
-F "width=768" \
-F "height=480" \
-F "num_frames=41" \
-F "fps=24" \
-F "num_inference_steps=20" \
-F "guidance_scale=3.0" \
-F "seed=42"
```
See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
pipeline selection, requests, defaults, and advanced options.
## Helios
examples/online_serving/text_to_video/run_curl_ltx2.sh
-1
View File
@@ -22,7 +22,6 @@ create_response=$(
-F "num_frames=41" \
-F "fps=24" \
-F "num_inference_steps=20" \
-F "guidance_scale=3.0" \
-F "seed=42"
)
examples/online_serving/text_to_video/run_server_ltx2.sh
-8
View File
@@ -21,8 +21,6 @@ set -euo pipefail
MODEL="${MODEL:-Lightricks/LTX-2}"
PORT="${PORT:-8098}"
FLOW_SHIFT="${FLOW_SHIFT:-1.0}"
BOUNDARY_RATIO="${BOUNDARY_RATIO:-1.0}"
PRESET="${1:-baseline}"
@@ -66,19 +64,13 @@ case "$PRESET" in
echo "Environment variables:"
echo " MODEL - Model path (default: Lightricks/LTX-2)"
echo " PORT - Server port (default: 8098)"
echo " FLOW_SHIFT - Scheduler flow shift (default: 1.0)"
echo " BOUNDARY_RATIO - Boundary ratio (default: 1.0)"
exit 1
;;
esac
echo "Model: $MODEL"
echo "Port: $PORT"
echo "Flow shift: $FLOW_SHIFT"
echo "Boundary ratio: $BOUNDARY_RATIO"
vllm serve "$MODEL" --omni \
--port "$PORT" \
--flow-shift "$FLOW_SHIFT" \
--boundary-ratio "$BOUNDARY_RATIO" \
"${EXTRA_ARGS[@]}"
recipes/LTX/LTX-2.3.md
-154
View File
@@ -1,154 +0,0 @@
# LTX-2.3
> 22B parameter text-to-video + audio generation model served via vLLM-Omni
## Summary
- Vendor: Lightricks
- Model: `diffusers/LTX-2.3-Diffusers`
- 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 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
- Upstream raw checkpoints: <https://huggingface.co/Lightricks/LTX-2.3>
- Diffusers-format checkpoint: <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
## Serving
### Command
```bash
vllm serve diffusers/LTX-2.3-Diffusers \
--omni \
--model-class-name LTX23Pipeline \
--stage-init-timeout 600
```
For image-to-video:
```bash
vllm serve diffusers/LTX-2.3-Diffusers \
--omni \
--model-class-name LTX23ImageToVideoPipeline \
--stage-init-timeout 600
```
### T2V Verification
```bash
# Health check
curl http://localhost:8000/health
# Generate a 3-second video (81 frames at 24fps)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=A majestic bald eagle soaring over a misty mountain valley at dawn, golden sunlight breaking through clouds" \
-F "negative_prompt=blurry, low quality, distorted, watermark" \
-F "model=diffusers/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"
# Generate a 10-second video (241 frames)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=A cozy Japanese ramen shop at night in the rain, steam rising from bowls, neon signs reflecting on wet cobblestone streets" \
-F "model=diffusers/LTX-2.3-Diffusers" \
-F "num_frames=241" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-F "guidance_scale=4.0"
# Generate a 20-second video (481 frames)
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=An underwater coral reef teeming with tropical fish, sea turtles gliding gracefully, National Geographic documentary style" \
-F "model=diffusers/LTX-2.3-Diffusers" \
-F "num_frames=481" \
-F "fps=24" \
-F "size=768x512" \
-F "num_inference_steps=30" \
-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=diffusers/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: 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
`diffusers/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`: 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:
- 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
## GPU
### 1x NVIDIA RTX PRO 6000 Blackwell (96GB)
#### Environment
- OS: Ubuntu 22.04
- Python: 3.10+
- Driver / runtime: CUDA 13.0, Driver 580.126.09
- vLLM version: Match the PR or release checkout being validated
- vLLM-Omni version or commit: Use the commit you are deploying from
### Benchmarking
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_vllm_omni.json` with
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. Each case in that file includes a JSON **`mark`** (`hardware_marks` + `full_model` + `diffusion`) for local `-m` filtering. That config captures the
single-device eager baseline and CFG-parallel=2 case for a small 384x512,
25-frame, 20-step workload.
recipes/LTX/LTX-2.md
+295 -113
View File
@@ -1,121 +1,303 @@
# LTX-2
# LTX-2 Family
> Text-to-video and image-to-video serving
> LTX-2 and LTX-2.3 text-to-video and image-to-video generation with synchronized audio
## Summary
## Pipelines
- Vendor: Lightricks
- Model: `Lightricks/LTX-2`
- Task: Text-to-video generation and image-to-video generation
- Mode: Online serving with the OpenAI-compatible video API
- Maintainer: Community
| Model | Checkpoint | Task | `--model-class-name` | Batching |
|---|---|---|---|---|
| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
## When to use this recipe
`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
image selects I2V. Both one-stage repositories declare this class, so
`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
Use this recipe when you want to deploy the base `Lightricks/LTX-2` model
for both text-to-video and image-to-video generation. The same model ID is used for
both tasks; switch the pipeline class with `--model-class-name` depending on
whether you want T2V or I2V.
`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
or provide one initial image for I2V. As in the official LTX pipelines,
distilled mode is selected explicitly rather than inferred from the checkpoint
name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
## API Migration
Only `req` may be passed positionally to `LTX2Pipeline`; every optional
`forward` argument is keyword-only:
```python
# No longer supported
pipe(req, prompt)
pipe(req, image, prompt)
# Supported
pipe(req, prompt=prompt)
pipe(req, image=image, prompt=prompt)
```
The consolidation also removes these registry names without aliases:
| Removed name | Replacement |
|---|---|
| `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
| `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
| `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
Passing any second positional argument now raises `TypeError`. These changes
affect direct Python callers and explicit `--model-class-name` overrides;
offline and serving entrypoints already use named fields and are unaffected.
## One-Stage Defaults
| Parameter | LTX-2 | LTX-2.3 |
|---|---:|---:|
| Width × height | 768 × 512 | 768 × 512 |
| Frames / frame rate | 121 / 24 | 121 / 24 |
| Denoise steps | 40 | 30 |
| Video/audio CFG | 3.0 / 7.0 | 3.0 / 7.0 |
| Video/audio STG | 1.0 / 1.0 | 1.0 / 1.0 |
| Video/audio modality guidance | 3.0 / 3.0 | 3.0 / 3.0 |
| Video/audio rescale | 0.7 / 0.7 | 0.7 / 0.7 |
| Video/audio STG blocks | `[29]` / `[29]` | `[28]` / `[28]` |
The model recipe also supplies the default negative prompt. Top-level
`guidance_scale` overrides both video and audio CFG values. The online video
API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
default to `121`.
## Distilled Two-Stage Defaults
| Parameter | Value |
|---|---:|
| Final width × height | 1536 × 1024 |
| Stage 1 width × height | 768 × 512 |
| Frames / frame rate | 121 / 24 |
| Stage 1 / Stage 2 denoise steps | 8 / 3 |
| Guidance | Fixed positive-only |
The dimensions passed to the API are the final output dimensions and must be
divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
`8` when supplied.
## Serving
Start either checkpoint:
```bash
vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
```
```bash
vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
```
Start the distilled checkpoint with its two-stage entry:
```bash
vllm serve rootonchair/LTX-2-19b-distilled --omni \
--model-class-name LTX2DistilledPipeline --stage-init-timeout 600
```
The same server handles T2V and I2V. A T2V request using the selected recipe's
default guidance is:
```bash
curl -X POST http://localhost:8000/v1/videos/sync \
-F "prompt=A cinematic close-up of ocean waves at golden hour." \
-F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
-F "size=768x512" \
-F "num_frames=121" \
-F "fps=24" \
-F "seed=42" \
-o ltx_t2v.mp4
```
For I2V, provide exactly one initial image:
```bash
curl -X POST http://localhost:8000/v1/videos/sync \
-F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
-F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
-F "input_reference=@/absolute/path/to/reference.png" \
-F "size=768x512" \
-F "num_frames=121" \
-F "fps=24" \
-F "seed=42" \
-o ltx_i2v.mp4
```
Use `image_reference` for a URL or JSON-safe image reference. Do not provide it
together with `input_reference`.
## Guidance
LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
cross-modality guidance, and rescaling:
| Parameter | Default | Effect | Alias |
|---|---:|---|---|
| `video_cfg_scale` | 3.0 | Video text CFG; `1.0` disables it | `video_cfg_guidance_scale` |
| `audio_cfg_scale` | 7.0 | Audio text CFG; `1.0` disables it | `audio_cfg_guidance_scale` |
| `video_stg_scale` | 1.0 | Video STG; `0.0` disables it | `video_stg_guidance_scale` |
| `audio_stg_scale` | 1.0 | Audio STG; `0.0` disables it | `audio_stg_guidance_scale` |
| `video_modality_scale` | 3.0 | Audio-to-video guidance; `1.0` disables it | `a2v_guidance_scale` |
| `audio_modality_scale` | 3.0 | Video-to-audio guidance; `1.0` disables it | `v2a_guidance_scale` |
| `video_rescale_scale` | 0.7 | Video guidance rescale; `0.0` disables it | — |
| `audio_rescale_scale` | 0.7 | Audio guidance rescale; `0.0` disables it | — |
| `video_stg_blocks` | `[29]` / `[28]` | Perturbed video transformer blocks | — |
| `audio_stg_blocks` | `[29]` / `[28]` | Perturbed audio transformer blocks | — |
The STG block defaults are shown as LTX-2 / LTX-2.3. A canonical name wins
over its alias. Explicit top-level `guidance_scale` wins over both CFG fields;
omit it for independent video/audio CFG.
Pass these fields in online `extra_params` or offline `--extra-body`:
```bash
# Add to the curl request above
-F 'extra_params={"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
# Add to text_to_video.py or image_to_video.py
--extra-body '{"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
```
### Python API
Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the
resulting `DiffusionRequestBatch` as the only positional argument:
```python
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
params = OmniDiffusionSamplingParams(
width=768,
height=512,
num_frames=121,
frame_rate=24.0,
num_inference_steps=30,
seed=42,
extra_args={"video_cfg_scale": 3.0, "audio_cfg_scale": 7.0},
)
req = DiffusionRequestBatch([
OmniDiffusionRequest(
prompt={"prompt": "Cherry blossoms moving in a light breeze"},
sampling_params=params,
request_id="ltx-example",
)
])
t2v_output = pipe(req)
i2v_output = pipe(req, image=image)
```
Direct keyword arguments are low-level fallbacks, for example
`pipe(req, prompt=prompt, width=768)`. Request fields take precedence unless
noted below.
### Complete `forward` Surface
| Argument | Type/default | Meaning and constraints |
|---|---|---|
| `req` | `DiffusionRequestBatch`, required | Only positional argument; contains prompts and per-request sampling parameters. |
| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
| `width` | `int`, `None` | Same precedence and alignment as `height`. |
| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
| `decode_noise_scale` | float or list, `None` | Same list rules; request value wins. Defaults to `decode_timestep`. |
| `output_type` | string, `"np"` | Request value wins. `"np"` decodes output; `"latent"` skips VAE/vocoder decode. |
| `return_dict` | `bool`, `True` | Compatibility slot; only `True` is accepted. |
| `attention_kwargs` | dict, `None` | Per-call values are unsupported; configure attention at engine startup. |
| `max_sequence_length` | `int`, `None` | Request → direct value → tokenizer limit. |
Request-object naming differs only in a few places: `num_videos_per_prompt`
maps to `num_outputs_per_prompt`; images and prompt text/embeddings live in the
request prompt payload; LTX guidance fields live in sampling `extra_args`.
### Recipe-Specific Request Capabilities
| Override | One-stage | Distilled two-stage |
|---|---|---|
| Guidance | Supported | Fixed positive-only |
| Negative prompt/embeddings | Supported | Rejected |
| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
| Video/audio latents | Supported | Rejected |
These capability checks apply equally to direct `forward` keywords and
values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
being silently ignored.
### Custom Sigma Schedules
One-stage Python requests may set final scheduler boundaries directly:
```python
params = OmniDiffusionSamplingParams(sigmas=[1.0, 0.75, 0.5, 0.25])
```
Each nonterminal value produces one denoise step, and a terminal `0.0` is
appended when omitted. This schedule overrides `num_inference_steps`. All
requests in a fused batch must use the same list. The video form API and
bundled offline CLI do not currently expose `sigmas`.
### Constraints
- LTX rejects non-default `timesteps`, `flow_shift`, `guidance_rescale`,
`noise_scale`, `attention_kwargs`, and `return_dict=False`. Use final
`sigmas`, modality rescale fields, startup attention configuration, and the
standard `DiffusionOutput` respectively.
- Fused one-stage requests must resolve to identical LTX guidance and sigma
schedules. Keep concurrent requests guidance-homogeneous or use
`--max-num-seqs 1`.
- `--cfg-parallel-size 2` supports only the CFG-only plan: disable STG and
modality guidance and set both rescale fields to `0.0`.
## Operational Notes
- LTX-2 one-stage previously loaded and peaked at about 73.5 GiB on one H200
141GB; remeasure on the deployed commit and hardware.
- LTX-2.3 includes a 22B transformer, Gemma encoder, two VAEs, and a vocoder.
Start on a 96GB-class GPU or use CPU/layerwise offload on smaller devices.
- The distilled upsampler uses the configured pipeline dtype and participates
in component discovery/device placement; it remains resident during denoise
offload because it is used only at the phase boundary.
- The output audio sample rate comes from the loaded components and is not a
request parameter.
- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
`tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
- Current two-stage support is limited to the distilled checkpoint; official
LTX two-stage/HQ execution remains out of scope.
## References
- Model: <https://huggingface.co/Lightricks/LTX-2>
- User guide:
[`docs/user_guide/examples/online_serving/text_to_video.md`](../../docs/user_guide/examples/online_serving/text_to_video.md)
- Example guide:
[`examples/online_serving/text_to_video/README.md`](../../examples/online_serving/text_to_video/README.md)
[`examples/online_serving/image_to_video/README.md`](../../examples/online_serving/image_to_video/README.md)
## Hardware Support
## GPU
### 1x H200 141GB (Text-to-Video)
#### Environment
- OS: Ubuntu 22.04.5 LTS
- Python: 3.12+
- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU
- vLLM version: Match the repository requirements from your current checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Command
```bash
vllm serve Lightricks/LTX-2 \
--omni \
--model-class-name LTX2Pipeline
```
#### Verification
```bash
curl http://localhost:8000/health
curl -X POST http://localhost:8000/v1/videos \
-F "prompt=A cinematic close-up of ocean waves at golden hour." \
-F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
-F "width=512" \
-F "height=768" \
-F "num_frames=121" \
-F "fps=24" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-o ltx2_t2v_h200.mp4
```
#### Notes
- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference.
- Key flags:
- `--model-class-name LTX2Pipeline` selects the text-to-video pipeline explicitly.
- Known limitations:
- This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine.
### 1x H200 141GB (Image-to-Video)
#### Environment
- OS: Ubuntu 22.04.5 LTS
- Python: 3.12+
- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU
- vLLM version: Match the repository requirements from your current checkout
- vLLM-Omni version or commit: Use the commit you are deploying from
#### Command
```bash
vllm serve Lightricks/LTX-2 \
--omni \
--model-class-name LTX2ImageToVideoPipeline
```
#### Verification
```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 "width=512" \
-F "height=768" \
-F "num_frames=121" \
-F "fps=24" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-o ltx2_i2v_h200.mp4
```
#### Notes
- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference.
- Key flags:
- `--model-class-name LTX2ImageToVideoPipeline` switches the same `Lightricks/LTX-2` weights into image-to-video mode.
- `input_reference=@...` uploads the reference image directly
- Known limitations:
- Use `image_reference` when you want to pass a URL or JSON-safe image reference
instead of uploading a file.
- Do not send `input_reference` and `image_reference` in the same request.
- This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine.
- <https://huggingface.co/Lightricks/LTX-2>
- <https://huggingface.co/Lightricks/LTX-2.3>
- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
- [Request batching](../../docs/user_guide/diffusion/request_batching.md)
- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
- [I2V offline example](../../examples/offline_inference/image_to_video/README.md)
recipes/README.md
+1 -2
View File
@@ -46,8 +46,7 @@ recipes/
| [`inclusionAI/Ming-omni-tts.md`](./inclusionAI/Ming-omni-tts.md) | Offline + online dense Ming TTS/audio generation | 1x H100 80GB / 1x AMD MI300X (ROCm 7.2) |
| [`IndexTeam/IndexTTS-2.md`](./IndexTeam/IndexTTS-2.md) | Online serving for voice-cloned TTS with optional emotion control | 1x L4 24GB or larger CUDA GPU |
| [`krea/Krea-2.md`](./krea/Krea-2.md) | Text-to-image (Turbo + Raw), offline + online, with LoRA | 1x H100 80GB |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | LTX-2/LTX-2.3 text-to-video and image-to-video with synchronized audio | H200 141GB / 96GB-class GPU |
| [`MammothModa2/MammothModa2.md`](./MammothModa2/MammothModa2.md) | Preview and Dev text-to-image (AR → DiT); Dev text/image understanding | Preview: 1x L40S 48GB / 1x ≥40GB GPU; Dev: 1x NVIDIA GPU with sufficient cache headroom |
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
| [`Robbyant/LingBot-Video.md`](./Robbyant/LingBot-Video.md) | Native dense and MoE text-to-video serving | Dense: 1x L20X; MoE: 1x L20X (~67.7 GiB peak) |
tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py
+1 -2
View File
@@ -8,7 +8,6 @@ import torch
from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
@@ -25,7 +24,7 @@ _DMD2_BASE = {
WanT2VDMD2Pipeline: Wan22Pipeline,
WanI2VDMD2Pipeline: Wan22I2VPipeline,
LTX2T2VDMD2Pipeline: LTX2Pipeline,
LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
LTX2I2VDMD2Pipeline: LTX2Pipeline,
FluxDMD2Pipeline: FluxPipeline,
QwenImageDMD2Pipeline: QwenImagePipeline,
}
tests/diffusion/models/dmd2/test_dmd2_scheduler.py
+1 -2
View File
@@ -8,7 +8,6 @@ import torch
from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
)
@@ -27,7 +26,7 @@ _DMD2_BASE = {
WanT2VDMD2Pipeline: Wan22Pipeline,
WanI2VDMD2Pipeline: Wan22I2VPipeline,
LTX2T2VDMD2Pipeline: LTX2Pipeline,
LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
LTX2I2VDMD2Pipeline: LTX2Pipeline,
FluxDMD2Pipeline: FluxPipeline,
QwenImageDMD2Pipeline: QwenImagePipeline,
}
tests/diffusion/models/ltx2/test_ltx2_conditioning.py
+214 -30
View File
@@ -7,6 +7,7 @@ from types import SimpleNamespace
import pytest
import torch
from PIL import Image
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -16,17 +17,31 @@ def _make_ltx23_request_pipe(cls):
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 99
pipe.vae_spatial_compression_ratio = 32
pipe.vae_temporal_compression_ratio = 8
return pipe
class TestLTXImageToVideoForwardStages:
def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
def test_i2v_pil_preprocessing_matches_official_center_crop(self):
from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
image = Image.fromarray(pixels.numpy())
actual = _preprocess_i2v_pil_images(image, height=4, width=4)
expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
torch.testing.assert_close(actual, expected)
def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
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)
pipe = _make_ltx23_request_pipe(LTX2Pipeline)
image = torch.zeros(3, 8, 8)
req = DiffusionRequestBatch(
[
@@ -48,27 +63,45 @@ class TestLTXImageToVideoForwardStages:
)
seen = {}
def fake_forward_impl(req_arg, request_inputs, **kwargs):
def fake_run_recipe(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)
object.__setattr__(pipe, "_run_recipe", fake_run_recipe)
output = pipe.forward(req, noise_scale=0.5)
output = pipe.forward(req)
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
assert seen["kwargs"]["request_sigmas"] is None
def test_unified_entry_selects_t2v_without_images_and_rejects_mixed_batches(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
request_inputs = SimpleNamespace(latents=None)
t2v_req = SimpleNamespace(prompts=["text only", {"prompt": "also text only"}])
assert pipe._resolve_request_image(t2v_req, None, request_inputs) is None
mixed_req = SimpleNamespace(
prompts=[
"text only",
{"prompt": "image conditioned", "multi_modal_data": {"image": torch.zeros(3, 8, 8)}},
]
)
with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
pipe._resolve_request_image(mixed_req, None, request_inputs)
def test_denoise_timestep_kwargs_masks_video_only(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
ts = torch.tensor([2.0, 4.0])
denoise_ctx = SimpleNamespace(
conditioning_mask=torch.tensor([[1.0, 0.0]]),
@@ -79,30 +112,32 @@ class TestLTXImageToVideoForwardStages:
ts,
SimpleNamespace(cfg_parallel_ready=False),
denoise_ctx,
video_token_count=2,
audio_token_count=1,
)
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["audio_timestep"], ts[:, None])
torch.testing.assert_close(kwargs["sigma"], ts)
class TestLTXImageToVideoConditioning:
def test_ltx23_i2v_supports_image_input(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
assert LTX23ImageToVideoPipeline.support_image_input is True
assert LTX2Pipeline.support_image_input is True
def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
image = object()
assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
assert LTX2Pipeline._resolve_single_prompt_image([image]) is image
with pytest.raises(ValueError, match="exactly one image per prompt"):
LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
LTX2Pipeline._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 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
image = torch.zeros(1, 3, 4, 4)
additional = {
@@ -111,19 +146,24 @@ class TestLTXImageToVideoConditioning:
"image": torch.ones_like(image),
}
assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
assert LTX2Pipeline._resolve_additional_image(additional) is image
def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
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
object.__setattr__(
pipe,
"_encode_i2v_image_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
)
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
raise AssertionError("packed I2V latents should not be noised")
@@ -134,7 +174,7 @@ class TestLTXImageToVideoConditioning:
latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
out, conditioning_mask = pipe.prepare_latents(
image=None,
image=torch.zeros(1, 3, 1, 1),
batch_size=1,
num_channels_latents=2,
height=1,
@@ -147,14 +187,14 @@ class TestLTXImageToVideoConditioning:
)
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, latents)
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.vae_spatial_compression_ratio = 1
pipe.vae_temporal_compression_ratio = 1
@@ -165,8 +205,16 @@ class TestLTXImageToVideoConditioning:
latents_std=torch.ones(2),
config=SimpleNamespace(scaling_factor=1.0),
)
object.__setattr__(
pipe,
"_encode_i2v_image_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
)
sampled_shapes = []
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
sampled_shapes.append(tuple(shape))
return torch.ones(shape, device=device, dtype=dtype)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
@@ -175,7 +223,7 @@ class TestLTXImageToVideoConditioning:
latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
out, conditioning_mask = pipe.prepare_latents(
image=None,
image=torch.zeros(1, 3, 1, 1),
batch_size=1,
num_channels_latents=2,
height=1,
@@ -187,20 +235,64 @@ class TestLTXImageToVideoConditioning:
latents=latents,
)
torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
assert sampled_shapes == [(1, 3, 2)]
def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
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(
encode=lambda image: image,
latents_mean=torch.zeros(2),
latents_std=torch.ones(2),
)
monkeypatch.setattr(
ltx2_conditioning,
"retrieve_latents",
lambda *_args, **_kwargs: torch.tensor([[[[[10.0]]], [[[11.0]]]]]),
)
sampled_shapes = []
def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
sampled_shapes.append(tuple(shape))
return torch.ones(shape, device=device, dtype=dtype)
monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
out, conditioning_mask = pipe.prepare_latents(
image=torch.zeros(1, 3, 1, 1),
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"),
)
assert sampled_shapes == [(1, 3, 2)]
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 import LTX23ImageToVideoPipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23ImageToVideoPipeline)
pipe = object.__new__(LTX2Pipeline)
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,)
sigmas = torch.tensor([1.0, 0.5])
pipe.scheduler = FakeScheduler()
latents = torch.tensor([[[1.0], [2.0], [3.0]]])
@@ -209,11 +301,103 @@ class TestLTXImageToVideoConditioning:
out = pipe._step_video_latents_i2v(
noise_pred,
latents,
torch.tensor(0.5),
0,
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)
torch.testing.assert_close(out[:, 1:], latents[:, 1:] - 0.5 * noise_pred[:, 1:])
def test_shared_step_adapter_uses_official_euler_for_t2v(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
pipeline = SimpleNamespace(scheduler=SimpleNamespace(sigmas=torch.tensor([1.0, 0.5])))
adapter = LTXVideoAudioStepAdapter(
pipeline,
SimpleNamespace(sigmas=torch.tensor([1.0, 0.25])),
latent_num_frames=1,
latent_height=1,
latent_width=1,
image_conditioned=False,
)
video = torch.tensor([[[2.0]]])
audio = torch.tensor([[[4.0]]])
((video, audio),) = adapter.step(
(torch.tensor([[[2.0]]]), torch.tensor([[[4.0]]])),
(torch.tensor(1.0), torch.tensor(1.0)),
(video, audio),
)
torch.testing.assert_close(video, torch.tensor([[[1.0]]]))
torch.testing.assert_close(audio, torch.tensor([[[1.0]]]))
def test_shared_step_adapter_uses_official_euler_for_i2v_audio(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
class FakePipeline:
scheduler = object()
def __init__(self):
self.step_indices = []
def _step_video_latents_i2v(
self,
noise_pred,
latents,
step_index,
latent_num_frames,
latent_height,
latent_width,
):
del noise_pred, latent_num_frames, latent_height, latent_width
self.step_indices.append(step_index)
return latents
pipeline = FakePipeline()
adapter = LTXVideoAudioStepAdapter(
pipeline,
SimpleNamespace(sigmas=torch.tensor([1.0, 0.5, 0.0])),
latent_num_frames=1,
latent_height=1,
latent_width=1,
image_conditioned=True,
)
video = torch.tensor([[[1.0]]])
audio = torch.tensor([[[2.0]]])
video_velocity = torch.zeros_like(video)
audio_velocity = torch.tensor([[[4.0]]])
((video, audio),) = adapter.step(
(video_velocity, audio_velocity),
(torch.tensor(1.0), torch.tensor(1.0)),
(video, audio),
)
torch.testing.assert_close(audio, torch.tensor([[[0.0]]]))
((video, audio),) = adapter.step(
(video_velocity, audio_velocity),
(torch.tensor(0.5), torch.tensor(0.5)),
(video, audio),
)
torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
assert pipeline.step_indices == [0, 1]
def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX2Pipeline)
denoise_ctx = LTXDenoiseContext(
latents=torch.empty(1, 3, 1),
audio_latents=torch.empty(1, 1, 1),
video_coords=torch.empty(1),
audio_coords=torch.empty(1),
conditioning_mask=torch.tensor([[1.0, 0.0, 0.0]]),
)
actual = pipe._video_guidance_model_sigma(torch.tensor(0.75), denoise_ctx)
torch.testing.assert_close(actual, torch.tensor([[[0.0], [0.75], [0.75]]]))
tests/diffusion/models/ltx2/test_ltx2_latents.py
+51 -7
View File
@@ -4,7 +4,7 @@ import pytest
import torch
from vllm_omni.diffusion.models.ltx2 import ltx2_latents as latent_ops
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
@@ -23,8 +23,51 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
return pipeline
def test_prepare_video_latents_samples_directly_in_packed_token_space():
pipeline = _make_pipeline(LTX2Pipeline)
pipeline.vae_spatial_compression_ratio = 8
pipeline.vae_temporal_compression_ratio = 8
pipeline.transformer_spatial_patch_size = 2
pipeline.transformer_temporal_patch_size = 1
expected_generator = torch.Generator().manual_seed(42)
expected = torch.randn((1, 32, 16), generator=expected_generator)
actual = pipeline.prepare_latents(
batch_size=1,
num_channels_latents=4,
height=64,
width=64,
num_frames=9,
dtype=torch.float32,
device=torch.device("cpu"),
generator=torch.Generator().manual_seed(42),
)
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_prepare_audio_latents_samples_directly_in_packed_token_space():
pipeline = _make_pipeline(LTX2Pipeline)
expected_generator = torch.Generator().manual_seed(42)
expected = torch.randn((1, 3, 32), generator=expected_generator)
actual, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
num_channels_latents=2,
num_mel_bins=64,
audio_latent_length=3,
dtype=torch.float32,
device=torch.device("cpu"),
generator=torch.Generator().manual_seed(42),
)
assert original_num_frames == 3
assert padded_num_frames == 3
torch.testing.assert_close(actual, expected, rtol=0, atol=0)
def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=2)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
batch_size=1,
@@ -66,8 +109,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack():
padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
unpadded = latent_ops.unpad_audio_latents(padded, 10)
unpacked = latent_ops.unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
expected = latent_ops.unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
unpacked = latent_ops.unpack_audio_latents(unpadded, num_mel_bins=2)
expected = latent_ops.unpack_audio_latents(original, num_mel_bins=2)
assert unpacked.shape == (1, 2, 10, 2)
assert not (unpacked == 999.0).any()
@@ -75,7 +118,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack():
def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
pipeline.preserve_sp_padded_audio_duration = True
latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
@@ -97,7 +141,7 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
latents = torch.zeros(1, 2, 13, 4)
audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
@@ -106,7 +150,7 @@ def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
def test_prepare_audio_latents_rejects_incompatible_provided_length():
pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
latents = torch.zeros(1, 11, 4)
with pytest.raises(ValueError, match="incompatible audio frame count"):
tests/diffusion/models/ltx2/test_ltx2_parallel.py
+54 -89
View File
@@ -15,9 +15,9 @@ 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 import LTX23Pipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
pipe = object.__new__(LTX23Pipeline)
pipe = object.__new__(LTX2Pipeline)
video_sample = torch.tensor([[[1.0, -2.0]]])
audio_sample = torch.tensor([[[0.5, 3.0]]])
video_pos = torch.tensor([[[0.2, -0.3]]])
@@ -50,93 +50,53 @@ class TestCFGParallelHelpers:
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 ltx2_guidance
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 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(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
expected_video = ltx2_guidance.combine_velocity_via_x0(
video_sample,
video_pos,
video_neg,
video_sigma,
scale,
)
expected_audio = ltx2_guidance.combine_velocity_via_x0(
audio_sample,
audio_pos,
audio_neg,
audio_sigma,
scale,
def test_cfg_parallel_rejects_non_cfg_guidance(self):
from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
LTX_GUIDANCE_EXECUTOR,
LTXGuidancePlan,
)
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
for rank, expected_branch in ((0, "positive"), (1, "negative")):
calls = []
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
def fake_predict_noise(**kwargs):
calls.append(kwargs["branch"])
if kwargs["branch"] == "positive":
return video_pos, audio_pos
return video_neg, audio_neg
with pytest.raises(ValueError, match="CFG-only guidance"):
LTX_GUIDANCE_EXECUTOR.validate_cfg_world_size(plan, 2)
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,
)
def test_cfg_parallel_dummy_warms_supported_cfg_only_plan(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
assert calls == [expected_branch]
torch.testing.assert_close(video_combined, expected_video)
torch.testing.assert_close(audio_combined, expected_audio)
monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
pipe = object.__new__(LTX2Pipeline)
req = SimpleNamespace(is_dummy_run=lambda: True)
request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
assert "_cfg_video_latents" not in pipe.__dict__
assert "_cfg_audio_latents" not in pipe.__dict__
cfg_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
assert cfg_parallel_ready is True
assert pipe._guidance_plan.names == ("cond", "uncond")
assert pipe._guidance_plan.spec.video.cfg_scale == 3.0
assert pipe._guidance_plan.spec.audio.cfg_scale == 7.0
class TestCFGParallelForwardPath:
"""Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
@pytest.mark.parametrize("pipeline_cls_name", ["LTX2Pipeline", "LTX23Pipeline"])
def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch, pipeline_cls_name):
from vllm_omni.diffusion.models.ltx2 import ltx2_pipeline_runtime
def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 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
pipeline_cls = getattr(ltx23, pipeline_cls_name)
pipe = object.__new__(pipeline_cls)
pipe = object.__new__(ltx23.LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 4
monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
pipe.vae_spatial_compression_ratio = 32
pipe.vae_temporal_compression_ratio = 8
monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
class StopAtEncodePromptError(Exception):
pass
@@ -224,7 +184,7 @@ class TestCFGParallelForwardPath:
("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(
def test_forward_cfg_parallel_uses_shared_video_audio_euler_step(
self,
monkeypatch,
cfg_rank,
@@ -233,13 +193,13 @@ class TestCFGParallelForwardPath:
audio_sampling_rate,
expected_frame_rate,
):
from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_pipeline_runtime
from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_runtime
from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 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)
pipe = object.__new__(ltx23.LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.device = torch.device("cpu")
pipe.tokenizer_max_length = 1
@@ -278,7 +238,7 @@ class TestCFGParallelForwardPath:
return [audio_pos, audio_neg]
raise AssertionError(f"Unexpected gathered tensor: {tensor}")
monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank)
monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
@@ -314,7 +274,7 @@ class TestCFGParallelForwardPath:
def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
assert padding_side == "left"
assert prompt_embeds.shape[0] == 2
assert prompt_embeds.shape[0] == 1
return prompt_embeds, prompt_embeds, prompt_attention_mask
rope_video_fps: list[float] = []
@@ -380,7 +340,16 @@ class TestCFGParallelForwardPath:
num_frames=1,
frame_rate=frame_rate_input,
num_inference_steps=2,
guidance_scale=4.0,
extra_args={
"video_cfg_scale": 3.0,
"audio_cfg_scale": 7.0,
"video_stg_scale": 0.0,
"audio_stg_scale": 0.0,
"video_modality_scale": 1.0,
"audio_modality_scale": 1.0,
"video_rescale_scale": 0.0,
"audio_rescale_scale": 0.0,
},
latents=video_latents,
audio_latents=audio_latents,
output_type="latent",
@@ -388,35 +357,31 @@ class TestCFGParallelForwardPath:
request_id="ltx23-cfg-parallel-forward-test",
)
output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
output = pipe.forward(DiffusionRequestBatch(requests=[req]), sigmas=[0.25, 0.25])[0]
expected_video_noise = ltx2_guidance.combine_velocity_via_x0(
video_latents,
video_pos,
video_neg,
pipe.scheduler.sigmas[0],
4.0,
3.0,
)
expected_audio_noise = ltx2_guidance.combine_velocity_via_x0(
audio_latents,
audio_pos,
audio_neg,
pipe.scheduler.sigmas[0],
4.0,
7.0,
)
scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
assert scheduler_call_names == ["video", "audio", "video", "audio"]
assert pipe.scheduler.calls == []
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))
sigma_delta = pipe.scheduler.sigmas[-1] - pipe.scheduler.sigmas[0]
expected_video = video_latents + expected_video_noise * sigma_delta
expected_audio = audio_latents + expected_audio_noise * sigma_delta
torch.testing.assert_close(video_out, expected_video.reshape(1, 2, 1, 1, 1))
torch.testing.assert_close(audio_out, expected_audio.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.
tests/diffusion/models/ltx2/test_ltx2_pipeline.py
+1001 -203
View File
File diff suppressed because it is too large. Load diff
tests/diffusion/models/ltx2/test_ltx2_transformer.py
+16 -2
View File
@@ -10,11 +10,12 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
def test_ltx_rms_norm_no_affine_matches_official_torch_module():
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
assert "weight" not in dict(norm.named_parameters())
assert "weight" in dict(norm.named_buffers())
assert norm.weight is None
assert "weight" not in dict(norm.named_buffers())
assert "weight" not in norm.state_dict()
@@ -50,3 +51,16 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
matched.append(name)
assert matched == ["transformer_blocks.0", "transformer_blocks.1"]
@pytest.mark.parametrize("rope_type", ["interleaved", "split"])
def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""]
video_timestep = root_plan["timestep"]
audio_timestep = root_plan["audio_timestep"]
assert video_timestep.split_dim == audio_timestep.split_dim == 1
assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
assert not video_timestep.split_output
assert not audio_timestep.split_output
tests/diffusion/models/ltx2/test_ltx2_vae.py
+24 -16
View File
@@ -14,7 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
class TestLTXOutputRank:
@pytest.mark.parametrize(
("distributed_vae_state", "expected_decode_calls"),
[(False, 0), (True, 1), (RuntimeError("unavailable"), 0)],
[(False, 0), (True, 1), (RuntimeError("unavailable"), None)],
)
def test_non_output_rank_only_enters_collective_vae_decode(
self,
@@ -22,7 +22,8 @@ class TestLTXOutputRank:
distributed_vae_state,
expected_decode_calls,
):
import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx
import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
class FakeVae:
dtype = torch.float32
@@ -40,23 +41,30 @@ class TestLTXOutputRank:
self.decode_calls += 1
return (torch.ones(1, 1),)
monkeypatch.setattr(ltx.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(ltx.LTX2Pipeline)
monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
pipe = object.__new__(LTX2Pipeline)
torch.nn.Module.__init__(pipe)
pipe.vae = FakeVae()
output = pipe._decode_output(
latents=torch.ones(1, 1),
audio_latents=torch.ones(1, 1),
output_type="np",
connector_prompt_embeds=torch.ones(1, 1),
generator=None,
device=torch.device("cpu"),
decode_timestep=0.0,
decode_noise_scale=None,
prompt_batch_size=1,
)
decode_kwargs = {
"latents": torch.ones(1, 1),
"audio_latents": torch.ones(1, 1),
"output_type": "np",
"connector_prompt_embeds": torch.ones(1, 1),
"generator": None,
"device": torch.device("cpu"),
"decode_timestep": 0.0,
"decode_noise_scale": None,
"prompt_batch_size": 1,
}
if isinstance(distributed_vae_state, Exception):
with pytest.raises(RuntimeError, match="unavailable"):
pipe._decode_output(**decode_kwargs)
assert pipe.vae.decode_calls == 0
return
output = pipe._decode_output(**decode_kwargs)
assert pipe.vae.decode_calls == expected_decode_calls
assert output.output[0].numel() == 0
tests/e2e/accuracy/ltx/__init__.py
+1
View File
@@ -0,0 +1 @@
"""LTX family official-reference accuracy tests."""
tests/e2e/accuracy/ltx/run_ltx_reference.py
+331
View File
@@ -0,0 +1,331 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Generate raw LTX video and audio outputs with the official or Omni runtime."""
from __future__ import annotations
import argparse
import json
import os
import sys
from pathlib import Path
from typing import Any
import numpy as np
import torch
from PIL import Image
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser()
parser.add_argument("--backend", choices=("official", "omni"), required=True)
parser.add_argument("--request", type=Path, required=True)
parser.add_argument("--output-dir", type=Path, required=True)
parser.add_argument("--model")
parser.add_argument("--model-class-name")
parser.add_argument("--official-root", type=Path)
parser.add_argument("--checkpoint", type=Path)
parser.add_argument("--gemma-root", type=Path)
parser.add_argument("--enable-layerwise-offload", action="store_true")
return parser.parse_args()
def _save_outputs(
output_dir: Path,
*,
video: torch.Tensor,
audio: torch.Tensor,
audio_sample_rate: int,
metadata: dict[str, Any],
) -> None:
output_dir.mkdir(parents=True, exist_ok=True)
video = video.detach().float().cpu().clamp(0.0, 1.0)
audio = audio.detach().float().cpu()
np.save(output_dir / "video.npy", video.numpy())
np.save(output_dir / "audio.npy", audio.numpy())
frame_indices = sorted({0, video.shape[0] // 2, video.shape[0] - 1})
for index in frame_indices:
frame = video[index].mul(255.0).round().to(torch.uint8).numpy()
Image.fromarray(frame).save(output_dir / f"frame_{index:04d}.png")
metadata.update(
{
"video_shape": list(video.shape),
"audio_shape": list(audio.shape),
"audio_sample_rate": int(audio_sample_rate),
"frame_indices": frame_indices,
}
)
(output_dir / "metadata.json").write_text(json.dumps(metadata, indent=2) + "\n")
def _insert_official_paths(official_root: Path) -> None:
for relative_path in ("packages/ltx-core/src", "packages/ltx-pipelines/src"):
path = str((official_root / relative_path).resolve())
if path not in sys.path:
sys.path.insert(0, path)
def _configure_official_sdpa(pipeline: Any) -> None:
"""Use PyTorch SDPA for official connector and denoiser attention."""
from ltx_core.loader.attention_ops import set_attention_module_op
from ltx_core.model.transformer.attention import PytorchAttention
class AllValidSDPA(PytorchAttention):
def __call__(
self,
query: torch.Tensor,
key: torch.Tensor,
value: torch.Tensor,
heads: int,
mask: torch.Tensor | None = None,
) -> torch.Tensor:
if mask is not None and torch.count_nonzero(mask).item():
raise ValueError("The LTX accuracy guard cannot discard a non-empty attention mask")
return super().__call__(query, key, value, heads, mask=None)
attention = AllValidSDPA()
module_op = set_attention_module_op(
attention=attention,
masked_attention=attention,
)
owners_and_attributes = (
(pipeline.stage, "_transformer_builder"),
(pipeline.prompt_encoder, "_embeddings_processor_builder"),
)
for owner, attribute in owners_and_attributes:
builder = getattr(owner, attribute)
setattr(
owner,
attribute,
builder.with_module_ops((*builder.module_ops, module_op)),
)
@torch.inference_mode()
def _run_official(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.official_root is None or args.checkpoint is None or args.gemma_root is None:
raise ValueError("Official backend requires --official-root, --checkpoint, and --gemma-root")
_insert_official_paths(args.official_root)
# vLLM disables cuDNN SDPA during import; mirror the worker's Gemma dispatch.
torch.backends.cuda.enable_cudnn_sdp(False)
from ltx_core.components.guiders import MultiModalGuiderParams
from ltx_pipelines.ti2vid_one_stage import TI2VidOneStagePipeline
from ltx_pipelines.utils.args import ImageConditioningInput
from ltx_pipelines.utils.types import OffloadMode
pipeline = TI2VidOneStagePipeline(
checkpoint_path=str(args.checkpoint),
gemma_root=str(args.gemma_root),
loras=(),
offload_mode=OffloadMode.CPU if args.enable_layerwise_offload else OffloadMode.NONE,
)
_configure_official_sdpa(pipeline)
image_path = request.get("image")
images = (
[]
if image_path is None
else [
ImageConditioningInput(
path=str(image_path),
frame_idx=0,
strength=1.0,
# Both runtimes must receive the same source pixels. The
# official CLI's optional H.264 preprocessing is not part of
# the seeded model trajectory under test.
crf=0,
)
]
)
video, audio = pipeline(
prompt=request["prompt"],
negative_prompt=request["negative_prompt"],
seed=request["seed"],
height=request["height"],
width=request["width"],
num_frames=request["num_frames"],
frame_rate=request["fps"],
num_inference_steps=request["num_inference_steps"],
video_guider_params=MultiModalGuiderParams(
cfg_scale=request["video_cfg_scale"],
stg_scale=request["video_stg_scale"],
rescale_scale=request["video_rescale_scale"],
modality_scale=request["video_modality_scale"],
skip_step=0,
stg_blocks=request["video_stg_blocks"],
),
audio_guider_params=MultiModalGuiderParams(
cfg_scale=request["audio_cfg_scale"],
stg_scale=request["audio_stg_scale"],
rescale_scale=request["audio_rescale_scale"],
modality_scale=request["audio_modality_scale"],
skip_step=0,
stg_blocks=request["audio_stg_blocks"],
),
images=images,
max_batch_size=4,
)
video_tensor = torch.cat([chunk.detach().cpu() for chunk in video], dim=0)
_save_outputs(
args.output_dir,
video=_canonical_video(video_tensor),
audio=audio.waveform,
audio_sample_rate=audio.sampling_rate,
metadata={
"backend": "official",
"attention_backend": "torch_sdpa",
"official_revision": os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION"),
"checkpoint": str(args.checkpoint),
},
)
def _unwrap_omni_output(output: Any) -> tuple[Any, Any, int]:
from vllm_omni.outputs import OmniRequestOutput
audio = None
audio_sample_rate = None
frames = output[0] if isinstance(output, list) and output else output
if isinstance(frames, OmniRequestOutput):
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio")
audio_sample_rate = multimodal_output.get("audio_sample_rate")
if frames.is_pipeline_output and isinstance(frames.request_output, OmniRequestOutput):
frames = frames.request_output
multimodal_output = frames.multimodal_output or {}
audio = multimodal_output.get("audio", audio)
audio_sample_rate = multimodal_output.get("audio_sample_rate", audio_sample_rate)
if isinstance(frames, OmniRequestOutput):
if not frames.images:
raise ValueError("No video frames found in OmniRequestOutput")
frames = frames.images
if isinstance(frames, list) and len(frames) == 1:
frames = frames[0]
if isinstance(frames, tuple) and len(frames) == 2:
frames, audio = frames
elif isinstance(frames, dict):
audio = frames.get("audio", audio)
audio_sample_rate = frames.get("audio_sample_rate", audio_sample_rate)
frames = frames.get("frames", frames.get("video"))
if frames is None or audio is None or audio_sample_rate is None:
raise ValueError("Omni output did not contain video, audio, and audio_sample_rate")
return frames, audio, int(audio_sample_rate)
def _canonical_video(video: Any) -> torch.Tensor:
if isinstance(video, list):
if len(video) == 1:
return _canonical_video(video[0])
tensors = [torch.as_tensor(np.asarray(item) if not isinstance(item, torch.Tensor) else item) for item in video]
if all(item.ndim == 3 for item in tensors):
video = torch.stack(tensors)
elif all(item.ndim == 4 for item in tensors):
video = torch.cat(tensors)
else:
raise ValueError(f"Cannot combine video items with dimensions {[item.ndim for item in tensors]}")
tensor = torch.as_tensor(np.asarray(video) if not isinstance(video, torch.Tensor) else video).detach().cpu()
if tensor.ndim == 5:
tensor = tensor[0]
if tensor.ndim != 4:
raise ValueError(f"Expected a 4D video tensor, got {tuple(tensor.shape)}")
if tensor.shape[-1] in (3, 4):
tensor = tensor[..., :3]
elif tensor.shape[1] in (3, 4):
tensor = tensor[:, :3].permute(0, 2, 3, 1)
elif tensor.shape[0] in (3, 4):
tensor = tensor[:3].permute(1, 2, 3, 0)
else:
raise ValueError(f"Cannot infer video channel dimension from {tuple(tensor.shape)}")
tensor = tensor.float()
if tensor.numel() and tensor.max() > 1:
tensor = tensor / 255.0
elif tensor.numel() and tensor.min() < 0:
tensor = tensor.clamp(-1.0, 1.0).add(1.0).mul(0.5)
return tensor.clamp(0.0, 1.0)
@torch.inference_mode()
def _run_omni(args: argparse.Namespace, request: dict[str, Any]) -> None:
if args.model is None or args.model_class_name is None:
raise ValueError("Omni backend requires --model and --model-class-name")
attention_config = {"default": {"backend": "TORCH_SDPA"}}
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
from vllm_omni.platforms import current_omni_platform
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(request["seed"])
omni = Omni(
model=args.model,
model_class_name=args.model_class_name,
enforce_eager=True,
enable_layerwise_offload=args.enable_layerwise_offload,
diffusion_attention_config=attention_config,
parallel_config=DiffusionParallelConfig(),
)
try:
model_class_name = get_model_class_name(omni)
sampling_params = OmniDiffusionSamplingParams(
height=request["height"],
width=request["width"],
generator=generator,
guidance_scale=None,
num_inference_steps=request["num_inference_steps"],
num_frames=request["num_frames"],
fps=request["fps"],
frame_rate=float(request["fps"]),
output_type="np",
)
guidance = {
key: value for key, value in request.items() if key.startswith("video_") or key.startswith("audio_")
}
apply_declared_extra_args(sampling_params, get_extra_body_params(model_class_name), guidance)
prompt: dict[str, Any] = {
"prompt": request["prompt"],
"negative_prompt": request["negative_prompt"],
}
image_path = request.get("image")
if image_path is not None:
with Image.open(str(image_path)) as source_image:
image = source_image.convert("RGB")
image.load()
prompt["multi_modal_data"] = {"image": image}
output = omni.generate(prompt, sampling_params)
video, audio, audio_sample_rate = _unwrap_omni_output(output)
audio_tensor = torch.as_tensor(np.asarray(audio) if not isinstance(audio, torch.Tensor) else audio)
_save_outputs(
args.output_dir,
video=_canonical_video(video),
audio=audio_tensor,
audio_sample_rate=audio_sample_rate,
metadata={
"backend": "omni",
"attention_backend": "torch_sdpa",
"model": args.model,
"model_class_name": model_class_name,
},
)
finally:
omni.shutdown()
def main() -> None:
args = _parse_args()
request = json.loads(args.request.read_text())
if args.backend == "official":
_run_official(args, request)
else:
_run_omni(args, request)
if __name__ == "__main__":
main()
tests/e2e/accuracy/ltx/test_ltx_official_similarity.py
+440
View File
@@ -0,0 +1,440 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E accuracy guard against a pinned Lightricks LTX pipeline revision.
The comparison runs both runtimes through PyTorch SDPA and uses
``max_batch_size=4`` in the official reference to match Omni's fused guidance
batch. Video and audio guidance use the official non-HQ one-stage defaults;
only the generation shape and step count are reduced for CI runtime.
"""
from __future__ import annotations
import json
import os
import shutil
import subprocess
import sys
import time
from dataclasses import dataclass
from pathlib import Path
import numpy as np
import pytest
import torch
from huggingface_hub import hf_hub_download, snapshot_download
from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure
from tests.e2e.accuracy.helpers import reset_artifact_dir
from tests.helpers.mark import hardware_test
OFFICIAL_REPOSITORY = "https://github.com/Lightricks/LTX-2.git"
OFFICIAL_REVISION = "9377758131b1ffde4b7f766804590a6617bf2ab9"
# Version selected by this revision's uv.lock. Keep it out of Omni's runtime and dev dependencies.
OFFICIAL_OPENIMAGEIO_VERSION = "3.1.11.0"
PROMPT = (
"A space shuttle launches vertically above a desert launch pad. Bright exhaust flames and a dense white "
"plume billow beneath it while the camera remains fixed."
)
NEGATIVE_PROMPT = (
"blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
"grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
"deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
"wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of field, "
"background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent lighting "
"direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny valley effect, "
"incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, mismatched lip sync, "
"silent or muted audio, distorted voice, robotic voice, echo, background noise, off-sync audio, incorrect "
"dialogue, added dialogue, repetitive speech, jittery movement, awkward pauses, incorrect timing, unnatural "
"transitions, inconsistent framing, tilted camera, flat lighting, inconsistent tone, cinematic oversaturation, "
"stylized filters, or AI artifacts."
)
# Both runtimes use PyTorch SDPA with the current Torch dispatch defaults.
ATTENTION_BACKEND = "torch_sdpa"
VIDEO_SSIM_MEAN_THRESHOLD = 0.95
VIDEO_SSIM_MIN_THRESHOLD = 0.90
VIDEO_PSNR_MEAN_THRESHOLD = 30.0
AUDIO_RELATIVE_L2_THRESHOLD = 0.2
AUDIO_COSINE_THRESHOLD = 0.95
@dataclass(frozen=True)
class LTXAccuracyCase:
name: str
model_id: str
model_revision: str
model_env: str
model_class_name: str
checkpoint_repo: str
checkpoint_filename: str
checkpoint_revision: str
checkpoint_env: str
stg_block: int
prompt: str = PROMPT
image_repo: str | None = None
image_filename: str | None = None
image_revision: str | None = None
CASES = (
LTXAccuracyCase(
name="ltx2",
model_id="Lightricks/LTX-2",
model_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
model_env="VLLM_TEST_LTX2_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2",
checkpoint_filename="ltx-2-19b-dev.safetensors",
checkpoint_revision="47da56e2ad66ce4125a9922b4a8826bf407f9d0a",
checkpoint_env="VLLM_TEST_LTX2_OFFICIAL_CHECKPOINT",
stg_block=29,
),
LTXAccuracyCase(
name="ltx2_3",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
stg_block=28,
),
LTXAccuracyCase(
name="ltx2_3_i2v",
model_id="diffusers/LTX-2.3-Diffusers",
model_revision="8eee8edcf067e838b843f926ec4d4cc9b2be1aaf",
model_env="VLLM_TEST_LTX23_MODEL",
model_class_name="LTX2Pipeline",
checkpoint_repo="Lightricks/LTX-2.3",
checkpoint_filename="ltx-2.3-22b-dev.safetensors",
checkpoint_revision="4229404625088d21c4f112eb640fb04a0900ee25",
checkpoint_env="VLLM_TEST_LTX23_OFFICIAL_CHECKPOINT",
stg_block=28,
image_repo="huggingface/documentation-images",
image_filename="diffusers/svd/rocket.png",
image_revision="645d8364f0c7a101180b364811b5a11a362e4010",
),
)
def _run(command: list[str], *, env: dict[str, str], timeout: int = 1800) -> None:
start = time.perf_counter()
subprocess.run(command, env=env, timeout=timeout, check=True)
print(f"{' '.join(command[:3])} finished in {time.perf_counter() - start:.1f}s")
def _clone_official_source(root: Path, revision: str) -> None:
root.mkdir(parents=True)
repository = os.environ.get("VLLM_TEST_LTX_OFFICIAL_REPOSITORY", OFFICIAL_REPOSITORY)
subprocess.run(["git", "init", "-q", str(root)], check=True)
subprocess.run(["git", "-C", str(root), "remote", "add", "origin", repository], check=True)
last_error: subprocess.CalledProcessError | None = None
for attempt in range(3):
try:
subprocess.run(
["git", "-C", str(root), "fetch", "--depth", "1", "origin", revision],
check=True,
)
last_error = None
break
except subprocess.CalledProcessError as error:
last_error = error
if attempt < 2:
time.sleep(5 * (attempt + 1))
if last_error is not None:
raise last_error
subprocess.run(["git", "-C", str(root), "checkout", "-q", "--detach", "FETCH_HEAD"], check=True)
def _git_revision(root: Path) -> str | None:
try:
return subprocess.check_output(
["git", "-C", str(root), "rev-parse", "HEAD"],
stderr=subprocess.DEVNULL,
text=True,
).strip()
except subprocess.CalledProcessError:
return None
def _official_source(artifact_root: Path) -> tuple[Path, str]:
revision = os.environ.get("VLLM_TEST_LTX_OFFICIAL_REVISION", OFFICIAL_REVISION)
configured_root = os.environ.get("VLLM_TEST_LTX_OFFICIAL_ROOT")
root = Path(configured_root) if configured_root else artifact_root / f"official-source-{revision[:12]}"
actual_revision = _git_revision(root) if root.exists() else None
if actual_revision != revision and configured_root:
raise AssertionError(f"Official source revision mismatch: {actual_revision} != {revision}")
if actual_revision != revision:
if root.exists():
shutil.rmtree(root)
_clone_official_source(root, revision)
actual_revision = _git_revision(root)
assert actual_revision == revision, f"Official source revision mismatch: {actual_revision} != {revision}"
return root, revision
def _official_runner_prefix() -> list[str]:
"""Run the reference with its missing binary dependency isolated from CI."""
uv = shutil.which("uv")
assert uv is not None, "uv is required to run the pinned official LTX reference"
return [
uv,
"run",
"--no-project",
"--with",
f"openimageio=={OFFICIAL_OPENIMAGEIO_VERSION}",
"--python",
sys.executable,
"python",
]
def _resolve_model(case: LTXAccuracyCase) -> Path:
configured_model = os.environ.get(case.model_env)
if configured_model and Path(configured_model).exists():
return Path(configured_model)
model_id = configured_model or case.model_id
revision = os.environ.get(f"{case.model_env}_REVISION")
if revision is None and model_id == case.model_id:
revision = case.model_revision
return Path(
snapshot_download(
repo_id=model_id,
revision=revision,
allow_patterns=[
"model_index.json",
"audio_vae/*",
"connectors/*",
"processor/*",
"scheduler/*",
"text_encoder/config.json",
"text_encoder/generation_config.json",
"text_encoder/model*",
"tokenizer/*",
"transformer/*",
"vae/*",
"vocoder/*",
],
)
)
def _resolve_gemma_root(model: Path) -> Path:
configured_root = os.environ.get("VLLM_TEST_LTX_GEMMA_ROOT")
if configured_root:
root = Path(configured_root)
assert root.is_dir(), f"Gemma root not found: {root}"
return root
return model
def _resolve_checkpoint(case: LTXAccuracyCase, model: Path) -> Path:
configured_checkpoint = os.environ.get(case.checkpoint_env)
if configured_checkpoint:
checkpoint = Path(configured_checkpoint)
assert checkpoint.is_file(), f"Official checkpoint not found: {checkpoint}"
return checkpoint
model_checkpoint = model / case.checkpoint_filename
if model_checkpoint.is_file():
return model_checkpoint
return Path(
hf_hub_download(
repo_id=case.checkpoint_repo,
filename=case.checkpoint_filename,
revision=case.checkpoint_revision,
)
)
def _resolve_image(case: LTXAccuracyCase) -> Path | None:
if case.image_filename is None:
return None
if case.image_repo is None or case.image_revision is None:
raise ValueError(f"Incomplete image source for LTX accuracy case {case.name!r}.")
configured_image = os.environ.get("VLLM_TEST_LTX_I2V_IMAGE")
if configured_image:
image = Path(configured_image)
assert image.is_file(), f"LTX I2V conditioning image not found: {image}"
return image
return Path(
hf_hub_download(
repo_id=case.image_repo,
repo_type="dataset",
filename=case.image_filename,
revision=case.image_revision,
)
)
def _request(case: LTXAccuracyCase, image: Path | None) -> dict[str, object]:
request: dict[str, object] = {
"prompt": case.prompt,
"negative_prompt": NEGATIVE_PROMPT,
"width": 512,
"height": 384,
"num_frames": 25,
"fps": 24,
"num_inference_steps": 20,
"seed": 42,
"video_cfg_scale": 3.0,
"audio_cfg_scale": 7.0,
"video_stg_scale": 1.0,
"audio_stg_scale": 1.0,
"video_modality_scale": 3.0,
"audio_modality_scale": 3.0,
"video_rescale_scale": 0.7,
"audio_rescale_scale": 0.7,
"video_stg_blocks": [case.stg_block],
"audio_stg_blocks": [case.stg_block],
}
if image is not None:
request["image"] = str(image.resolve())
return request
def _video_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float]:
assert reference.shape == prediction.shape
assert reference.ndim == 4 and reference.shape[-1] == 3
ssim_scores: list[float] = []
psnr_scores: list[float] = []
for reference_frame, prediction_frame in zip(reference, prediction, strict=True):
reference_tensor = torch.from_numpy(reference_frame).permute(2, 0, 1).unsqueeze(0)
prediction_tensor = torch.from_numpy(prediction_frame).permute(2, 0, 1).unsqueeze(0)
ssim_scores.append(float(StructuralSimilarityIndexMeasure(data_range=1.0)(prediction_tensor, reference_tensor)))
psnr_scores.append(float(PeakSignalNoiseRatio(data_range=1.0)(prediction_tensor, reference_tensor)))
difference = np.abs(reference.astype(np.float64) - prediction.astype(np.float64))
return {
"ssim_mean": float(np.mean(ssim_scores)),
"ssim_min": float(np.min(ssim_scores)),
"psnr_mean_db": float(np.mean(psnr_scores)),
"max_abs": float(difference.max()),
"mean_abs": float(difference.mean()),
}
def _canonical_audio(audio: np.ndarray) -> np.ndarray:
while audio.ndim > 2 and audio.shape[0] == 1:
audio = audio[0]
return audio.astype(np.float64)
def _audio_metrics(reference: np.ndarray, prediction: np.ndarray) -> dict[str, float | bool]:
reference = _canonical_audio(reference)
prediction = _canonical_audio(prediction)
assert reference.shape == prediction.shape
difference = reference - prediction
reference_norm = max(float(np.linalg.norm(reference)), 1e-12)
prediction_norm = max(float(np.linalg.norm(prediction)), 1e-12)
return {
"bitwise_equal": bool(np.array_equal(reference, prediction)),
"max_abs": float(np.abs(difference).max()),
"mean_abs": float(np.abs(difference).mean()),
"relative_l2": float(np.linalg.norm(difference) / reference_norm),
"cosine_similarity": float(np.vdot(reference.ravel(), prediction.ravel()) / (reference_norm * prediction_norm)),
}
@pytest.mark.parametrize("case", CASES, ids=lambda case: case.name)
@pytest.mark.full_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx_one_stage_matches_official(case: LTXAccuracyCase, accuracy_artifact_root: Path) -> None:
"""Compare official and Omni raw AV outputs from the same E2E request."""
output_root = reset_artifact_dir(accuracy_artifact_root / "ltx_official" / case.name)
official_root, official_revision = _official_source(accuracy_artifact_root / "ltx_official")
model = _resolve_model(case)
gemma_root = _resolve_gemma_root(model)
checkpoint = _resolve_checkpoint(case, model)
image = _resolve_image(case)
request_path = output_root / "request.json"
request_path.write_text(json.dumps(_request(case, image), indent=2) + "\n")
runner = Path(__file__).with_name("run_ltx_reference.py")
runner_args = [
str(runner),
"--request",
str(request_path),
]
if os.environ.get("VLLM_TEST_LTX_ENABLE_LAYERWISE_OFFLOAD", "").lower() in {"1", "true", "yes", "on"}:
runner_args.append("--enable-layerwise-offload")
env = os.environ.copy()
env["VLLM_TEST_LTX_OFFICIAL_REVISION"] = official_revision
env["PYTHONUNBUFFERED"] = "1"
repository_root = Path(__file__).resolve().parents[4]
existing_pythonpath = env.get("PYTHONPATH")
env["PYTHONPATH"] = (
str(repository_root) if not existing_pythonpath else f"{repository_root}{os.pathsep}{existing_pythonpath}"
)
official_output = output_root / "official"
_run(
_official_runner_prefix()
+ runner_args
+ [
"--backend",
"official",
"--output-dir",
str(official_output),
"--official-root",
str(official_root),
"--checkpoint",
str(checkpoint),
"--gemma-root",
str(gemma_root),
],
env=env,
)
omni_output = output_root / "omni"
_run(
[sys.executable]
+ runner_args
+ [
"--backend",
"omni",
"--output-dir",
str(omni_output),
"--model",
str(model),
"--model-class-name",
case.model_class_name,
],
env=env,
)
official_metadata = json.loads((official_output / "metadata.json").read_text())
omni_metadata = json.loads((omni_output / "metadata.json").read_text())
assert official_metadata["attention_backend"] == ATTENTION_BACKEND
assert omni_metadata["attention_backend"] == ATTENTION_BACKEND
assert official_metadata["audio_sample_rate"] == omni_metadata["audio_sample_rate"]
video_metrics = _video_metrics(
np.load(official_output / "video.npy"),
np.load(omni_output / "video.npy"),
)
audio_metrics = _audio_metrics(
np.load(official_output / "audio.npy"),
np.load(omni_output / "audio.npy"),
)
result = {
"case": case.name,
"task": "i2v" if image is not None else "t2v",
"attention_backend": ATTENTION_BACKEND,
"official_revision": official_revision,
"model_revision": case.model_revision,
"checkpoint_revision": case.checkpoint_revision,
"video": video_metrics,
"audio": audio_metrics,
}
(output_root / "metrics.json").write_text(json.dumps(result, indent=2) + "\n")
print(json.dumps(result, indent=2))
assert video_metrics["ssim_mean"] >= VIDEO_SSIM_MEAN_THRESHOLD
assert video_metrics["ssim_min"] >= VIDEO_SSIM_MIN_THRESHOLD
assert video_metrics["psnr_mean_db"] >= VIDEO_PSNR_MEAN_THRESHOLD
assert audio_metrics["relative_l2"] <= AUDIO_RELATIVE_L2_THRESHOLD
assert audio_metrics["cosine_similarity"] >= AUDIO_COSINE_THRESHOLD
tests/e2e/accuracy/test_ltx2_3_video_similarity.py
-413
View File
@@ -1,413 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
SSIM/PSNR accuracy tests for LTX-2.3.
1. **Transformer parity** (``test_ltx2_3_transformer_matches_diffusers``):
Swaps our custom transformer into diffusers' ``LTX2Pipeline`` to measure
numerical parity in isolation. Thresholds: SSIM >= 0.95, PSNR >= 28 dB.
Result: SSIM 0.999987 (bit-identical).
2. **Full pipeline** (``test_ltx2_3_pipeline_matches_diffusers``):
Runs the full vLLM-Omni serving stack (``OmniServer`` -> HTTP API) and
compares per-frame against stock diffusers. Currently skipped because
the OmniServer subprocess creates a different RNG state than in-process
diffusers, producing different initial latents from the same seed.
This is a test infrastructure limitation, not a model accuracy issue.
"""
from __future__ import annotations
import gc
import os
import tempfile
from pathlib import Path
import diffusers
import numpy as np
import pytest
import requests
import torch
from PIL import Image
from tests.e2e.accuracy.helpers import compute_image_ssim_psnr, model_output_dir
from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
from tests.helpers.mark import hardware_test
from tests.helpers.runtime import OmniServer
# Parse diffusers version for compatibility check
_DIFFUSERS_VERSION = tuple(int(x) for x in diffusers.__version__.split(".")[:2] if x.isdigit())
_DIFFUSERS_038 = _DIFFUSERS_VERSION >= (0, 38)
MODEL_ID = "diffusers/LTX-2.3-Diffusers"
MODEL_ENV_VAR = "VLLM_TEST_LTX23_MODEL"
PROMPT = "A lighthouse on a rocky cliff at sunset, waves crashing below, golden hour lighting"
NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
WIDTH = 512
HEIGHT = 384
NUM_FRAMES = 25 # ~1 second at 24fps
NUM_INFERENCE_STEPS = 20
GUIDANCE_SCALE = 4.0
SEED = 42
# Transformer-swap test: near-identical output expected
TRANSFORMER_SSIM_THRESHOLD = 0.95
TRANSFORMER_PSNR_THRESHOLD = 28.0
# Full-pipeline test: allows minor divergence from RNG / pipeline differences
PIPELINE_SSIM_THRESHOLD = 0.94
PIPELINE_PSNR_THRESHOLD = 28.0
def _model_name() -> str:
return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
def _local_files_only(model: str) -> bool:
return Path(model).exists()
# ---------------------------------------------------------------------------
# Frame extraction helpers
# ---------------------------------------------------------------------------
def _video_to_frames(video_np: np.ndarray) -> list[Image.Image]:
"""Convert numpy video to list of PIL Images."""
while video_np.ndim > 4:
video_np = video_np[0]
if video_np.dtype in (np.float32, np.float64, np.float16):
video_np = np.clip(video_np * 255, 0, 255).astype(np.uint8)
return [Image.fromarray(video_np[t]) for t in range(video_np.shape[0])]
def _extract_diffusers_frames(result) -> list[Image.Image]:
"""Extract frames from diffusers pipeline output."""
video = result.frames
if isinstance(video, np.ndarray):
return _video_to_frames(video)
if isinstance(video, list):
if isinstance(video[0], list):
return [img.convert("RGB") for img in video[0]]
if isinstance(video[0], Image.Image):
return [img.convert("RGB") for img in video]
raise ValueError(f"Unexpected output type: {type(video)}")
def _extract_mp4_frames(mp4_bytes: bytes) -> list[Image.Image]:
"""Extract frames from an MP4 video using ffmpeg."""
import subprocess
with tempfile.TemporaryDirectory() as tmpdir:
mp4_path = os.path.join(tmpdir, "video.mp4")
with open(mp4_path, "wb") as f:
f.write(mp4_bytes)
# Extract video frames as PNG files using ffmpeg
frame_pattern = os.path.join(tmpdir, "frame_%04d.png")
subprocess.run(
["ffmpeg", "-i", mp4_path, "-vsync", "0", frame_pattern],
capture_output=True,
check=True,
)
# Load frames in order
frames = []
i = 1
while True:
fpath = os.path.join(tmpdir, f"frame_{i:04d}.png")
if not os.path.exists(fpath):
break
frames.append(Image.open(fpath).convert("RGB").copy())
i += 1
return frames
# ---------------------------------------------------------------------------
# Comparison helper
# ---------------------------------------------------------------------------
def _assert_video_similarity(
*,
model_name: str,
vllm_frames: list[Image.Image],
diffusers_frames: list[Image.Image],
ssim_threshold: float,
psnr_threshold: float,
) -> tuple[float, float]:
"""Compare video frames and assert SSIM/PSNR meet thresholds."""
min_frames = min(len(vllm_frames), len(diffusers_frames))
assert min_frames > 0, "No frames to compare"
ssim_scores = []
psnr_scores = []
for i in range(min_frames):
ssim_val, psnr_val = compute_image_ssim_psnr(
prediction=vllm_frames[i],
reference=diffusers_frames[i],
)
ssim_scores.append(ssim_val)
psnr_scores.append(psnr_val)
avg_ssim = sum(ssim_scores) / len(ssim_scores)
avg_psnr = sum(psnr_scores) / len(psnr_scores)
print(f"\n{model_name} video similarity ({min_frames} frames):")
print(f" SSIM: avg={avg_ssim:.6f}, min={min(ssim_scores):.6f}, threshold>={ssim_threshold:.6f}")
print(f" PSNR: avg={avg_psnr:.6f} dB, min={min(psnr_scores):.6f} dB, threshold>={psnr_threshold:.6f} dB")
assert avg_ssim >= ssim_threshold, f"SSIM below threshold: got {avg_ssim:.6f}, expected >= {ssim_threshold:.6f}."
assert avg_psnr >= psnr_threshold, f"PSNR below threshold: got {avg_psnr:.6f}, expected >= {psnr_threshold:.6f}."
return avg_ssim, avg_psnr
# ---------------------------------------------------------------------------
# Diffusers baseline (shared by both tests)
# ---------------------------------------------------------------------------
def _run_diffusers_baseline(model: str, output_dir: Path) -> list[Image.Image]:
"""Generate video using stock diffusers LTX2Pipeline."""
from diffusers import LTX2Pipeline
run_pre_test_cleanup()
pipe = None
try:
pipe = LTX2Pipeline.from_pretrained(
model, torch_dtype=torch.bfloat16, local_files_only=_local_files_only(model)
).to("cuda")
generator = torch.Generator(device="cuda").manual_seed(SEED)
result = pipe(
prompt=PROMPT,
negative_prompt=NEGATIVE_PROMPT,
width=WIDTH,
height=HEIGHT,
num_frames=NUM_FRAMES,
num_inference_steps=NUM_INFERENCE_STEPS,
guidance_scale=GUIDANCE_SCALE,
generator=generator,
output_type="np",
)
frames = _extract_diffusers_frames(result)
for i, f in enumerate(frames):
f.save(output_dir / f"diffusers_frame_{i:04d}.png")
return frames
finally:
del pipe
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
run_post_test_cleanup()
# ---------------------------------------------------------------------------
# Test 1: Transformer-swap parity
# ---------------------------------------------------------------------------
def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Image]:
"""Run diffusers pipeline with our custom transformer swapped in."""
from contextlib import nullcontext
from diffusers import LTX2Pipeline
from vllm.config import VllmConfig, set_current_vllm_config
from vllm.distributed.parallel_state import init_distributed_environment, initialize_model_parallel
from vllm_omni.diffusion.models.ltx2 import (
create_transformer_from_config,
load_transformer_config,
)
vllm_config = VllmConfig()
ctx = set_current_vllm_config(vllm_config)
ctx.__enter__()
if not torch.distributed.is_initialized():
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29503")
os.environ.setdefault("RANK", "0")
os.environ.setdefault("WORLD_SIZE", "1")
init_distributed_environment(world_size=1, rank=0, local_rank=0)
initialize_model_parallel(tensor_model_parallel_size=1)
local = _local_files_only(model)
pipe = LTX2Pipeline.from_pretrained(model, torch_dtype=torch.bfloat16, local_files_only=local)
transformer_config = load_transformer_config(model, "transformer", local)
our_transformer = create_transformer_from_config(transformer_config)
diffusers_state = dict(pipe.transformer.named_parameters())
def _weight_iter():
for name, param in diffusers_state.items():
yield name, param.data
our_transformer.load_weights(_weight_iter())
our_transformer = our_transformer.to(dtype=torch.bfloat16, device="cuda").eval()
# Compatibility shims for diffusers pipeline
our_transformer.dtype = torch.bfloat16
if not hasattr(our_transformer, "cache_context"):
our_transformer.cache_context = lambda name: nullcontext()
del pipe.transformer
pipe.transformer = our_transformer
for name, component in pipe.components.items():
if name != "transformer" and hasattr(component, "to"):
try:
component.to("cuda")
except Exception:
pass
generator = torch.Generator(device="cuda").manual_seed(SEED)
result = pipe(
prompt=PROMPT,
negative_prompt=NEGATIVE_PROMPT,
width=WIDTH,
height=HEIGHT,
num_frames=NUM_FRAMES,
num_inference_steps=NUM_INFERENCE_STEPS,
guidance_scale=GUIDANCE_SCALE,
generator=generator,
output_type="np",
)
frames = _extract_diffusers_frames(result)
for i, f in enumerate(frames):
f.save(output_dir / f"vllm_transformer_frame_{i:04d}.png")
del pipe, result, our_transformer
gc.collect()
if torch.cuda.is_available():
torch.accelerator.empty_cache()
return frames
@pytest.mark.full_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@pytest.mark.skipif(
not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
)
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx2_3_transformer_matches_diffusers(accuracy_artifact_root: Path) -> None:
"""Transformer-level parity: swap our transformer into diffusers pipeline.
Isolates transformer numerical accuracy from pipeline-level differences.
Both runs use diffusers' denoising loop, CFG, scheduler, and RNG.
"""
model = _model_name()
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
vllm_frames = _run_with_custom_transformer(model=model, output_dir=output_dir)
_assert_video_similarity(
model_name=f"{MODEL_ID} (transformer-swap)",
vllm_frames=vllm_frames,
diffusers_frames=diffusers_frames,
ssim_threshold=TRANSFORMER_SSIM_THRESHOLD,
psnr_threshold=TRANSFORMER_PSNR_THRESHOLD,
)
# ---------------------------------------------------------------------------
# Test 2: Full pipeline (OmniServer → HTTP API vs diffusers)
# ---------------------------------------------------------------------------
def _run_vllm_omni_serving(model: str, output_dir: Path) -> list[Image.Image]:
"""Generate video via the full vLLM-Omni serving stack."""
server_args = [
"--model-class-name",
"LTX23Pipeline",
"--stage-init-timeout",
"600",
]
with OmniServer(model, server_args, use_omni=True) as server:
# Submit generation request
response = requests.post(
f"http://{server.host}:{server.port}/v1/videos",
files={
"prompt": (None, PROMPT),
"negative_prompt": (None, NEGATIVE_PROMPT),
"model": (None, server.model),
"num_frames": (None, str(NUM_FRAMES)),
"fps": (None, "24"),
"size": (None, f"{WIDTH}x{HEIGHT}"),
"num_inference_steps": (None, str(NUM_INFERENCE_STEPS)),
"guidance_scale": (None, str(GUIDANCE_SCALE)),
"seed": (None, str(SEED)),
},
timeout=120,
)
response.raise_for_status()
video_id = response.json()["id"]
# Poll for completion
import time
for _ in range(120):
status_resp = requests.get(
f"http://{server.host}:{server.port}/v1/videos/{video_id}",
timeout=30,
)
status_resp.raise_for_status()
status = status_resp.json()["status"]
if status == "completed":
break
if status in ("error", "failed"):
raise RuntimeError(f"Video generation failed: {status_resp.json()}")
time.sleep(5)
else:
raise TimeoutError(f"Video generation timed out after 600s (id={video_id})")
# Download video content
content_resp = requests.get(
f"http://{server.host}:{server.port}/v1/videos/{video_id}/content",
timeout=120,
)
content_resp.raise_for_status()
mp4_bytes = content_resp.content
# Save MP4
mp4_path = output_dir / "vllm_omni_pipeline.mp4"
with open(mp4_path, "wb") as f:
f.write(mp4_bytes)
# Extract frames
frames = _extract_mp4_frames(mp4_bytes)
for i, frame in enumerate(frames):
frame.save(output_dir / f"vllm_pipeline_frame_{i:04d}.png")
return frames
@pytest.mark.full_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@pytest.mark.skipif(
not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
)
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_ltx2_3_pipeline_matches_diffusers(accuracy_artifact_root: Path) -> None:
"""Full-pipeline parity: vLLM-Omni serving stack vs diffusers.
Runs the complete vLLM-Omni OmniServer (subprocess, HTTP API, video
encoding) and compares per-frame against stock diffusers output.
Follows the Wan2.2 / Qwen Image pattern with seed-based determinism.
"""
model = _model_name()
output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
vllm_frames = _run_vllm_omni_serving(model=model, output_dir=output_dir)
_assert_video_similarity(
model_name=f"{MODEL_ID} (full-pipeline)",
vllm_frames=vllm_frames,
diffusers_frames=diffusers_frames,
ssim_threshold=PIPELINE_SSIM_THRESHOLD,
psnr_threshold=PIPELINE_PSNR_THRESHOLD,
)
tests/e2e/online_serving/test_ltx2_expansion.py
+12 -16
View File
@@ -4,8 +4,7 @@
Expansion tests for LTX-2 two-stage pipelines.
Coverage:
- HSDP with LTX2TwoStagesPipeline (text-to-video)
- HSDP with LTX2ImageToVideoTwoStagesPipeline (image-to-video)
- HSDP with LTX2DistilledPipeline (text-to-video and image-to-video)
"""
import os
@@ -27,7 +26,6 @@ HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
def _cases():
cases = []
# T2V: LTX2TwoStagesPipeline
cases.append(
pytest.param(
OmniServerParams(
@@ -35,15 +33,15 @@ def _cases():
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2TwoStagesPipeline",
"LTX2DistilledPipeline",
],
),
False,
id="t2v_hsdp",
marks=PARALLEL_MARKS,
)
)
# I2V: LTX2ImageToVideoTwoStagesPipeline
cases.append(
pytest.param(
OmniServerParams(
@@ -51,9 +49,10 @@ def _cases():
server_args=[
*HSDP_ARGS,
"--model-class-name",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2DistilledPipeline",
],
),
True,
id="i2v_hsdp",
marks=PARALLEL_MARKS,
)
@@ -62,25 +61,22 @@ def _cases():
return cases
@pytest.mark.parametrize("omni_server", _cases(), indirect=True)
@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"])
def test_ltx2_two_stage_hsdp(
omni_server: OmniServer,
is_i2v: bool,
openai_client: OpenAIClientHandler,
):
is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args)
# The two-stage pipeline generates at the requested resolution then 2x
# upsamples via LTX2LatentUpsamplerModel, so output dimensions differ
# from the request. Omit height/width from form_data so the assertion
# helper skips the dimension check (it only asserts when the key is
# present). The pipeline falls back to its own defaults.
# Keep CI small while exercising the fixed 8+3-step distilled recipe.
# Height and width describe the final output; Stage 1 runs at half size.
form_data = {
"prompt": PROMPT,
"model": omni_server.model,
"height": 128,
"width": 128,
"num_frames": 9,
"fps": 8,
"num_inference_steps": 2,
"guidance_scale": 1.0,
"num_inference_steps": 8,
"seed": 42,
}
tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py
-243
View File
@@ -1,243 +0,0 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import hashlib
import os
import subprocess
import sys
from pathlib import Path
import numpy as np
import pytest
from PIL import Image
from tests.helpers.mark import hardware_test
REPO_ROOT = Path(__file__).resolve().parents[3]
T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
I2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
T2V_PROMPT = (
"At sunrise, a glowing paper lantern boat drifts through a narrow canal between mossy stone walls, "
"soft fog above the water, the camera slowly gliding forward as golden reflections shimmer across "
"the ripples, cinematic, realistic, highly detailed."
)
T2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion, distorted, oversaturated, artifacts"
I2V_PROMPT = "A cinematic dolly shot of a boat drifting on calm water at sunset"
I2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion"
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
def _get_ltx2_model() -> str:
return "Lightricks/LTX-2"
def _md5(path: Path) -> str:
digest = hashlib.md5(usedforsecurity=False)
with path.open("rb") as f:
for chunk in iter(lambda: f.read(1024 * 1024), b""):
digest.update(chunk)
return digest.hexdigest()
def _make_deterministic_test_image(path: Path) -> None:
"""Create a deterministic 256x256 test image for I2V tests."""
rng = np.random.RandomState(42)
img = Image.fromarray(rng.randint(0, 255, (256, 256, 3), dtype=np.uint8))
img.save(path)
def _run_and_check(cmd: list[str], env: dict, output_path: Path, expected_md5: str) -> None:
result = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True, check=False)
assert result.returncode == 0, (
f"Command failed (exit {result.returncode}).\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
)
generated_md5 = _md5(output_path)
assert generated_md5 == expected_md5, (
f"Unexpected output md5: {generated_md5} != {expected_md5}.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
)
# ── T2V tests ──
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=2)
def test_ltx2_t2v_cfg_parallel(tmp_path: Path):
"""T2V with CFG=4.0, cfg-parallel-size=2."""
output = tmp_path / "t2v_cfg4.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
cmd = [
sys.executable,
str(T2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--prompt",
T2V_PROMPT,
"--negative-prompt",
T2V_NEGATIVE_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"145",
"--num-inference-steps",
"6",
"--guidance-scale",
"4.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--cfg-parallel-size",
"2",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="08e606b9c522fee4b6f30cee8b77db40")
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ltx2_t2v_no_cfg(tmp_path: Path):
"""T2V with CFG=1.0 (no classifier-free guidance)."""
output = tmp_path / "t2v_nocfg.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
cmd = [
sys.executable,
str(T2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--prompt",
T2V_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"145",
"--num-inference-steps",
"6",
"--guidance-scale",
"1.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="a83994b94b6e67c54a524e0383c45ce8")
# ── I2V tests ──
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.parallel
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=2)
def test_ltx2_i2v_cfg_parallel(tmp_path: Path):
"""I2V with CFG=4.0, cfg-parallel-size=2."""
test_image = tmp_path / "test_input.png"
_make_deterministic_test_image(test_image)
output = tmp_path / "i2v_cfg4.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
cmd = [
sys.executable,
str(I2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--model-class-name",
"LTX2ImageToVideoPipeline",
"--image",
str(test_image),
"--prompt",
I2V_PROMPT,
"--negative-prompt",
I2V_NEGATIVE_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"73",
"--num-inference-steps",
"6",
"--guidance-scale",
"4.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--cfg-parallel-size",
"2",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="aed7e56084b36373244d8f839b16d115")
@pytest.mark.full_model
@pytest.mark.diffusion
@pytest.mark.slow
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_ltx2_i2v_no_cfg(tmp_path: Path):
"""I2V with CFG=1.0 (no classifier-free guidance)."""
test_image = tmp_path / "test_input.png"
_make_deterministic_test_image(test_image)
output = tmp_path / "i2v_nocfg.mp4"
env = os.environ.copy()
env.setdefault("CUDA_VISIBLE_DEVICES", "0")
cmd = [
sys.executable,
str(I2V_EXAMPLE),
"--model",
_get_ltx2_model(),
"--model-class-name",
"LTX2ImageToVideoPipeline",
"--image",
str(test_image),
"--prompt",
I2V_PROMPT,
"--height",
"256",
"--width",
"256",
"--num-frames",
"73",
"--num-inference-steps",
"6",
"--guidance-scale",
"1.0",
"--frame-rate",
"24",
"--fps",
"24",
"--seed",
"42",
"--enforce-eager",
"--output",
str(output),
]
_run_and_check(cmd, env, output, expected_md5="81b21ede12753e9e14a357a6c548b666")
tests/model_extras/test_model_extras.py
+30
View File
@@ -181,6 +181,36 @@ def test_magi_human_extra_registry_declares_request_and_response_params() -> Non
assert should_init_extra_args_for_non_diffusion_stages("MagiHumanPipeline") is False
@pytest.mark.core_model
@pytest.mark.cpu
def test_ltx_extra_registry_declares_official_guidance_params() -> None:
expected = frozenset(
{
"video_cfg_scale",
"audio_cfg_scale",
"video_cfg_guidance_scale",
"audio_cfg_guidance_scale",
"video_stg_scale",
"audio_stg_scale",
"video_stg_guidance_scale",
"audio_stg_guidance_scale",
"video_modality_scale",
"audio_modality_scale",
"a2v_guidance_scale",
"v2a_guidance_scale",
"video_rescale_scale",
"audio_rescale_scale",
"video_stg_blocks",
"audio_stg_blocks",
}
)
assert get_extra_body_params("LTX2Pipeline") == expected
assert get_extra_output_params("LTX2Pipeline") == frozenset()
assert get_extra_body_params("LTX2DistilledPipeline") == expected
assert get_extra_output_params("LTX2DistilledPipeline") == frozenset()
@pytest.mark.core_model
@pytest.mark.cpu
def test_ming_flash_omni_extra_registry_declares_request_and_response_params() -> None:
tests/model_tests/diffusion/test_alignment.py
+1 -5
View File
@@ -30,13 +30,9 @@ EXCLUDED_MODELS = [
"OvisImagePipeline",
"WanPipeline",
"WanVACEPipeline",
"LTX2ImageToVideoPipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2DistilledPipeline",
"LTX2T2VDMD2Pipeline",
"LTX2I2VDMD2Pipeline",
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
"StableAudioPipeline",
"WanImageToVideoPipeline",
"WanS2VPipeline",
vllm_omni/diffusion/model_loader/diffusers_loader.py
+1 -1
View File
@@ -634,7 +634,7 @@ class DiffusersPipelineLoader:
# Discover pipeline components (DiT, encoders, VAEs) via
# ModuleDiscovery, which consults SupportsComponentDiscovery
# when available and falls back to well-known attribute names.
# This supports nested pipelines (e.g. LTX2TwoStagesPipeline
# This supports nested pipelines (e.g. LTX2DistilledPipeline
# where the transformer lives at "pipe.transformer").
discovered_modules = ModuleDiscovery.discover(model)
vllm_omni/diffusion/models/ltx2/__init__.py
+2 -14
View File
@@ -9,28 +9,16 @@ from vllm_omni.diffusion.models.ltx2.ltx2_components import (
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
LTX2I2VDMD2Pipeline,
LTX2ImageToVideoPipeline,
LTX2Pipeline,
LTX2T2VDMD2Pipeline,
LTX23ImageToVideoPipeline,
LTX23Pipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
LTX2ImageToVideoTwoStagesPipeline,
LTX2TwoStagesPipeline,
)
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
__all__ = [
"LTX2Pipeline",
"LTX2T2VDMD2Pipeline",
"LTX2ImageToVideoPipeline",
"LTX2I2VDMD2Pipeline",
"LTX2LatentUpsamplePipeline",
"LTX2TwoStagesPipeline",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX23Pipeline",
"LTX23ImageToVideoPipeline",
"LTX2DistilledPipeline",
"get_ltx2_post_process_func",
"load_transformer_config",
"create_transformer_from_config",
vllm_omni/diffusion/models/ltx2/ltx2_components.py
+191 -18
View File
@@ -7,6 +7,7 @@ from __future__ import annotations
import inspect
import json
import logging
import os
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
@@ -14,11 +15,14 @@ from typing import TYPE_CHECKING, Any
import torch
from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
from diffusers.pipelines.ltx2 import LTX2TextConnectors
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
from diffusers.video_processor import VideoProcessor
from huggingface_hub import hf_hub_download
from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video
from vllm_omni.diffusion.distributed.utils import get_local_device
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -28,7 +32,12 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
if TYPE_CHECKING:
from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
from .ltx2_transformer import LTX2VideoTransformer3DModel
from .ltx2_transformer import (
LTX2VideoTransformer3DModel,
apply_interleaved_rotary_emb,
apply_split_rotary_emb,
to_ltx_padding_mask,
)
try:
from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
@@ -44,7 +53,9 @@ _LTX_COMPONENT_SUBFOLDERS = (
"audio_vae",
"vocoder",
"scheduler",
"latent_upsampler",
)
logger = logging.getLogger(__name__)
@dataclass(frozen=True)
@@ -81,6 +92,164 @@ LTX23_COMPONENT_PROFILE = LTXComponentProfile(
vocoder_fallback_cls=LTX2Vocoder,
)
LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
name="ltx2_distilled",
dit_modules=("transformer",),
encoder_modules=("text_encoder", "connectors"),
vae_modules=("vae", "audio_vae"),
resident_modules=("vocoder", "latent_upsampler"),
video_vae_cls=DistributedAutoencoderKLLTX2Video,
)
_COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = {
("one_stage", "2"): LTX2_COMPONENT_PROFILE,
("one_stage", "2.3"): LTX23_COMPONENT_PROFILE,
("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE,
("dmd2", "2"): LTX2_COMPONENT_PROFILE,
("dmd2", "2.3"): LTX23_COMPONENT_PROFILE,
}
def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile:
"""Resolve component construction independently from execution recipes."""
try:
return _COMPONENT_PROFILES[(pipeline_kind, model_version)]
except KeyError as exc:
raise ValueError(f"Unsupported LTX component kind/version: {pipeline_kind!r}/{model_version!r}.") from exc
def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]:
"""Load small checkpoint metadata without relying on repository names."""
if os.path.isdir(model):
path = os.path.join(model, filename)
if not os.path.isfile(path):
return {}
else:
try:
path = hf_hub_download(repo_id=model, filename=filename)
except Exception:
return {}
try:
with open(path) as config_file:
value = json.load(config_file)
except (OSError, TypeError, ValueError):
return {}
return value if isinstance(value, dict) else {}
def detect_ltx_model_version(model: str) -> str:
"""Detect LTX-2 versus LTX-2.3 from checkpoint component metadata.
Official checkpoints use ``model_version`` metadata. Diffusers repositories
expose the equivalent distinction through the BWE vocoder introduced by
LTX-2.3. Unknown conversions retain the official LTX-2 fallback.
"""
model_index = _load_ltx_metadata_json(model, "model_index.json")
if str(model_index.get("model_version", "")).startswith("2.3"):
return "2.3"
vocoder_entry = model_index.get("vocoder")
if isinstance(vocoder_entry, (list, tuple)) and vocoder_entry:
vocoder_class = str(vocoder_entry[-1])
elif isinstance(vocoder_entry, dict):
vocoder_class = str(vocoder_entry.get("_class_name", ""))
else:
vocoder_class = ""
if vocoder_class == "LTX2VocoderWithBWE":
return "2.3"
vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json")
if str(vocoder_config.get("model_version", "")).startswith("2.3"):
return "2.3"
if vocoder_config.get("_class_name") == "LTX2VocoderWithBWE":
return "2.3"
logger.info("Using LTX-2 defaults for checkpoint %s", model)
return "2"
class _LTXConnectorAttnProcessor:
"""Preserve official connector math around Omni attention dispatch."""
def __init__(self, *, has_learned_registers: bool) -> None:
self.has_learned_registers = has_learned_registers
def __call__(
self,
attn: Any,
hidden_states: torch.Tensor,
encoder_hidden_states: torch.Tensor | None = None,
attention_mask: torch.Tensor | None = None,
query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
) -> torch.Tensor:
encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None
query = attn.norm_q(attn.to_q(hidden_states))
key = attn.norm_k(attn.to_k(encoder_hidden_states))
value = attn.to_v(encoder_hidden_states)
if query_rotary_emb is not None:
key_rotary_emb = key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
if attn.rope_type == "interleaved":
query = apply_interleaved_rotary_emb(query, query_rotary_emb)
key = apply_interleaved_rotary_emb(key, key_rotary_emb)
elif attn.rope_type == "split":
query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim)
key = apply_split_rotary_emb(key, key_rotary_emb, head_dim=attn.head_dim)
else:
raise ValueError(f"Unsupported LTX connector RoPE type: {attn.rope_type}")
batch_size, _, inner_dim = query.shape
head_dim = inner_dim // attn.heads
kv_heads = attn.inner_kv_dim // attn.head_dim
query = query.view(batch_size, -1, attn.heads, head_dim)
key = key.view(batch_size, -1, kv_heads, head_dim)
value = value.view(batch_size, -1, kv_heads, head_dim)
# The connector replaces padding tokens with learned registers before
# entering its blocks, so every key is valid and the old padding mask
# becomes an all-keep no-op.
if self.has_learned_registers:
attention_mask = None
elif attention_mask is not None and attn.omni_attention.attn_backend.get_name().upper() == "FLASH_ATTN":
attention_mask = to_ltx_padding_mask(attention_mask)
attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
hidden_states = attn.omni_attention(query, key, value, attn_metadata)
hidden_states = hidden_states.reshape(batch_size, -1, inner_dim)
if gate_logits is not None:
hidden_states = hidden_states.unflatten(2, (attn.heads, -1))
hidden_states = hidden_states * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = attn.to_out[0](hidden_states)
return attn.to_out[1](hidden_states)
def _install_connector_attention(connectors: LTX2TextConnectors) -> None:
for connector_name in ("video_connector", "audio_connector"):
connector = getattr(connectors, connector_name, None)
has_learned_registers = getattr(connector, "learnable_registers", None) is not None
for block_index, block in enumerate(getattr(connector, "transformer_blocks", ())):
attention = getattr(block, "attn1", None)
if attention is not None:
attention.omni_attention = OmniAttention(
num_heads=attention.heads,
head_size=attention.head_dim,
num_kv_heads=attention.inner_kv_dim // attention.head_dim,
softmax_scale=1.0 / (attention.head_dim**0.5),
causal=False,
prefix=f"connectors.{connector_name}.transformer_blocks.{block_index}.attn1",
role="ltx2.connector",
role_category="self",
skip_sequence_parallel=True,
disable_kv_quant=True,
)
attention.set_processor(_LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers))
def _detect_vocoder_output_sample_rate(model: str) -> int | None:
"""Read the generated waveform sample rate from the vocoder config."""
@@ -188,6 +357,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
local_files_only=local_files_only,
dtype=dtype,
)
_install_connector_attention(pipeline.connectors)
pipeline.vae = _load_component(
profile.video_vae_cls,
model,
@@ -221,6 +391,18 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
dtype=dtype,
)
if "latent_upsampler" in profile.resident_modules:
# BlurDownsample constructs an integer kernel that must be initialized
# on CPU; component placement is handled uniformly after construction.
with torch.device("cpu"):
pipeline.latent_upsampler = _load_component(
LTX2LatentUpsamplerModel,
model,
"latent_upsampler",
local_files_only=local_files_only,
dtype=dtype,
)
transformer_config = load_transformer_config(model, "transformer", local_files_only)
quant_config = getattr(od_config, "quantization_config", None)
pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
@@ -249,8 +431,6 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
tokenizer_max_length = getattr(encoder_config, "max_seq_len", None)
pipeline.tokenizer_max_length = int(tokenizer_max_length or 1024)
pipeline._guidance_scale = None
pipeline._guidance_rescale = None
pipeline._interrupt = False
@@ -262,22 +442,15 @@ def load_transformer_config(
"""Load an LTX transformer config from a local model or the HF Hub."""
if local_files_only:
config_path = os.path.join(model_path, subfolder, "config.json")
if os.path.exists(config_path):
with open(config_path) as config_file:
return json.load(config_file)
if not os.path.exists(config_path):
raise FileNotFoundError(f"LTX transformer config not found: {config_path}")
else:
try:
from huggingface_hub import hf_hub_download
config_path = hf_hub_download(
repo_id=model_path,
filename=f"{subfolder}/config.json",
)
with open(config_path) as config_file:
return json.load(config_file)
except Exception:
pass
return {}
config_path = hf_hub_download(
repo_id=model_path,
filename=f"{subfolder}/config.json",
)
with open(config_path) as config_file:
return json.load(config_file)
def create_transformer_from_config(
vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py
+331 -130
View File
@@ -5,17 +5,20 @@
from __future__ import annotations
import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import numpy as np
import PIL.Image
import torch
import torch.nn.functional as F
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 . import ltx2_latents as latent_ops
from .ltx2_denoise import I2VVideoAudioScheduler
from .ltx2_guidance import euler_step_from_velocity
if TYPE_CHECKING:
from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext
@@ -44,6 +47,41 @@ def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_outputs: int) ->
return tensor.repeat_interleave(num_outputs, dim=0)
def _preprocess_i2v_pil_images(
images: PIL.Image.Image | list[PIL.Image.Image],
*,
height: int,
width: int,
device: torch.device | None = None,
dtype: torch.dtype | None = None,
) -> torch.Tensor:
"""Match official LTX aspect-preserving resize and center crop."""
image_list = images if isinstance(images, list) else [images]
processed = []
for image in image_list:
pixels = (
torch.from_numpy(np.array(image.convert("RGB"), copy=True))
.permute(2, 0, 1)
.unsqueeze(0)
.to(device=device, dtype=torch.float32)
)
source_height, source_width = pixels.shape[-2:]
scale = max(height / source_height, width / source_width)
resized_height = math.ceil(source_height * scale)
resized_width = math.ceil(source_width * scale)
pixels = F.interpolate(
pixels,
size=(resized_height, resized_width),
mode="bilinear",
align_corners=False,
)
crop_top = (resized_height - height) // 2
crop_left = (resized_width - width) // 2
pixels = pixels[:, :, crop_top : crop_top + height, crop_left : crop_left + width]
processed.append((pixels / 127.5 - 1.0).to(dtype=dtype))
return torch.cat(processed, dim=0)
class LTXTextConditioningMixin:
"""Shared Gemma encoding and text connector orchestration."""
@@ -113,21 +151,64 @@ class LTXTextConditioningMixin:
batch_size = len(prompt) if prompt is not None else 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,
num_videos_per_prompt=num_videos_per_prompt,
if do_classifier_free_guidance and prompt_embeds is None and negative_prompt_embeds is None:
negative_prompt = negative_prompt or ""
negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
if prompt is not None and type(prompt) is not type(negative_prompt):
raise TypeError(
f"`negative_prompt` should be the same type as `prompt`, but got {type(negative_prompt)} !="
f" {type(prompt)}."
)
if isinstance(negative_prompt, list) and batch_size != len(negative_prompt):
raise ValueError(
f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has batch size"
f" {batch_size}."
)
# Official LTX encodes positive and negative prompts in one Gemma batch;
# separate bf16 GEMMs produce a different conditioning trajectory.
combined_embeds, combined_mask = self._get_gemma_prompt_embeds(
prompt=prompt + negative_prompt,
num_videos_per_prompt=1,
max_sequence_length=max_sequence_length,
scale_factor=scale_factor,
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,
)
prompt_embeds = combined_embeds[:batch_size]
negative_prompt_embeds = combined_embeds[batch_size:]
prompt_attention_mask = combined_mask[:batch_size]
negative_prompt_attention_mask = combined_mask[batch_size:]
if num_videos_per_prompt > 1:
prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
negative_prompt_embeds = _repeat_prompt_tensor_for_outputs(
negative_prompt_embeds,
num_videos_per_prompt,
)
prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
prompt_attention_mask,
num_videos_per_prompt,
)
negative_prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
negative_prompt_attention_mask,
num_videos_per_prompt,
)
else:
if prompt_embeds is None:
prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
prompt=prompt,
num_videos_per_prompt=num_videos_per_prompt,
max_sequence_length=max_sequence_length,
scale_factor=scale_factor,
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 ""
@@ -243,9 +324,10 @@ class LTXTextConditioningMixin:
class LTXI2VConditioningMixin:
"""First-frame conditioning behavior common to LTX2 and LTX2.3."""
"""First-frame conditioning with optional unified T2V/I2V dispatch."""
support_image_input = True
unified_text_image_entry = False
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
@@ -254,10 +336,6 @@ class LTXI2VConditioningMixin:
resample="bilinear",
)
def forward(self, req: Any, image: Any | None = None, *args: Any, **kwargs: Any) -> Any:
"""Preserve the public I2V positional API while sharing the T2V runner."""
return super().forward(req, *args, image=image, **kwargs)
@staticmethod
def _resolve_single_prompt_image(raw_image: Any) -> Any:
if isinstance(raw_image, list):
@@ -283,11 +361,12 @@ class LTXI2VConditioningMixin:
image: Any | None,
request_inputs: LTXRequestInputs,
) -> Any | None:
if image is not None or not req.prompts:
prompts = getattr(req, "prompts", None)
if image is not None or not prompts:
return image
raw_images = []
for prompt_item in req.prompts:
for prompt_item in prompts:
if isinstance(prompt_item, str):
raw_image = None
else:
@@ -300,8 +379,12 @@ class LTXI2VConditioningMixin:
raw_image = PIL.Image.open(raw_image).convert("RGB")
raw_images.append(raw_image)
if all(raw_image is None for raw_image in raw_images) and self.unified_text_image_entry:
return None
if any(raw_image is None for raw_image in raw_images) and request_inputs.latents is None:
raise ValueError("Image is required for LTX I2V generation.")
if not self.unified_text_image_entry:
raise ValueError("Image is required for LTX I2V generation.")
raise ValueError("LTX request batches cannot mix text-to-video and image-to-video prompts.")
if len(raw_images) == 1:
return raw_images[0]
return raw_images or image
@@ -311,75 +394,19 @@ class LTXI2VConditioningMixin:
request_inputs: LTXRequestInputs,
image: Any | None = None,
) -> None:
if image is None and request_inputs.latents is None:
if not self.unified_text_image_entry and image is None and request_inputs.latents is None:
raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
super()._check_forward_inputs(request_inputs, image=image)
def prepare_latents(
def _encode_i2v_image_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]:
num_frames, height, width = latent_ops.resolve_video_latent_shape(
height,
width,
num_frames,
vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
)
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 = latent_ops.normalize_latents(
latents,
self.vae.latents_mean,
self.vae.latents_std,
self.vae.config.scaling_factor,
)
latents = latent_ops.create_noised_state(
latents,
noise_scale * (1 - conditioning_mask),
generator,
)
latents = latent_ops.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 = latent_ops.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: torch.Tensor,
*,
batch_size: int,
generator: torch.Generator | list[torch.Generator] | None,
dtype: torch.dtype | None,
) -> torch.Tensor:
"""Encode one conditioning frame per prompt into normalized latents."""
image_batch_size = image.shape[0]
if image_batch_size == 0:
raise ValueError("`image` batch is empty.")
@@ -415,38 +442,173 @@ class LTXI2VConditioningMixin:
for img in image
]
init_latents = torch.cat(init_latents, dim=0).to(dtype)
init_latents = torch.cat(init_latents, dim=0).to(dtype=dtype)
if num_videos_per_prompt > 1:
init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
init_latents = latent_ops.normalize_latents(
return latent_ops.normalize_latents(
init_latents,
self.vae.latents_mean,
self.vae.latents_std,
)
@staticmethod
def _replace_first_latent_frame(
latents: torch.Tensor,
image_latents: torch.Tensor,
) -> torch.Tensor:
expected_shape = (latents.shape[0], latents.shape[1], 1, latents.shape[3], latents.shape[4])
if tuple(image_latents.shape) != expected_shape:
raise ValueError(
f"Encoded LTX conditioning image has shape {tuple(image_latents.shape)}, expected {expected_shape}."
)
latents = latents.clone()
latents[:, :, :1] = image_latents
return latents
def prepare_latents(
self,
image: torch.Tensor | None = None,
batch_size: int = 1,
num_channels_latents: int = 128,
height: int = 512,
width: int = 768,
num_frames: int = 121,
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,
) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
if image is None and self.unified_text_image_entry:
return super().prepare_latents(
batch_size=batch_size,
num_channels_latents=num_channels_latents,
height=height,
width=width,
num_frames=num_frames,
noise_scale=noise_scale,
dtype=dtype,
device=device,
generator=generator,
latents=latents,
)
num_frames, height, width = latent_ops.resolve_video_latent_shape(
height,
width,
num_frames,
vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
)
mask_shape = (batch_size, 1, num_frames, height, width)
if latents is not None:
unpacked_latents = latents.ndim == 5
if latents.ndim == 5:
expected_shape = (batch_size, num_channels_latents, num_frames, height, width)
if tuple(latents.shape) != expected_shape:
raise ValueError(
f"Provided unpacked `latents` have shape {tuple(latents.shape)}, expected {expected_shape}."
)
conditioning_mask = latents.new_zeros(mask_shape)
conditioning_mask[:, :, 0] = 1.0
latents = latent_ops.normalize_latents(
latents,
self.vae.latents_mean,
self.vae.latents_std,
self.vae.config.scaling_factor,
)
if image is not None:
image_latents = self._encode_i2v_image_latents(
image,
batch_size=batch_size,
generator=generator,
dtype=dtype,
)
# Official distilled Stage 2 starts from the upsampled Stage 1
# state, then reapplies the source image at final resolution.
latents = self._replace_first_latent_frame(latents, image_latents)
latents = latent_ops.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 = latent_ops.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,)}."
)
if unpacked_latents:
# Official LTX patchifies the latent state before drawing noise.
# Sampling the packed shape preserves its seed-to-output contract.
latents = latent_ops.create_noised_state(
latents,
noise_scale * (1 - conditioning_mask).unsqueeze(-1),
generator,
)
elif image is not None:
image_latents = self._encode_i2v_image_latents(
image,
batch_size=batch_size,
generator=generator,
dtype=dtype,
)
unpacked = latent_ops.unpack_latents(
latents,
num_frames,
height,
width,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
unpacked = self._replace_first_latent_frame(unpacked, image_latents)
latents = latent_ops.pack_latents(
unpacked,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
return latents.to(device=device, dtype=dtype), conditioning_mask
if image is None:
raise ValueError("`image` must be provided when `latents` is None.")
init_latents = self._encode_i2v_image_latents(
image,
batch_size=batch_size,
generator=generator,
dtype=dtype,
)
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)
init_latents = latent_ops.pack_latents(
init_latents,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
conditioning_mask = latent_ops.pack_latents(
conditioning_mask,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
).squeeze(-1)
latents = latent_ops.pack_latents(
latents,
self.transformer_spatial_patch_size,
self.transformer_temporal_patch_size,
)
noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype)
latents = init_latents * conditioning_mask.unsqueeze(-1) + noise * (1 - conditioning_mask).unsqueeze(-1)
return latents, conditioning_mask
def _step_video_latents_i2v(
self,
noise_pred_video: torch.Tensor,
latents: torch.Tensor,
timestep: torch.Tensor,
step_index: int,
latent_num_frames: int,
latent_height: int,
latent_width: int,
@@ -469,12 +631,14 @@ class LTXI2VConditioningMixin:
)
noise_pred_video = noise_pred_video[:, :, 1:]
noise_latents = latents_unpacked[:, :, 1:]
pred_latents = self.scheduler.step(
noise_pred_video,
timestep,
# Official LTX advances the unconditioned frames directly on the
# sigma trajectory; the first frame remains fixed conditioning.
pred_latents = euler_step_from_velocity(
noise_latents,
return_dict=False,
)[0]
noise_pred_video,
self.scheduler.sigmas,
step_index,
)
latents_unpacked = torch.cat([latents_unpacked[:, :, :1], pred_latents], dim=2)
return latent_ops.pack_latents(
latents_unpacked,
@@ -491,12 +655,31 @@ class LTXI2VConditioningMixin:
noise_scale: float,
image: Any | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
if request_inputs.latents is None:
if image is None and self.unified_text_image_entry:
return super()._prepare_video_latents_stage(
request_inputs,
prompt_context,
device=device,
noise_scale=noise_scale,
image=None,
)
if image is not 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)
elif isinstance(image, PIL.Image.Image) or (
isinstance(image, list) and image and isinstance(image[0], PIL.Image.Image)
):
image = _preprocess_i2v_pil_images(
image,
height=request_inputs.height,
width=request_inputs.width,
device=device,
dtype=prompt_context.positive_connector_prompt_embeds.dtype,
)
else:
image = self.video_processor.preprocess(
image,
@@ -509,32 +692,17 @@ class LTXI2VConditioningMixin:
)
return self.prepare_latents(
image,
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
self.transformer.config.in_channels,
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,
image=image,
batch_size=prompt_context.batch_size * request_inputs.num_videos_per_prompt,
num_channels_latents=self.transformer.config.in_channels,
height=request_inputs.height,
width=request_inputs.width,
num_frames=request_inputs.num_frames,
noise_scale=noise_scale,
dtype=prompt_context.positive_connector_prompt_embeds.dtype,
device=device,
generator=request_inputs.generator,
latents=request_inputs.latents,
)
def _prepare_denoise_context_for_cfg(
@@ -544,6 +712,8 @@ class LTXI2VConditioningMixin:
) -> LTXDenoiseContext:
denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
if denoise_ctx.conditioning_mask is None:
if self.unified_text_image_entry:
return denoise_ctx
raise ValueError("LTX I2V denoising requires a conditioning mask.")
mask_batch = denoise_ctx.conditioning_mask.shape[0]
@@ -564,15 +734,46 @@ class LTXI2VConditioningMixin:
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
*,
video_token_count: int,
audio_token_count: int,
) -> dict[str, torch.Tensor]:
kwargs = super()._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx)
if denoise_ctx.conditioning_mask is None:
if self.unified_text_image_entry:
return super()._denoise_timestep_kwargs(
ts,
forward_ctx,
denoise_ctx,
video_token_count=video_token_count,
audio_token_count=audio_token_count,
)
raise ValueError("LTX I2V denoising requires a conditioning mask.")
kwargs = super()._denoise_timestep_kwargs(
ts,
forward_ctx,
denoise_ctx,
video_token_count=video_token_count,
audio_token_count=audio_token_count,
)
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("LTX I2V denoising requires a conditioning mask.")
kwargs.update(
timestep=ts.unsqueeze(-1) * (1 - conditioning_mask),
audio_timestep=ts,
timestep=ts.reshape(-1, 1) * (1 - conditioning_mask),
)
return kwargs
def _video_guidance_model_sigma(
self,
sigma: torch.Tensor,
denoise_ctx: LTXDenoiseContext,
) -> torch.Tensor:
if denoise_ctx.conditioning_mask is None:
if self.unified_text_image_entry:
return super()._video_guidance_model_sigma(sigma, denoise_ctx)
raise ValueError("LTX I2V guidance requires a conditioning mask.")
# The conditioned tokens are evaluated at timestep zero by the model,
# so velocity-to-x0 guidance must use that same per-token sigma.
return sigma.reshape(-1, 1, 1) * (1 - denoise_ctx.conditioning_mask).unsqueeze(-1)
vllm_omni/diffusion/models/ltx2/ltx2_denoise.py
+114 -63
View File
@@ -6,20 +6,22 @@
from __future__ import annotations
import copy
import math
from collections.abc import Callable, Iterable
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any, Protocol
import numpy as np
import torch
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from .ltx2_latents import LTXAVState
from .ltx2_guidance import euler_step_from_velocity
from .ltx2_latents import LTXAVState, unpack_audio_latents, unpad_audio_latents
if TYPE_CHECKING:
from .ltx2_conditioning import LTXPromptContext
from .ltx2_recipes import LTXPhaseRecipe
from .ltx2_request import LTXRequestInputs
@@ -41,7 +43,7 @@ class LTXForwardContext:
padded_audio_num_frames: int
timesteps: torch.Tensor
audio_scheduler: Any
video_audio_scheduler: Any
video_audio_step_adapter: Any
@property
def batch_size(self) -> int:
@@ -71,6 +73,7 @@ class LTXPhaseResult:
forward_context: LTXForwardContext
video: torch.Tensor
audio: torch.Tensor
audio_for_next_phase: torch.Tensor | None = None
class LTXDenoisePipeline(Protocol):
@@ -122,61 +125,81 @@ def calculate_shift(
return image_seq_len * slope + intercept
class VideoAudioScheduler:
"""Composite scheduler dispatching video and audio updates."""
class LTXVideoAudioStepAdapter:
"""Expose the shared LTX Euler update through the distributed scheduler API."""
def __init__(self, video_scheduler, audio_scheduler):
self.video_scheduler = video_scheduler
self.audio_scheduler = audio_scheduler
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self.video_scheduler.step(
noise_pred[0],
t[0],
latents[0],
return_dict=False,
generator=generator,
)[0]
audio_out = self.audio_scheduler.step(
noise_pred[1],
t[1],
latents[1],
return_dict=False,
generator=generator,
)[0]
return ((video_out, audio_out),)
class I2VVideoAudioScheduler:
"""Update the unconditioned video frames and the full audio state."""
def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
self.video_scheduler = pipeline.scheduler
self.audio_scheduler = audio_scheduler
def __init__(
self,
pipeline: Any,
audio_scheduler: Any,
latent_num_frames: int,
latent_height: int,
latent_width: int,
*,
image_conditioned: bool,
) -> None:
self._pipeline = pipeline
self._audio_scheduler = audio_scheduler
self._latent_num_frames = latent_num_frames
self._latent_height = latent_height
self._latent_width = latent_width
self._image_conditioned = image_conditioned
self._step_index = 0
def step(self, noise_pred, t, latents, return_dict=False, generator=None):
video_out = self._pipeline._step_video_latents_i2v(
noise_pred[0],
latents[0],
t[0],
self._latent_num_frames,
self._latent_height,
self._latent_width,
)
audio_out = self.audio_scheduler.step(
noise_pred[1],
t[1],
del t, return_dict, generator
if self._image_conditioned:
video_out = self._pipeline._step_video_latents_i2v(
noise_pred[0],
latents[0],
self._step_index,
self._latent_num_frames,
self._latent_height,
self._latent_width,
)
else:
video_out = euler_step_from_velocity(
latents[0],
noise_pred[0],
self._pipeline.scheduler.sigmas,
self._step_index,
)
audio_out = euler_step_from_velocity(
latents[1],
return_dict=False,
generator=generator,
)[0]
noise_pred[1],
self._audio_scheduler.sigmas,
self._step_index,
)
self._step_index += 1
return ((video_out, audio_out),)
def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> torch.Tensor:
"""Build the official LTX one-stage sigma schedule in torch fp32."""
config = scheduler.config
base_anchor = config.get("base_image_seq_len", 1024)
max_anchor = config.get("max_image_seq_len", 4096)
base_shift = config.get("base_shift", 0.95)
max_shift = config.get("max_shift", 2.05)
sigmas = torch.linspace(1.0, 0.0, steps + 1)
slope = (max_shift - base_shift) / (max_anchor - base_anchor)
# Official LTX2 one-stage intentionally uses the max sequence anchor, so the shift stays at max_shift.
sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor)
exp_shift = math.exp(sigma_shift)
sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0)
# Official non-distilled checkpoints explicitly set this to 0.1. Missing
# or None means the loaded scheduler disabled terminal stretching.
terminal = config.get("shift_terminal")
if terminal is not None:
non_zero = sigmas != 0
one_minus_sigmas = 1.0 - sigmas[non_zero]
scale = one_minus_sigmas[-1] / (1.0 - terminal)
sigmas[non_zero] = 1.0 - one_minus_sigmas / scale
return sigmas.to(dtype=torch.float32, device=device)
def _set_scheduler_sigmas(scheduler: Any, sigmas: torch.Tensor) -> torch.Tensor:
sigmas = sigmas.to(torch.float32)
timesteps = sigmas[:-1] * scheduler.config.get("num_train_timesteps", 1000)
@@ -198,16 +221,20 @@ def prepare_scheduler_stage(
latent_num_frames: int,
latent_height: int,
latent_width: int,
use_official_sigma_schedule: bool,
image_conditioned: bool = False,
) -> tuple[Any, Any, torch.Tensor]:
if sigmas is not None and timesteps is not None:
raise ValueError("Only one of `sigmas` or `timesteps` may be provided.")
audio_scheduler = copy.deepcopy(pipeline.scheduler)
video_audio_scheduler = pipeline._make_video_audio_scheduler(
video_audio_step_adapter = LTXVideoAudioStepAdapter(
pipeline,
audio_scheduler,
latent_num_frames,
latent_height,
latent_width,
image_conditioned=image_conditioned,
)
if sigmas is not None:
scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device)
@@ -217,12 +244,14 @@ def prepare_scheduler_stage(
scheduler_sigmas = torch.cat([scheduler_sigmas, scheduler_sigmas.new_zeros(1)])
timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas)
_set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone())
return audio_scheduler, video_audio_scheduler, timesteps_tensor
return audio_scheduler, video_audio_step_adapter, timesteps_tensor
if sigmas is None and timesteps is None and use_official_sigma_schedule:
scheduler_sigmas = _official_ltx_sigmas(pipeline.scheduler, request_inputs.num_inference_steps, device)
timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas)
_set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone())
return audio_scheduler, video_audio_step_adapter, timesteps_tensor
sigmas = np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps)
# Official LTX2 one-stage omits the latent, so mu stays at max_shift across resolutions.
# Mu varies with latent token count only when the scheduler receives a latent, as in official HQ two-stage;
# explicit sigma schedules bypass mu entirely.
mu = calculate_shift(
pipeline.scheduler.config.get("max_image_seq_len", 4096),
pipeline.scheduler.config.get("base_image_seq_len", 1024),
@@ -246,7 +275,7 @@ def prepare_scheduler_stage(
sigmas=sigmas,
mu=mu,
)
return audio_scheduler, video_audio_scheduler, timesteps_tensor
return audio_scheduler, video_audio_step_adapter, timesteps_tensor
def prepare_rope_coords_stage(
@@ -283,15 +312,25 @@ def build_transformer_kwargs(
encoder_attention_mask: torch.Tensor | None,
audio_encoder_attention_mask: torch.Tensor | None,
ts: torch.Tensor,
attention_kwargs: dict[str, Any] | None = None,
) -> dict[str, Any]:
del encoder_attention_mask, audio_encoder_attention_mask
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,
**pipeline._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx),
"encoder_attention_mask": encoder_attention_mask,
"audio_encoder_attention_mask": audio_encoder_attention_mask,
**pipeline._denoise_timestep_kwargs(
ts,
forward_ctx,
denoise_ctx,
video_token_count=hidden_states.shape[1],
audio_token_count=audio_hidden_states.shape[1],
),
# This is valid only because LTX connectors replace every padding token
# with a learned register, making all output context tokens valid.
"encoder_attention_mask": None,
"audio_encoder_attention_mask": None,
"num_frames": forward_ctx.latent_num_frames,
"height": forward_ctx.latent_height,
"width": forward_ctx.latent_width,
@@ -299,7 +338,7 @@ def build_transformer_kwargs(
"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,
"attention_kwargs": forward_ctx.attention_kwargs if attention_kwargs is None else attention_kwargs,
"return_dict": False,
}
@@ -317,7 +356,7 @@ def step_denoised_latents(
(timestep, timestep),
(denoise_ctx.latents, denoise_ctx.audio_latents),
do_true_cfg=pipeline.do_classifier_free_guidance,
per_request_scheduler=forward_ctx.video_audio_scheduler,
per_request_scheduler=forward_ctx.video_audio_step_adapter,
)
return pipeline._synchronize_cfg_parallel_step_output(
latents,
@@ -338,11 +377,12 @@ class LTXPhaseExecutor:
sigmas: list[float] | None,
timesteps: list[int] | None,
attention_kwargs: dict[str, Any] | None,
phase_recipe: LTXPhaseRecipe,
image: Any | None = None,
prompt_context: LTXPromptContext | None = None,
) -> LTXPhaseResult:
pipeline._check_forward_inputs(request_inputs, image=image)
cfg_parallel_ready = pipeline._setup_forward_runtime(request_inputs, attention_kwargs)
cfg_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs)
device = pipeline.device
if prompt_context is None:
prompt_context = pipeline._prepare_prompt_context(
@@ -372,7 +412,7 @@ class LTXPhaseExecutor:
noise_scale=noise_scale,
)
)
audio_scheduler, video_audio_scheduler, timesteps_tensor = prepare_scheduler_stage(
audio_scheduler, video_audio_step_adapter, timesteps_tensor = prepare_scheduler_stage(
pipeline,
request_inputs,
device=device,
@@ -381,6 +421,8 @@ class LTXPhaseExecutor:
latent_num_frames=latent_num_frames,
latent_height=latent_height,
latent_width=latent_width,
use_official_sigma_schedule=phase_recipe.use_official_sigma_schedule,
image_conditioned=conditioning_mask is not None,
)
forward_ctx = LTXForwardContext(
req=req,
@@ -397,7 +439,7 @@ class LTXPhaseExecutor:
padded_audio_num_frames=padded_audio_num_frames,
timesteps=timesteps_tensor,
audio_scheduler=audio_scheduler,
video_audio_scheduler=video_audio_scheduler,
video_audio_step_adapter=video_audio_step_adapter,
)
video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents)
denoise_ctx = LTXDenoiseContext(
@@ -427,4 +469,13 @@ class LTXPhaseExecutor:
state.video,
state.audio,
)
return LTXPhaseResult(forward_context=forward_ctx, video=latents, audio=audio_latents)
normalized_audio = unpack_audio_latents(
unpad_audio_latents(state.audio, forward_ctx.original_audio_num_frames),
num_mel_bins=forward_ctx.latent_mel_bins,
)
return LTXPhaseResult(
forward_context=forward_ctx,
video=latents,
audio=audio_latents,
audio_for_next_phase=normalized_audio,
)
vllm_omni/diffusion/models/ltx2/ltx2_guidance.py
+362 -311
View File
@@ -1,16 +1,16 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Guidance strategies for the LTX model family."""
"""Official multi-modal guidance for the LTX model family."""
from __future__ import annotations
from typing import TYPE_CHECKING, Any, Protocol
import math
from dataclasses import dataclass
from typing import TYPE_CHECKING, Any
import torch
from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.parallel_state import (
get_cfg_group,
get_classifier_free_guidance_rank,
@@ -22,358 +22,409 @@ if TYPE_CHECKING:
from .ltx2_latents import LTXAVState
class LTXGuidanceStrategy(Protocol):
"""Model-version guidance behavior consumed by the shared pipeline."""
@dataclass(frozen=True)
class LTXModalityGuidance:
"""Guidance parameters applied to one LTX output modality."""
def validate_cfg_world_size(self, cfg_world_size: int) -> None: ...
cfg_scale: float = 1.0
stg_scale: float = 0.0
modality_scale: float = 1.0
rescale_scale: float = 0.0
stg_blocks: tuple[int, ...] = ()
def prepare_denoise_context(
self,
pipeline: Any,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext: ...
@property
def do_cfg(self) -> bool:
return not math.isclose(self.cfg_scale, 1.0)
def timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> dict[str, torch.Tensor]: ...
@property
def do_stg(self) -> bool:
return not math.isclose(self.stg_scale, 0.0) and bool(self.stg_blocks)
def combine_cfg_noise(
self,
pipeline: Any,
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
true_cfg_scale: float,
cfg_normalize: bool,
context: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]: ...
def predict_noise(
self,
pipeline: Any,
index: int,
timestep: torch.Tensor,
state: LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]: ...
@property
def do_modality_guidance(self) -> bool:
return not math.isclose(self.modality_scale, 1.0)
class _LTXGuidanceStrategyBase:
def validate_cfg_world_size(self, cfg_world_size: int) -> None:
del cfg_world_size
@dataclass(frozen=True)
class LTXGuidanceSpec:
"""Video and audio guidance requested for one denoise phase."""
def prepare_denoise_context(
self,
pipeline: Any,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext:
del pipeline, forward_ctx
return denoise_ctx
video: LTXModalityGuidance = LTXModalityGuidance()
audio: LTXModalityGuidance = LTXModalityGuidance()
def timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> dict[str, torch.Tensor]:
del forward_ctx, denoise_ctx
return {"timestep": ts}
@classmethod
def positive_only(cls) -> LTXGuidanceSpec:
return cls()
@property
def do_cfg(self) -> bool:
return self.video.do_cfg or self.audio.do_cfg
@property
def do_stg(self) -> bool:
return self.video.do_stg or self.audio.do_stg
@property
def do_modality_guidance(self) -> bool:
return self.video.do_modality_guidance or self.audio.do_modality_guidance
@property
def do_rescale(self) -> bool:
return self.video.rescale_scale != 0.0 or self.audio.rescale_scale != 0.0
def combine_velocity_via_x0(
@dataclass(frozen=True)
class LTXDenoisePass:
"""One conditioned Transformer evaluation in a guidance batch."""
name: str
negative_video_context: bool = False
negative_audio_context: bool = False
@dataclass(frozen=True)
class LTXGuidancePlan:
"""Concrete Transformer passes derived from a guidance specification."""
spec: LTXGuidanceSpec
passes: tuple[LTXDenoisePass, ...]
@classmethod
def build(cls, spec: LTXGuidanceSpec) -> LTXGuidancePlan:
passes = [LTXDenoisePass("cond")]
if spec.do_cfg:
passes.append(
LTXDenoisePass(
"uncond",
negative_video_context=spec.video.do_cfg,
negative_audio_context=spec.audio.do_cfg,
)
)
if spec.do_stg:
passes.append(LTXDenoisePass("ptb"))
if spec.do_modality_guidance:
passes.append(LTXDenoisePass("mod"))
return cls(spec=spec, passes=tuple(passes))
@property
def names(self) -> tuple[str, ...]:
return tuple(item.name for item in self.passes)
@property
def cfg_parallel_compatible(self) -> bool:
return self.names == ("cond", "uncond") and not self.spec.do_rescale
def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
"""Convert velocity to x0 using the official fp32 arithmetic order."""
sigma = sigma.to(torch.float32)
return (sample.to(torch.float32) - velocity.to(torch.float32) * sigma).to(sample.dtype)
def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
"""Convert x0 back to velocity using the official fp32 arithmetic order."""
sigma = sigma.to(torch.float32)
return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype)
def euler_step_from_velocity(
sample: torch.Tensor,
positive_velocity: torch.Tensor,
negative_velocity: torch.Tensor,
sigma: torch.Tensor,
guidance_scale: float,
velocity: torch.Tensor,
sigmas: torch.Tensor,
step_index: int,
) -> torch.Tensor:
"""Apply CFG in x0 space and convert the result back to velocity."""
x0_cond = sample - positive_velocity * sigma
x0_uncond = sample - negative_velocity * sigma
x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond)
return (sample - x0_guided) / sigma
sigma = sigmas[step_index].to(torch.float32)
sigma_next = sigmas[step_index + 1].to(torch.float32)
return (sample.float() + velocity.float() * (sigma_next - sigma)).to(sample.dtype)
class LTXLegacyVelocityGuidance(_LTXGuidanceStrategyBase):
"""Existing LTX2 Diffusers-style velocity-space CFG behavior."""
def combine_cfg_noise(
self,
pipeline: Any,
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
true_cfg_scale: float,
cfg_normalize: bool,
context: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
del context
video_pos, audio_pos = positive_noise_pred
video_neg, audio_neg = negative_noise_pred
video_combined = CFGParallelMixin.combine_cfg_noise(
pipeline,
video_pos,
video_neg,
true_cfg_scale,
cfg_normalize,
)
audio_combined = CFGParallelMixin.combine_cfg_noise(
pipeline,
audio_pos,
audio_neg,
true_cfg_scale,
cfg_normalize,
)
if pipeline.guidance_rescale and pipeline.guidance_rescale > 0:
video_combined = rescale_noise_cfg(
video_combined,
video_pos,
guidance_rescale=pipeline.guidance_rescale,
)
audio_combined = rescale_noise_cfg(
audio_combined,
audio_pos,
guidance_rescale=pipeline.guidance_rescale,
)
return video_combined, audio_combined
def predict_noise(
self,
pipeline: Any,
index: int,
timestep: torch.Tensor,
state: LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]:
del index
prompt_context = forward_ctx.prompt_context
video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
expanded_timestep = timestep.expand(video_input.shape[0])
positive_kwargs = pipeline._build_transformer_kwargs(
forward_ctx,
denoise_ctx,
hidden_states=video_input,
audio_hidden_states=audio_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=expanded_timestep,
)
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,
}
if pipeline.do_classifier_free_guidance
else None
)
return pipeline.predict_noise_maybe_with_cfg(
do_true_cfg=pipeline.do_classifier_free_guidance,
true_cfg_scale=forward_ctx.request_inputs.guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
)
def combine_guided_x0(
*,
cond: torch.Tensor,
uncond_text: torch.Tensor | float,
uncond_perturbed: torch.Tensor | float,
uncond_modality: torch.Tensor | float,
guidance: LTXModalityGuidance,
) -> torch.Tensor:
dtype = cond.dtype
cond = cond.float()
uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text
uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed
uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality
pred = cond + (guidance.cfg_scale - 1) * (cond - uncond_text)
if guidance.do_stg:
pred = pred + guidance.stg_scale * (cond - uncond_perturbed)
pred = pred + (guidance.modality_scale - 1) * (cond - uncond_modality)
if guidance.rescale_scale != 0.0:
# The pinned official one-stage entry runs one generated sample and
# uses its full-tensor standard deviation. In Omni, dimension 0 may
# instead contain independently scheduled requests. Reduce each item
# separately so its result cannot depend on unrelated co-batched work.
reduce_dims = tuple(range(1, pred.ndim))
cond_std = cond.std(dim=reduce_dims, keepdim=True)
pred_std = pred.std(dim=reduce_dims, keepdim=True)
eps = torch.finfo(pred.dtype).eps
factor = cond_std / pred_std.clamp_min(eps)
factor = torch.where(pred_std > eps, factor, torch.ones_like(factor))
factor = guidance.rescale_scale * factor + (1 - guidance.rescale_scale)
pred = pred * factor
return pred.to(dtype)
class LTXOfficialX0Guidance(_LTXGuidanceStrategyBase):
"""Official LTX x0-space CFG behavior currently used by LTX2.3."""
def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor:
return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1))
def validate_cfg_world_size(self, cfg_world_size: int) -> None:
if cfg_world_size not in (1, 2):
raise ValueError(
f"LTX x0 guidance supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}."
)
def _mask_for_pass(
pass_count: int,
batch_size: int,
pass_index: int,
reference: torch.Tensor,
) -> torch.Tensor:
mask = reference.new_ones((pass_count * batch_size, 1, 1))
start = pass_index * batch_size
mask[start : start + batch_size] = 0
return mask
def build_perturbation_kwargs(
plan: LTXGuidancePlan,
batch_size: int,
reference: torch.Tensor,
) -> dict[str, Any]:
kwargs: dict[str, Any] = {}
if "ptb" in plan.names:
mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("ptb"), reference)
if plan.spec.video.do_stg:
kwargs["video_self_attention_mask"] = mask
kwargs["video_self_attention_blocks"] = plan.spec.video.stg_blocks
if plan.spec.audio.do_stg:
kwargs["audio_self_attention_mask"] = mask
kwargs["audio_self_attention_blocks"] = plan.spec.audio.stg_blocks
if "mod" in plan.names:
mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("mod"), reference)
kwargs["a2v_cross_attention_mask"] = mask
kwargs["v2a_cross_attention_mask"] = mask
return kwargs
class LTXGuidanceExecutor:
"""Build and execute official LTX multi-guidance Transformer passes."""
@staticmethod
def validate_cfg_world_size(plan: LTXGuidancePlan, cfg_world_size: int) -> None:
if cfg_world_size == 1:
return
if cfg_world_size != 2:
raise ValueError(f"LTX CFG parallelism supports cfg_parallel_size 1 or 2, got {cfg_world_size}.")
if not plan.cfg_parallel_compatible:
raise ValueError("LTX CFG parallelism only supports CFG-only guidance without rescale.")
@staticmethod
def prepare_denoise_context(
self,
pipeline: Any,
forward_ctx: LTXForwardContext,
plan: LTXGuidancePlan,
cfg_parallel_ready: bool,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext:
if pipeline.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)
)
if len(plan.passes) > 1 and not cfg_parallel_ready:
denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, len(plan.passes))
denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, len(plan.passes))
return denoise_ctx
@staticmethod
def timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
video_token_count: int,
audio_token_count: int,
) -> dict[str, torch.Tensor]:
del forward_ctx, denoise_ctx
return {"timestep": ts, "sigma": ts}
return {
"timestep": ts.reshape(-1, 1).expand(-1, video_token_count),
"audio_timestep": ts.reshape(-1, 1).expand(-1, audio_token_count),
"sigma": ts,
"audio_sigma": ts,
}
def combine_cfg_noise(
self,
pipeline: Any,
positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
true_cfg_scale: float,
cfg_normalize: bool,
context: dict[str, Any],
) -> tuple[torch.Tensor, torch.Tensor]:
required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma")
if any(context.get(name) is None for name in required):
raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.")
video_pos, audio_pos = positive_noise_pred
video_neg, audio_neg = negative_noise_pred
video_combined = combine_velocity_via_x0(
context["video_latents"],
video_pos,
video_neg,
context["video_sigma"],
true_cfg_scale,
@staticmethod
def combine_cfg_velocity(
sample: torch.Tensor,
positive_velocity: torch.Tensor,
negative_velocity: torch.Tensor,
sigma: torch.Tensor,
guidance_scale: float,
*,
model_sigma: torch.Tensor | None = None,
) -> torch.Tensor:
guidance = LTXModalityGuidance(cfg_scale=guidance_scale)
model_sigma = sigma if model_sigma is None else model_sigma
guided_x0 = combine_guided_x0(
cond=x0_from_velocity(sample, positive_velocity, model_sigma),
uncond_text=x0_from_velocity(sample, negative_velocity, model_sigma),
uncond_perturbed=0.0,
uncond_modality=0.0,
guidance=guidance,
)
audio_combined = combine_velocity_via_x0(
context["audio_latents"],
audio_pos,
audio_neg,
context["audio_sigma"],
true_cfg_scale,
)
if cfg_normalize:
video_combined = pipeline.cfg_normalize_function(video_pos, video_combined)
audio_combined = pipeline.cfg_normalize_function(audio_pos, audio_combined)
return video_combined, audio_combined
return velocity_from_x0(sample, guided_x0, sigma)
def predict_parallel_cfg(
self,
pipeline: Any,
true_cfg_scale: float,
positive_kwargs: dict[str, Any],
negative_kwargs: dict[str, Any],
cfg_normalize: bool = True,
output_slice: int | None = None,
**context: Any,
) -> tuple[torch.Tensor, torch.Tensor]:
def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
if output_slice is None:
return pred
return pred[0][:, :output_slice], pred[1][:, :output_slice]
cfg_world_size = get_classifier_free_guidance_world_size()
self.validate_cfg_world_size(cfg_world_size)
if cfg_world_size != 2:
raise ValueError(f"Parallel LTX x0 guidance requires cfg_parallel_size 2, but got {cfg_world_size}.")
cfg_rank = get_classifier_free_guidance_rank()
branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
local_video_pred, local_audio_pred = maybe_slice(pipeline.predict_noise(**branch_kwargs))
cfg_group = get_cfg_group()
gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True)
gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True)
return self.combine_cfg_noise(
pipeline,
(gathered_video[0], gathered_audio[0]),
(gathered_video[1], gathered_audio[1]),
true_cfg_scale,
cfg_normalize,
context,
)
def predict_noise(
self,
pipeline: Any,
plan: LTXGuidancePlan,
index: int,
timestep: torch.Tensor,
state: LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]:
prompt_context = forward_ctx.prompt_context
guidance_scale = forward_ctx.request_inputs.guidance_scale
audio_scheduler = forward_ctx.audio_scheduler
if forward_ctx.cfg_parallel_ready:
video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
ts = timestep.expand(video_input.shape[0])
positive_kwargs = pipeline._build_transformer_kwargs(
forward_ctx,
denoise_ctx,
hidden_states=video_input,
audio_hidden_states=audio_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,
}
return self.predict_parallel_cfg(
pipeline,
true_cfg_scale=guidance_scale,
positive_kwargs=positive_kwargs,
negative_kwargs=negative_kwargs,
cfg_normalize=False,
video_latents=state.video,
audio_latents=state.audio,
video_sigma=pipeline.scheduler.sigmas[index],
audio_sigma=audio_scheduler.sigmas[index],
)
video_input = torch.cat([state.video] * 2) if pipeline.do_classifier_free_guidance else state.video
video_input = video_input.to(prompt_context.connector_prompt_embeds.dtype)
audio_input = torch.cat([state.audio] * 2) if pipeline.do_classifier_free_guidance else state.audio
audio_input = audio_input.to(prompt_context.connector_prompt_embeds.dtype)
ts = timestep.expand(video_input.shape[0])
transformer_kwargs = pipeline._build_transformer_kwargs(
self.validate_cfg_world_size(plan, get_classifier_free_guidance_world_size())
prompt = forward_ctx.prompt_context
negative = get_classifier_free_guidance_rank() == 1
video_context = (
prompt.negative_connector_prompt_embeds
if negative and plan.spec.video.do_cfg
else prompt.positive_connector_prompt_embeds
)
audio_context = (
prompt.negative_connector_audio_prompt_embeds
if negative and plan.spec.audio.do_cfg
else prompt.positive_connector_audio_prompt_embeds
)
if video_context is None or audio_context is None:
raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
video_input = state.video.to(prompt.positive_connector_prompt_embeds.dtype)
audio_input = state.audio.to(prompt.positive_connector_audio_prompt_embeds.dtype)
kwargs = pipeline._build_transformer_kwargs(
forward_ctx,
denoise_ctx,
hidden_states=video_input,
audio_hidden_states=audio_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,
encoder_hidden_states=video_context,
audio_encoder_hidden_states=audio_context,
encoder_attention_mask=None,
audio_encoder_attention_mask=None,
ts=timestep.expand(video_input.shape[0]),
)
with pipeline._transformer_cache_context("cond_uncond"):
video_pred, audio_pred = pipeline.transformer(**transformer_kwargs)
video_pred = video_pred.float()
audio_pred = audio_pred.float()
if pipeline.do_classifier_free_guidance:
video_uncond, video_cond = video_pred.chunk(2)
video_pred = combine_velocity_via_x0(
local_video, local_audio = pipeline.predict_noise(**kwargs)
group = get_cfg_group()
video = group.all_gather(local_video, separate_tensors=True)
audio = group.all_gather(local_audio, separate_tensors=True)
video_sigma = pipeline.scheduler.sigmas[index]
return (
self.combine_cfg_velocity(
state.video,
video_cond,
video_uncond,
pipeline.scheduler.sigmas[index],
guidance_scale,
)
audio_uncond, audio_cond = audio_pred.chunk(2)
audio_pred = combine_velocity_via_x0(
video[0],
video[1],
video_sigma,
plan.spec.video.cfg_scale,
model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx),
),
self.combine_cfg_velocity(
state.audio,
audio_cond,
audio_uncond,
audio_scheduler.sigmas[index],
guidance_scale,
audio[0],
audio[1],
forward_ctx.audio_scheduler.sigmas[index],
plan.spec.audio.cfg_scale,
),
)
def predict_noise(
self,
pipeline: Any,
plan: LTXGuidancePlan,
index: int,
timestep: torch.Tensor,
state: LTXAVState,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]:
if forward_ctx.cfg_parallel_ready:
return self.predict_parallel_cfg(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx)
prompt = forward_ctx.prompt_context
video_contexts: list[torch.Tensor] = []
audio_contexts: list[torch.Tensor] = []
for denoise_pass in plan.passes:
video_context = (
prompt.negative_connector_prompt_embeds
if denoise_pass.negative_video_context
else prompt.positive_connector_prompt_embeds
)
return video_pred, audio_pred
audio_context = (
prompt.negative_connector_audio_prompt_embeds
if denoise_pass.negative_audio_context
else prompt.positive_connector_audio_prompt_embeds
)
if video_context is None or audio_context is None:
raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
video_contexts.append(video_context)
audio_contexts.append(audio_context)
pass_count = len(plan.passes)
video_input = _repeat_batch(state.video, pass_count).to(prompt.positive_connector_prompt_embeds.dtype)
audio_input = _repeat_batch(state.audio, pass_count).to(prompt.positive_connector_audio_prompt_embeds.dtype)
attention_kwargs = dict(forward_ctx.attention_kwargs or {})
perturbations = build_perturbation_kwargs(plan, state.video.shape[0], video_input)
if perturbations:
attention_kwargs["ltx_perturbation_kwargs"] = perturbations
kwargs = pipeline._build_transformer_kwargs(
forward_ctx,
denoise_ctx,
hidden_states=video_input,
audio_hidden_states=audio_input,
encoder_hidden_states=torch.cat(video_contexts),
audio_encoder_hidden_states=torch.cat(audio_contexts),
encoder_attention_mask=None,
audio_encoder_attention_mask=None,
ts=timestep.expand(video_input.shape[0]),
attention_kwargs=attention_kwargs,
)
with pipeline._transformer_cache_context("guided"):
video_velocity, audio_velocity = pipeline.transformer(**kwargs)
video_splits = dict(zip(plan.names, video_velocity.chunk(pass_count), strict=True))
audio_splits = dict(zip(plan.names, audio_velocity.chunk(pass_count), strict=True))
def guide_modality(
sample: torch.Tensor,
splits: dict[str, torch.Tensor],
sigma: torch.Tensor,
guidance: LTXModalityGuidance,
*,
model_sigma: torch.Tensor | None = None,
) -> torch.Tensor:
model_sigma = sigma if model_sigma is None else model_sigma
# Official guidance reduces contiguous BSC tensors. The fused
# transformer output may be channel-major after splitting, which
# changes fp32 reduction order and can cross bf16 rounding bounds.
x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()}
guided = combine_guided_x0(
cond=x0["cond"],
uncond_text=x0.get("uncond", 0.0),
uncond_perturbed=x0.get("ptb", 0.0),
uncond_modality=x0.get("mod", 0.0),
guidance=guidance,
)
return velocity_from_x0(sample, guided, sigma)
return (
guide_modality(
state.video,
video_splits,
pipeline.scheduler.sigmas[index],
plan.spec.video,
model_sigma=pipeline._video_guidance_model_sigma(
pipeline.scheduler.sigmas[index],
denoise_ctx,
),
),
guide_modality(
state.audio,
audio_splits,
forward_ctx.audio_scheduler.sigmas[index],
plan.spec.audio,
),
)
LTX_LEGACY_VELOCITY_GUIDANCE = LTXLegacyVelocityGuidance()
LTX_OFFICIAL_X0_GUIDANCE = LTXOfficialX0Guidance()
LTX_GUIDANCE_EXECUTOR = LTXGuidanceExecutor()
# Kept as a narrow compatibility helper for CFG-parallel tests and extensions.
combine_velocity_via_x0 = LTXGuidanceExecutor.combine_cfg_velocity
vllm_omni/diffusion/models/ltx2/ltx2_latents.py
+18 -46
View File
@@ -120,45 +120,14 @@ def create_noised_state(
return noise_scale * noise + (1 - noise_scale) * latents
def pack_audio_latents(
latents: torch.Tensor,
patch_size: int | None = None,
patch_size_t: int | None = None,
) -> torch.Tensor:
if patch_size is not None and patch_size_t is not None:
batch_size, _, latent_length, latent_mel_bins = latents.shape
post_patch_latent_length = latent_length / patch_size_t
post_patch_mel_bins = latent_mel_bins / patch_size
latents = latents.reshape(
batch_size,
-1,
post_patch_latent_length,
patch_size_t,
post_patch_mel_bins,
patch_size,
)
return latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor:
return latents.transpose(1, 2).flatten(2, 3)
def unpack_audio_latents(
latents: torch.Tensor,
latent_length: int,
num_mel_bins: int,
patch_size: int | None = None,
patch_size_t: int | None = None,
) -> torch.Tensor:
if patch_size is not None and patch_size_t is not None:
batch_size = latents.size(0)
latents = latents.reshape(
batch_size,
latent_length,
num_mel_bins,
-1,
patch_size_t,
patch_size,
)
return latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
return latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
@@ -224,18 +193,19 @@ def prepare_video_latents(
vae_spatial_compression_ratio=pipeline.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=pipeline.vae_temporal_compression_ratio,
)
shape = (batch_size, num_channels_latents, num_frames, height, width)
spatial_patch_size = pipeline.transformer_spatial_patch_size
temporal_patch_size = pipeline.transformer_temporal_patch_size
shape = (
batch_size,
(num_frames // temporal_patch_size) * (height // spatial_patch_size) * (width // spatial_patch_size),
num_channels_latents * temporal_patch_size * spatial_patch_size * spatial_patch_size,
)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return pack_latents(
latents,
pipeline.transformer_spatial_patch_size,
pipeline.transformer_temporal_patch_size,
)
return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
def prepare_audio_latents(
@@ -249,6 +219,7 @@ def prepare_audio_latents(
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
latents_normalized: bool = False,
) -> tuple[torch.Tensor, int, int]:
original_latent_length = audio_latent_length
latent_mel_bins = num_mel_bins // pipeline.audio_vae_mel_compression_ratio
@@ -260,11 +231,12 @@ def prepare_audio_latents(
latents = pack_audio_latents(latents)
if latents.ndim != 3:
raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].")
latents = normalize_audio_latents(
latents,
pipeline.audio_vae.latents_mean,
pipeline.audio_vae.latents_std,
)
if not latents_normalized:
latents = normalize_audio_latents(
latents,
pipeline.audio_vae.latents_mean,
pipeline.audio_vae.latents_std,
)
latents = create_noised_state(latents, noise_scale, generator)
if latents.shape[1] not in {original_latent_length, padded_latent_length}:
@@ -283,11 +255,11 @@ def prepare_audio_latents(
latents = torch.cat([latents, padding], dim=1)
return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins)
shape = (batch_size, padded_latent_length, num_channels_latents * latent_mel_bins)
if isinstance(generator, list) and len(generator) != batch_size:
raise ValueError(
f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
return pack_audio_latents(latents), original_latent_length, padded_latent_length
return latents, original_latent_length, padded_latent_length
vllm_omni/diffusion/models/ltx2/ltx2_recipes.py
+160 -5
View File
@@ -1,20 +1,175 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""User-visible one-stage defaults for the LTX model family."""
"""Declarative execution recipes for the LTX model family."""
from dataclasses import dataclass
from typing import Literal
from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance
LTX_DEFAULT_NEGATIVE_PROMPT = (
"blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
"grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
"deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
"wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of "
"field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent "
"lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny "
"valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, "
"mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, "
"off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward "
"pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, "
"inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
)
LTX_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0)
LTX_STAGE_2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0)
@dataclass(frozen=True)
class LTXOneStageRecipe:
class LTXPhaseRecipe:
"""One denoise phase and the transition used to construct its input."""
name: str
guidance: LTXGuidanceSpec
spatial_downscale: int = 1
sigmas: tuple[float, ...] | None = None
noise_scale: float = 0.0
input_transform: Literal["initial", "spatial_upsample"] = "initial"
allow_guidance_override: bool = True
use_official_sigma_schedule: bool = True
def __post_init__(self) -> None:
if self.spatial_downscale < 1:
raise ValueError("LTX phase spatial_downscale must be positive.")
if self.sigmas is not None and len(self.sigmas) < 2:
raise ValueError("An explicit LTX sigma schedule must contain at least two denoise sigmas.")
@property
def num_inference_steps(self) -> int | None:
return None if self.sigmas is None else len(self.sigmas) - 1
@dataclass(frozen=True)
class LTXPipelineRecipe:
"""Request defaults, ordered phases, output routing, and request capabilities."""
phases: tuple[LTXPhaseRecipe, ...]
height: int = 512
width: int = 768
num_frames: int = 121
frame_rate: float = 24.0
num_inference_steps: int = 40
guidance_scale: float = 4.0
negative_prompt: str = LTX_DEFAULT_NEGATIVE_PROMPT
video_output_phase: int = -1
audio_output_phase: int = -1
allow_request_sigmas: bool = True
allow_request_latents: bool = True
allow_negative_prompt: bool = True
fixed_num_inference_steps: bool = False
def __post_init__(self) -> None:
if not self.phases:
raise ValueError("An LTX pipeline recipe must contain at least one phase.")
if self.phases[0].input_transform != "initial":
raise ValueError("The first LTX phase must use the initial request input.")
phase_count = len(self.phases)
for output_phase in (self.video_output_phase, self.audio_output_phase):
if not -phase_count <= output_phase < phase_count:
raise ValueError(f"LTX output phase {output_phase} is outside a {phase_count}-phase recipe.")
@property
def request_guidance(self) -> LTXGuidanceSpec:
return self.phases[0].guidance
@property
def max_spatial_downscale(self) -> int:
return max(phase.spatial_downscale for phase in self.phases)
LTX2_ONE_STAGE_RECIPE = LTXOneStageRecipe()
LTX23_ONE_STAGE_RECIPE = LTXOneStageRecipe()
def _official_guidance(stg_block: int) -> LTXGuidanceSpec:
return LTXGuidanceSpec(
video=LTXModalityGuidance(
cfg_scale=3.0,
stg_scale=1.0,
modality_scale=3.0,
rescale_scale=0.7,
stg_blocks=(stg_block,),
),
audio=LTXModalityGuidance(
cfg_scale=7.0,
stg_scale=1.0,
modality_scale=3.0,
rescale_scale=0.7,
stg_blocks=(stg_block,),
),
)
LTX2_ONE_STAGE_RECIPE = LTXPipelineRecipe(
phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(29)),),
)
LTX23_ONE_STAGE_RECIPE = LTXPipelineRecipe(
num_inference_steps=30,
phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(28)),),
)
LTX_POSITIVE_ONLY_RECIPE = LTXPipelineRecipe(
phases=(
LTXPhaseRecipe(
name="generate",
guidance=LTXGuidanceSpec.positive_only(),
use_official_sigma_schedule=False,
),
),
)
LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe(
# Official distilled arguments describe the final output. Stage 1 applies
# spatial_downscale=2 and therefore runs at 512x768.
height=1024,
width=1536,
num_inference_steps=len(LTX_DISTILLED_SIGMAS) - 1,
negative_prompt="",
phases=(
LTXPhaseRecipe(
name="generate_lowres",
guidance=LTXGuidanceSpec.positive_only(),
spatial_downscale=2,
sigmas=LTX_DISTILLED_SIGMAS,
noise_scale=1.0,
allow_guidance_override=False,
use_official_sigma_schedule=False,
),
LTXPhaseRecipe(
name="refine",
guidance=LTXGuidanceSpec.positive_only(),
sigmas=LTX_STAGE_2_DISTILLED_SIGMAS,
noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0],
input_transform="spatial_upsample",
allow_guidance_override=False,
use_official_sigma_schedule=False,
),
),
video_output_phase=1,
audio_output_phase=1,
allow_request_sigmas=False,
allow_request_latents=False,
allow_negative_prompt=False,
fixed_num_inference_steps=True,
)
_PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = {
("one_stage", "2"): LTX2_ONE_STAGE_RECIPE,
("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE,
("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE,
("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE,
("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE,
}
def resolve_ltx_pipeline_recipe(pipeline_kind: str, model_version: str) -> LTXPipelineRecipe:
"""Resolve execution independently from component loading."""
try:
return _PIPELINE_RECIPES[(pipeline_kind, model_version)]
except KeyError as exc:
raise ValueError(f"Unsupported LTX pipeline kind/version: {pipeline_kind!r}/{model_version!r}.") from exc
vllm_omni/diffusion/models/ltx2/ltx2_request.py
+289 -32
View File
@@ -1,17 +1,21 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Request normalization shared by LTX pipeline variants."""
"""Request normalization and validation shared by LTX pipeline variants."""
from __future__ import annotations
from dataclasses import dataclass
from dataclasses import dataclass, replace
from typing import Any
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance
from .ltx2_recipes import LTXPipelineRecipe
@dataclass
class LTXRequestInputs:
@@ -24,8 +28,7 @@ class LTXRequestInputs:
num_frames: int
frame_rate: float
num_inference_steps: int
guidance_scale: float
guidance_rescale: float
guidance: LTXGuidanceSpec
num_videos_per_prompt: int
generator: torch.Generator | list[torch.Generator] | None
latents: torch.Tensor | None
@@ -38,6 +41,65 @@ class LTXRequestInputs:
decode_noise_scale: float | list[float] | None
output_type: str
max_sequence_length: int
audio_latents_normalized: bool = False
@property
def guidance_scale(self) -> float:
"""Compatibility view for callers that only understand video CFG."""
return self.guidance.video.cfg_scale
def validate_pipeline_request(
request_inputs: LTXRequestInputs,
*,
pipeline_recipe: LTXPipelineRecipe,
vae_spatial_compression_ratio: int,
vae_temporal_compression_ratio: int,
pipeline_name: str,
request_sigmas: list[float] | None = None,
) -> None:
"""Validate the request capabilities declared by an execution recipe."""
alignment = vae_spatial_compression_ratio * pipeline_recipe.max_spatial_downscale
if request_inputs.height % alignment != 0 or request_inputs.width % alignment != 0:
raise ValueError(
f"{pipeline_name} resolution must be divisible by {alignment}, got "
f"{request_inputs.width}x{request_inputs.height}."
)
if request_inputs.num_frames < 1 or (request_inputs.num_frames - 1) % vae_temporal_compression_ratio != 0:
raise ValueError(
f"{pipeline_name} `num_frames` must be {vae_temporal_compression_ratio} * k + 1, "
f"got {request_inputs.num_frames}."
)
if not pipeline_recipe.allow_request_latents and (
request_inputs.latents is not None or request_inputs.audio_latents is not None
):
raise ValueError(f"{pipeline_name} does not accept request-provided video or audio latents.")
if not pipeline_recipe.allow_request_sigmas and request_sigmas is not None:
raise ValueError(f"{pipeline_name} uses fixed phase sigma schedules.")
if pipeline_recipe.fixed_num_inference_steps and (
request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps
):
raise ValueError(
f"{pipeline_name} uses {pipeline_recipe.num_inference_steps} fixed Stage 1 denoise steps; "
f"got {request_inputs.num_inference_steps}."
)
negative_prompt = request_inputs.negative_prompt
has_negative_prompt = (
any(bool(value) for value in negative_prompt) if isinstance(negative_prompt, list) else bool(negative_prompt)
)
if not pipeline_recipe.allow_negative_prompt and (
has_negative_prompt or request_inputs.negative_prompt_embeds is not None
):
raise ValueError(f"{pipeline_name} uses positive-only guidance and does not accept negative conditioning.")
first_phase = pipeline_recipe.phases[0]
if not first_phase.allow_guidance_override and request_inputs.guidance != first_phase.guidance:
raise ValueError(f"{pipeline_name} uses fixed positive-only guidance and cannot override it.")
def _unwrap_request_tensor(value: Any) -> Any:
@@ -65,11 +127,187 @@ def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None:
return sampling.extra_args.get("audio_latents")
def _get_extra_arg(sampling: Any, *names: str) -> Any:
extra_args = sampling.extra_args or {}
for name in names:
value = extra_args.get(name)
if value is not None:
return value
return None
def _normalize_stg_blocks(value: Any, default: tuple[int, ...]) -> tuple[int, ...]:
if value is None:
return default
if isinstance(value, int):
return (value,)
if not isinstance(value, (list, tuple)):
raise ValueError(f"STG blocks must be an integer or a list of integers, got {value!r}.")
return tuple(int(item) for item in value)
def _resolve_negative_prompts(
prompts: list[Any],
fallback: str | list[str],
) -> str | list[str]:
if not prompts:
return fallback
if isinstance(fallback, list) and len(fallback) != len(prompts):
raise ValueError(
f"`negative_prompt` has batch size {len(fallback)}, but the request batch has size {len(prompts)}."
)
resolved = []
for index, item in enumerate(prompts):
request_value = None if isinstance(item, str) else item.get("negative_prompt")
if request_value is not None:
resolved.append(request_value)
else:
resolved.append(fallback[index] if isinstance(fallback, list) else fallback)
return resolved
def _resolve_modality_guidance(
sampling: Any,
modality: str,
default: LTXModalityGuidance,
) -> LTXModalityGuidance:
cfg_aliases = (f"{modality}_cfg_scale", f"{modality}_cfg_guidance_scale")
stg_aliases = (f"{modality}_stg_scale", f"{modality}_stg_guidance_scale")
modality_aliases = (
("video_modality_scale", "a2v_guidance_scale")
if modality == "video"
else ("audio_modality_scale", "v2a_guidance_scale")
)
cfg_scale = _get_extra_arg(sampling, *cfg_aliases)
stg_scale = _get_extra_arg(sampling, *stg_aliases)
modality_scale = _get_extra_arg(sampling, *modality_aliases)
rescale_scale = _get_extra_arg(sampling, f"{modality}_rescale_scale")
stg_blocks = _get_extra_arg(sampling, f"{modality}_stg_blocks")
return replace(
default,
cfg_scale=default.cfg_scale if cfg_scale is None else float(cfg_scale),
stg_scale=default.stg_scale if stg_scale is None else float(stg_scale),
modality_scale=default.modality_scale if modality_scale is None else float(modality_scale),
rescale_scale=default.rescale_scale if rescale_scale is None else float(rescale_scale),
stg_blocks=_normalize_stg_blocks(stg_blocks, default.stg_blocks),
)
def _resolve_guidance_spec(
sampling: Any,
default: LTXGuidanceSpec,
*,
guidance_scale: float | None,
) -> LTXGuidanceSpec:
video = _resolve_modality_guidance(sampling, "video", default.video)
audio = _resolve_modality_guidance(sampling, "audio", default.audio)
common_cfg = sampling.guidance_scale if sampling.guidance_scale_provided else guidance_scale
if common_cfg is not None:
video = replace(video, cfg_scale=float(common_cfg))
audio = replace(audio, cfg_scale=float(common_cfg))
return LTXGuidanceSpec(video=video, audio=audio)
class LTXRequestMixin:
"""Normalize serving requests without coupling them to a model version."""
supports_request_batch = False
supports_guidance_rescale = False
@staticmethod
def _validate_forward_options(
*,
timesteps: list[int] | None,
guidance_rescale: float | None,
noise_scale: float,
return_dict: bool,
attention_kwargs: dict[str, Any] | None,
) -> None:
if timesteps is not None:
raise ValueError("LTX does not support `timesteps`; use `sigmas` for a custom schedule.")
if guidance_rescale not in (None, 0.0):
raise ValueError(
"LTX does not support the common `guidance_rescale`; use `video_rescale_scale` and "
"`audio_rescale_scale`."
)
if noise_scale != 0.0:
raise ValueError("LTX only supports `noise_scale=0.0`.")
if return_dict is not True:
raise ValueError("LTX only supports `return_dict=True`.")
if attention_kwargs is not None:
raise ValueError("LTX does not support public `attention_kwargs`.")
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
*,
image: Any | 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 | None = None,
guidance_rescale: float | None = None,
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,
) -> DiffusionOutput | list[DiffusionOutput]:
"""Validate and normalize the public LTX request surface."""
self._validate_forward_options(
timesteps=timesteps,
guidance_rescale=guidance_rescale,
noise_scale=noise_scale,
return_dict=return_dict,
attention_kwargs=attention_kwargs,
)
if image is not None and not self.support_image_input:
raise ValueError(f"{self.__class__.__name__} does not support `image` input.")
return self._forward_request(
req,
image=image,
prompt=prompt,
negative_prompt=negative_prompt,
height=height,
width=width,
num_frames=num_frames,
frame_rate=frame_rate,
num_inference_steps=num_inference_steps,
sigmas=sigmas,
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,
)
@staticmethod
def _resolve_request_sigmas(
@@ -128,8 +366,7 @@ class LTXRequestMixin:
num_frames: int | None,
frame_rate: float | None,
num_inference_steps: int | None,
timesteps: list[int] | None,
guidance_scale: float,
guidance_scale: float | None,
num_videos_per_prompt: int | None,
generator: torch.Generator | list[torch.Generator] | None,
latents: torch.Tensor | None,
@@ -142,38 +379,55 @@ class LTXRequestMixin:
decode_noise_scale: float | list[float] | None,
output_type: str,
max_sequence_length: int | None,
guidance_rescale: float = 0.0,
) -> LTXRequestInputs:
sampling_params_list = req.sampling_params_list
sampling = sampling_params_list[0]
prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
if all(isinstance(item, str) or item.get("negative_prompt") is None for item in req.prompts):
negative_prompt = None
elif req.prompts:
negative_prompt = [
"" if isinstance(item, str) else (item.get("negative_prompt") or "") for item in req.prompts
]
for sampling in sampling_params_list:
if sampling.timesteps is not None:
raise ValueError("LTX does not support `timesteps`; use `sigmas` for a custom schedule.")
if sampling.guidance_rescale:
raise ValueError(
"LTX does not support the common `guidance_rescale`; use `video_rescale_scale` and "
"`audio_rescale_scale`."
)
if _get_extra_arg(sampling, "flow_shift") is not None:
raise ValueError("LTX does not support `flow_shift`; use `sigmas` for a custom schedule.")
height = sampling.height or height or self.one_stage_recipe.height
width = sampling.width or width or self.one_stage_recipe.width
num_frames = sampling.num_frames or num_frames or self.one_stage_recipe.num_frames
frame_rate = sampling.resolved_frame_rate or frame_rate or self.one_stage_recipe.frame_rate
num_inference_steps = (
sampling.num_inference_steps or num_inference_steps or self.one_stage_recipe.num_inference_steps
sampling = sampling_params_list[0]
is_dummy_run = req.is_dummy_run()
prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
negative_prompt = _resolve_negative_prompts(
req.prompts,
negative_prompt if negative_prompt is not None else self.pipeline_recipe.negative_prompt,
)
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.")
height = sampling.height or height or self.pipeline_recipe.height
width = sampling.width or width or self.pipeline_recipe.width
num_frames = sampling.num_frames or num_frames or self.pipeline_recipe.num_frames
frame_rate = sampling.resolved_frame_rate or frame_rate or self.pipeline_recipe.frame_rate
num_inference_steps = (
sampling.num_inference_steps or num_inference_steps or self.pipeline_recipe.num_inference_steps
)
if is_dummy_run and self.pipeline_recipe.fixed_num_inference_steps:
num_inference_steps = self.pipeline_recipe.num_inference_steps
num_inference_steps = max(int(num_inference_steps), 2)
num_videos_per_prompt = (
sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else num_videos_per_prompt or 1
)
max_sequence_length = sampling.max_sequence_length or max_sequence_length or self.tokenizer_max_length
if sampling.guidance_scale_provided:
guidance_scale = sampling.guidance_scale
if self.supports_guidance_rescale and sampling.guidance_rescale is not None:
guidance_rescale = sampling.guidance_rescale
guidance_specs = [
_resolve_guidance_spec(
item,
self.pipeline_recipe.request_guidance,
guidance_scale=guidance_scale,
)
for item in sampling_params_list
]
guidance = guidance_specs[0]
if is_dummy_run and not self.pipeline_recipe.phases[0].allow_guidance_override:
guidance = self.pipeline_recipe.request_guidance
if any(item != guidance for item in guidance_specs[1:]):
raise ValueError("Batched LTX requests must use identical video/audio guidance parameters.")
if self.supports_request_batch:
if generator is None:
@@ -235,6 +489,10 @@ class LTXRequestMixin:
]
if any(value is not None for value in request_negative_masks):
negative_prompt_attention_mask = torch.stack(request_negative_masks) # type: ignore[arg-type]
if prompt_embeds is not None:
prompt = None
if negative_prompt_embeds is not None:
negative_prompt = None
if sampling.decode_timestep is not None:
decode_timestep = sampling.decode_timestep
@@ -251,8 +509,7 @@ class LTXRequestMixin:
num_frames=int(num_frames),
frame_rate=float(frame_rate),
num_inference_steps=int(num_inference_steps),
guidance_scale=guidance_scale,
guidance_rescale=guidance_rescale,
guidance=guidance,
num_videos_per_prompt=int(num_videos_per_prompt),
generator=generator,
latents=latents,
vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py → vllm_omni/diffusion/models/ltx2/ltx2_runtime.py
+305 -65
View File
@@ -1,12 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Shared runtime surface for LTX pipeline variants."""
"""Shared recipe-driven runtime for LTX pipeline variants."""
from __future__ import annotations
from collections.abc import Iterable
from contextlib import nullcontext
from dataclasses import replace
from typing import Any, ClassVar
import torch
@@ -21,22 +22,41 @@ from vllm_omni.diffusion.models.interface import SupportsComponentDiscovery
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
from vllm_omni.platforms import current_omni_platform
from . import ltx2_latents as latent_ops
from .ltx2_components import LTXComponentProfile, initialize_pipeline_components
from .ltx2_components import (
LTXComponentProfile,
detect_ltx_model_version,
initialize_pipeline_components,
resolve_ltx_component_profile,
)
from .ltx2_conditioning import LTXPromptContext, LTXTextConditioningMixin
from .ltx2_denoise import (
LTXDenoiseContext,
LTXForwardContext,
LTXPhaseExecutor,
LTXPhaseResult,
VideoAudioScheduler,
build_transformer_kwargs,
step_denoised_latents,
)
from .ltx2_guidance import LTXGuidanceStrategy
from .ltx2_recipes import LTXOneStageRecipe
from .ltx2_request import LTXRequestInputs, LTXRequestMixin
from .ltx2_guidance import (
LTX_GUIDANCE_EXECUTOR,
LTXGuidanceExecutor,
LTXGuidancePlan,
LTXGuidanceSpec,
LTXModalityGuidance,
)
from .ltx2_recipes import (
LTXPhaseRecipe,
LTXPipelineRecipe,
resolve_ltx_pipeline_recipe,
)
from .ltx2_request import (
LTXRequestInputs,
LTXRequestMixin,
validate_pipeline_request,
)
def _expand_per_prompt_decode_value(
@@ -92,7 +112,7 @@ def _prepare_decode_timestep_conditioning(
)
class LTXPipelineRuntime(
class LTXRuntime(
LTXRequestMixin,
LTXTextConditioningMixin,
nn.Module,
@@ -101,27 +121,146 @@ class LTXPipelineRuntime(
SupportsComponentDiscovery,
DiffusionPipelineProfilerMixin,
):
"""Shared Omni runtime for explicitly composed LTX denoise phases."""
"""Shared Omni runtime for recipe-driven LTX denoise phases."""
pipeline_kind: ClassVar[str] = "one_stage"
component_profile: ClassVar[LTXComponentProfile]
guidance_strategy: ClassVar[LTXGuidanceStrategy]
one_stage_recipe: ClassVar[LTXOneStageRecipe]
pipeline_recipe: ClassVar[LTXPipelineRecipe]
guidance_executor: ClassVar[LTXGuidanceExecutor] = LTX_GUIDANCE_EXECUTOR
supports_request_batch = False
connector_batches_cfg = False
distributed_video_decode = True
support_image_input = False
dummy_run_num_frames = 2
dummy_run_num_frames = 1
preserve_sp_padded_audio_duration = False
reports_stage_durations = False
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
del prefix
self.model_version = detect_ltx_model_version(od_config.model)
self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version)
self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version)
self._dit_modules = list(self.component_profile.dit_modules)
self._encoder_modules = list(self.component_profile.encoder_modules)
self._vae_modules = list(self.component_profile.vae_modules)
self._resident_modules = list(self.component_profile.resident_modules)
if self.model_version == "2.3":
self.preserve_sp_padded_audio_duration = True
self.reports_stage_durations = True
super().__init__()
self._guidance_plan = LTXGuidancePlan.build(self.pipeline_recipe.request_guidance)
initialize_pipeline_components(self, od_config)
self.setup_diffusion_pipeline_profiler(
enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
)
def _forward_request(
self,
req: DiffusionRequestBatch,
*,
image: Any | 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,
guidance_scale: float | None = None,
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",
max_sequence_length: int | None = None,
) -> DiffusionOutput | 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,
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)
request_sigmas = self._resolve_request_sigmas(req, sigmas)
validate_pipeline_request(
request_inputs,
pipeline_recipe=self.pipeline_recipe,
vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
pipeline_name=self.__class__.__name__,
request_sigmas=request_sigmas,
)
return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image)
def _run_recipe(
self,
req: DiffusionRequestBatch,
request_inputs: LTXRequestInputs,
*,
request_sigmas: list[float] | None,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
"""Execute one- and multi-phase recipes through the same control flow."""
phase_results: list[LTXPhaseResult] = []
prompt_context = None
for phase_recipe in self.pipeline_recipe.phases:
self._enter_phase(phase_recipe)
phase_inputs = self._build_phase_inputs(
request_inputs,
phase_recipe,
phase_results[-1] if phase_results else None,
)
phase_result = self.run_phase(
req,
phase_inputs,
noise_scale=phase_recipe.noise_scale,
sigmas=list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else request_sigmas,
timesteps=None,
attention_kwargs=None,
phase_recipe=phase_recipe,
image=image,
prompt_context=prompt_context,
)
phase_results.append(phase_result)
prompt_context = phase_result.forward_context.prompt_context
final_context = phase_results[-1].forward_context
output_phase = LTXPhaseResult(
forward_context=final_context,
video=phase_results[self.pipeline_recipe.video_output_phase].video,
audio=phase_results[self.pipeline_recipe.audio_output_phase].audio,
)
return self.decode_phase(output_phase)
def _enter_phase(self, phase: LTXPhaseRecipe) -> None:
"""Hook for a future phase-weight strategy."""
del phase
def prepare_latents(
self,
batch_size: int = 1,
@@ -160,6 +299,7 @@ class LTXPipelineRuntime(
device: torch.device | None = None,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
latents_normalized: bool = False,
) -> tuple[torch.Tensor, int, int]:
return latent_ops.prepare_audio_latents(
self,
@@ -172,19 +312,16 @@ class LTXPipelineRuntime(
device,
generator,
latents,
latents_normalized,
)
@property
def guidance_scale(self):
return self._guidance_scale
@property
def guidance_rescale(self):
return self._guidance_rescale
return self._guidance_plan.spec.video.cfg_scale
@property
def do_classifier_free_guidance(self):
return self._guidance_scale is not None and self._guidance_scale > 1.0
return self._guidance_plan.spec.do_cfg
@property
def interrupt(self):
@@ -212,21 +349,29 @@ class LTXPipelineRuntime(
):
if kwargs is not None:
context = {**kwargs, **context}
return self.guidance_strategy.combine_cfg_noise(
self,
positive_noise_pred,
negative_noise_pred,
true_cfg_scale,
cfg_normalize,
context,
del cfg_normalize
required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma")
if any(context.get(name) is None for name in required):
raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.")
video_pos, audio_pos = positive_noise_pred
video_neg, audio_neg = negative_noise_pred
return (
self.guidance_executor.combine_cfg_velocity(
context["video_latents"],
video_pos,
video_neg,
context["video_sigma"],
true_cfg_scale,
),
self.guidance_executor.combine_cfg_velocity(
context["audio_latents"],
audio_pos,
audio_neg,
context["audio_sigma"],
true_cfg_scale,
),
)
def predict_noise_with_parallel_cfg(self, *args, **kwargs):
predict_parallel_cfg = getattr(self.guidance_strategy, "predict_parallel_cfg", None)
if predict_parallel_cfg is None:
raise NotImplementedError("The selected LTX guidance strategy does not implement parallel CFG.")
return predict_parallel_cfg(self, *args, **kwargs)
def _synchronize_cfg_parallel_step_output(
self,
latents: tuple[torch.Tensor, torch.Tensor],
@@ -238,23 +383,35 @@ class LTXPipelineRuntime(
# CUDA async execution otherwise permits numerical drift to accumulate
# across CFG-parallel denoise steps.
latents = tuple(tensor.contiguous() for tensor in latents)
device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
if device is not None:
torch.cuda.current_stream(device).synchronize()
current_omni_platform.synchronize()
return latents
def _setup_forward_runtime(
self,
req: DiffusionRequestBatch,
request_inputs: LTXRequestInputs,
attention_kwargs: dict[str, Any] | None,
) -> bool:
self._guidance_scale = request_inputs.guidance_scale
self._guidance_rescale = request_inputs.guidance_rescale
self._guidance_plan = LTXGuidancePlan.build(request_inputs.guidance)
del attention_kwargs
self._interrupt = False
cfg_world_size = get_classifier_free_guidance_world_size()
if (
req.is_dummy_run()
and cfg_world_size == 2
and self.do_classifier_free_guidance
and not self._guidance_plan.cfg_parallel_compatible
):
# The single-device warmup follows the full recipe. CFG2 executes
# the same per-modality CFG scales without STG/modality passes.
self._guidance_plan = LTXGuidancePlan.build(
LTXGuidanceSpec(
video=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.video.cfg_scale),
audio=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.audio.cfg_scale),
)
)
if self.do_classifier_free_guidance:
self.guidance_strategy.validate_cfg_world_size(cfg_world_size)
self.guidance_executor.validate_cfg_world_size(self._guidance_plan, cfg_world_size)
return self.do_classifier_free_guidance and cfg_world_size > 1
def _check_forward_inputs(
@@ -325,11 +482,8 @@ class LTXPipelineRuntime(
vae_decode_needs_all_ranks = False
is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None)
if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled):
try:
# Distributed tiled decode is collective, so every rank must enter it.
vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled())
except Exception:
pass
# Distributed tiled decode is collective, so every rank must enter it.
vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled())
should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks
if should_decode_video:
@@ -361,16 +515,16 @@ class LTXPipelineRuntime(
image: Any | None = None,
) -> tuple[torch.Tensor, torch.Tensor | None]:
latents = self.prepare_latents(
prompt_context.batch_size * request_inputs.num_videos_per_prompt,
self.transformer.config.in_channels,
request_inputs.height,
request_inputs.width,
request_inputs.num_frames,
noise_scale,
torch.float32,
device,
request_inputs.generator,
request_inputs.latents,
batch_size=prompt_context.batch_size * request_inputs.num_videos_per_prompt,
num_channels_latents=self.transformer.config.in_channels,
height=request_inputs.height,
width=request_inputs.width,
num_frames=request_inputs.num_frames,
noise_scale=noise_scale,
dtype=prompt_context.positive_connector_prompt_embeds.dtype,
device=device,
generator=request_inputs.generator,
latents=request_inputs.latents,
)
return latents, None
@@ -416,10 +570,11 @@ class LTXPipelineRuntime(
audio_latent_length=audio_num_frames,
num_mel_bins=num_mel_bins,
noise_scale=noise_scale,
dtype=torch.float32,
dtype=prompt_context.positive_connector_audio_prompt_embeds.dtype,
device=device,
generator=request_inputs.generator,
latents=request_inputs.audio_latents,
latents_normalized=request_inputs.audio_latents_normalized,
)
return audio_latents, original_num_frames, padded_num_frames, latent_mel_bins
@@ -439,29 +594,37 @@ class LTXPipelineRuntime(
padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size))
return requested_length if provided_length in {requested_length, padded_length} else provided_length
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_denoise_context_for_cfg(
self,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> LTXDenoiseContext:
return self.guidance_strategy.prepare_denoise_context(self, forward_ctx, denoise_ctx)
return self.guidance_executor.prepare_denoise_context(
self._guidance_plan,
forward_ctx.cfg_parallel_ready,
denoise_ctx,
)
def _denoise_timestep_kwargs(
self,
ts: torch.Tensor,
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
*,
video_token_count: int,
audio_token_count: int,
) -> dict[str, torch.Tensor]:
return self.guidance_strategy.timestep_kwargs(ts, forward_ctx, denoise_ctx)
del forward_ctx, denoise_ctx
return self.guidance_executor.timestep_kwargs(ts, video_token_count, audio_token_count)
def _video_guidance_model_sigma(
self,
sigma: torch.Tensor,
denoise_ctx: LTXDenoiseContext,
) -> torch.Tensor:
"""Return the video timestep used to convert velocity predictions to x0."""
del denoise_ctx
return sigma
def _build_transformer_kwargs(
self,
@@ -475,6 +638,7 @@ class LTXPipelineRuntime(
encoder_attention_mask: torch.Tensor | None,
audio_encoder_attention_mask: torch.Tensor | None,
ts: torch.Tensor,
attention_kwargs: dict[str, Any] | None = None,
) -> dict[str, Any]:
return build_transformer_kwargs(
self,
@@ -487,6 +651,7 @@ class LTXPipelineRuntime(
encoder_attention_mask=encoder_attention_mask,
audio_encoder_attention_mask=audio_encoder_attention_mask,
ts=ts,
attention_kwargs=attention_kwargs,
)
def _predict_noise_for_step(
@@ -497,8 +662,9 @@ class LTXPipelineRuntime(
forward_ctx: LTXForwardContext,
denoise_ctx: LTXDenoiseContext,
) -> tuple[torch.Tensor, torch.Tensor]:
return self.guidance_strategy.predict_noise(
return self.guidance_executor.predict_noise(
self,
self._guidance_plan,
index,
timestep,
state,
@@ -562,7 +728,6 @@ class LTXPipelineRuntime(
)
audio_latents = latent_ops.unpack_audio_latents(
audio_latents,
forward_ctx.original_audio_num_frames,
num_mel_bins=forward_ctx.latent_mel_bins,
)
return latents, audio_latents
@@ -576,6 +741,7 @@ class LTXPipelineRuntime(
sigmas: list[float] | None,
timesteps: list[int] | None,
attention_kwargs: dict[str, Any] | None,
phase_recipe: LTXPhaseRecipe,
image: Any | None = None,
prompt_context: LTXPromptContext | None = None,
) -> LTXPhaseResult:
@@ -588,10 +754,84 @@ class LTXPipelineRuntime(
sigmas=sigmas,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
phase_recipe=phase_recipe,
image=image,
prompt_context=prompt_context,
)
def _build_phase_inputs(
self,
request_inputs: LTXRequestInputs,
phase: LTXPhaseRecipe,
previous_phase: LTXPhaseResult | None,
) -> LTXRequestInputs:
"""Resolve one phase from immutable request inputs and prior AV state."""
divisor = phase.spatial_downscale
if request_inputs.height % divisor != 0 or request_inputs.width % divisor != 0:
raise ValueError(
f"LTX phase {phase.name!r} cannot scale resolution "
f"{request_inputs.width}x{request_inputs.height} by {divisor}."
)
height = request_inputs.height // divisor
width = request_inputs.width // divisor
latents = request_inputs.latents
audio_latents = request_inputs.audio_latents
audio_latents_normalized = request_inputs.audio_latents_normalized
decode_timestep = request_inputs.decode_timestep
decode_noise_scale = request_inputs.decode_noise_scale
if phase.input_transform == "spatial_upsample":
if previous_phase is None:
raise ValueError(f"LTX phase {phase.name!r} requires a previous phase to upsample.")
if previous_phase.video.ndim != 5:
raise ValueError(f"LTX spatial upsampling expects a 5D video latent, got {previous_phase.video.shape}.")
latents = self._spatial_upsample_phase(previous_phase.video)
if previous_phase.audio_for_next_phase is not None:
audio_latents = previous_phase.audio_for_next_phase
audio_latents_normalized = True
else:
audio_latents = previous_phase.audio
audio_latents_normalized = False
decode_timestep = 0.0
decode_noise_scale = None
expected_shape = latent_ops.resolve_video_latent_shape(
height,
width,
request_inputs.num_frames,
vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
)
if latents.shape[2:] != expected_shape:
raise ValueError(
f"LTX phase {phase.name!r} upsampler produced {tuple(latents.shape[2:])}, expected "
f"{expected_shape} for {width}x{height}."
)
elif phase.input_transform != "initial":
raise ValueError(f"Unsupported LTX phase input transform: {phase.input_transform!r}.")
guidance = request_inputs.guidance if phase.allow_guidance_override else phase.guidance
return replace(
request_inputs,
height=height,
width=width,
num_inference_steps=phase.num_inference_steps or request_inputs.num_inference_steps,
guidance=guidance,
latents=latents,
audio_latents=audio_latents,
audio_latents_normalized=audio_latents_normalized,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
)
def _spatial_upsample_phase(self, latents: torch.Tensor) -> torch.Tensor:
"""Apply the recipe-managed spatial upsampler to unpacked latents."""
latent_upsampler = getattr(self, "latent_upsampler", None)
if latent_upsampler is None:
raise RuntimeError("This LTX pipeline recipe requires a spatial latent upsampler.")
dtype = getattr(latent_upsampler, "dtype", latents.dtype)
return latent_upsampler(latents.to(device=self.device, dtype=dtype))
def decode_phase(self, phase: LTXPhaseResult) -> DiffusionOutput | list[DiffusionOutput]:
"""Decode one completed phase and restore per-request outputs."""
forward_ctx = phase.forward_context
vllm_omni/diffusion/models/ltx2/ltx2_transformer.py
+106 -39
View File
@@ -57,6 +57,9 @@ _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters
def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> nn.Module:
"""Bridge diffusers' RMSNorm API onto vLLM's `has_weight` variant."""
if not elementwise_affine:
return nn.RMSNorm(hidden_size, eps=eps, elementwise_affine=False)
kwargs: dict[str, Any] = {"eps": eps}
if "elementwise_affine" in _RMSNORM_INIT_PARAMS:
kwargs["elementwise_affine"] = elementwise_affine
@@ -79,10 +82,14 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
cos, sin = freqs
if cos.dtype != x.dtype:
cos = cos.to(dtype=x.dtype)
if sin.dtype != x.dtype:
sin = sin.to(dtype=x.dtype)
# Concrete pair count instead of -1 keeps SDPA shape static under torch.compile.
x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1) # [B, S, C // 2]
x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
out = x * cos + x_rotated * sin
return out
@@ -94,6 +101,10 @@ def apply_split_rotary_emb(
) -> torch.Tensor:
# `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile.
cos, sin = freqs
if cos.dtype != x.dtype:
cos = cos.to(dtype=x.dtype)
if sin.dtype != x.dtype:
sin = sin.to(dtype=x.dtype)
x_dtype = x.dtype
needs_reshape = False
@@ -112,7 +123,7 @@ def apply_split_rotary_emb(
r = last // 2
# (..., 2, r)
split_x = x.reshape(*x.shape[:-1], 2, r).float() # Explicitly upcast to float
split_x = x.reshape(*x.shape[:-1], 2, r)
first_x = split_x[..., :1, :] # (..., 1, r)
second_x = split_x[..., 1:, :] # (..., 1, r)
@@ -327,6 +338,23 @@ class TensorParallelRMSNorm(nn.Module):
return out.to(dtype=x_dtype)
def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor:
"""Convert an additive/expanded LTX mask to a 2D padding mask."""
if attention_mask.ndim > 2:
if attention_mask.is_floating_point():
valid = attention_mask >= 0
else:
valid = attention_mask.to(torch.bool)
batch_size = valid.shape[0]
key_length = valid.shape[-1]
attention_mask = valid.reshape(batch_size, -1, key_length).all(dim=1)
if attention_mask.is_floating_point():
attention_mask = attention_mask >= 0
if attention_mask.dtype != torch.bool:
attention_mask = attention_mask.to(torch.bool)
return attention_mask
class LTX2AudioVideoAttnProcessor:
r"""
Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model.
@@ -344,24 +372,6 @@ class LTX2AudioVideoAttnProcessor:
"Please upgrade your PyTorch installation."
)
@staticmethod
def _to_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor:
# Convert additive/expanded masks into a 2D padding mask for flash-attn.
if attention_mask.ndim > 2:
if attention_mask.is_floating_point():
valid = attention_mask >= 0
else:
valid = attention_mask.to(torch.bool)
b = valid.shape[0]
key_len = valid.shape[-1]
valid = valid.reshape(b, -1, key_len).all(dim=1)
attention_mask = valid
if attention_mask.is_floating_point():
attention_mask = attention_mask >= 0
if attention_mask.dtype != torch.bool:
attention_mask = attention_mask.to(torch.bool)
return attention_mask
@staticmethod
def _is_sp_enabled() -> bool:
if not is_forward_context_available():
@@ -390,12 +400,12 @@ class LTX2AudioVideoAttnProcessor:
# For cross-attention, encoder sequence length != query length, so drop the mask.
if encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]:
return None
return self._to_padding_mask(attention_mask)
return to_ltx_padding_mask(attention_mask)
attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
if attn.attn.attn_backend.get_name().upper() == "FLASH_ATTN":
attention_mask = self._to_padding_mask(attention_mask)
attention_mask = to_ltx_padding_mask(attention_mask)
return attention_mask
@staticmethod
@@ -467,6 +477,7 @@ class LTX2AudioVideoAttnProcessor:
attention_mask: torch.Tensor | None = None,
query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
perturbation_mask: torch.Tensor | None = None,
) -> torch.Tensor:
is_self_attention = encoder_hidden_states is None
batch_size, sequence_length, _ = hidden_states.shape if is_self_attention else encoder_hidden_states.shape
@@ -509,17 +520,21 @@ class LTX2AudioVideoAttnProcessor:
elif attn.rope_type == "split":
query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim)
key = apply_split_rotary_emb(
key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, head_dim=attn.head_dim
key,
key_rotary_emb if key_rotary_emb is not None else query_rotary_emb,
head_dim=attn.head_dim,
)
query = query.unflatten(2, (attn.heads, attn.head_dim))
key = key.unflatten(2, (attn.heads, attn.head_dim))
value = value.unflatten(2, (attn.heads, attn.head_dim))
value_heads = value.unflatten(2, (attn.heads, attn.head_dim))
attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
hidden_states = attn.attn(query, key, value, attn_metadata)
hidden_states = attn.attn(query, key, value_heads, attn_metadata)
hidden_states = hidden_states.flatten(2, 3)
hidden_states = hidden_states.to(query.dtype)
if perturbation_mask is not None:
hidden_states = hidden_states * perturbation_mask + value * (1 - perturbation_mask)
# LTX-2.3: per-head gated attention
if attn.to_gate_logits is not None:
@@ -560,6 +575,7 @@ class LTX2Attention(torch.nn.Module):
norm_elementwise_affine: bool = True,
rope_type: str = "interleaved",
apply_gated_attention: bool = False,
pack_qkv: bool = True,
processor=None,
quant_config: "QuantizationConfig | None" = None,
prefix: str = "",
@@ -591,7 +607,7 @@ class LTX2Attention(torch.nn.Module):
self.to_q = None
self.to_k = None
self.to_v = None
if cross_attention_dim is None:
if cross_attention_dim is None and pack_qkv:
self.to_qkv = QKVParallelLinear(
hidden_size=query_dim,
head_size=self.head_dim,
@@ -847,6 +863,9 @@ class LTX2VideoTransformerBlock(nn.Module):
qk_norm=qk_norm,
rope_type=rope_type,
apply_gated_attention=audio_gated_attn,
# The official checkpoint evaluates audio Q/K/V separately. Fusing
# the bf16 GEMMs changes denoise numerics enough to accumulate.
pack_qkv=False,
quant_config=quant_config,
prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1",
)
@@ -992,8 +1011,10 @@ class LTX2VideoTransformerBlock(nn.Module):
v2a_cross_attention_mask: torch.Tensor | None = None,
use_a2v_cross_attention: bool = True,
use_v2a_cross_attention: bool = True,
perturbation_mask: torch.Tensor | None = None,
all_perturbed: bool | None = None,
video_self_attention_perturbation_mask: torch.Tensor | None = None,
audio_self_attention_perturbation_mask: torch.Tensor | None = None,
a2v_cross_attention_perturbation_mask: torch.Tensor | None = None,
v2a_cross_attention_perturbation_mask: torch.Tensor | None = None,
) -> torch.Tensor:
batch_size = hidden_states.size(0)
@@ -1012,6 +1033,7 @@ class LTX2VideoTransformerBlock(nn.Module):
encoder_hidden_states=None,
query_rotary_emb=video_rotary_emb,
attention_mask=self_attention_mask,
perturbation_mask=video_self_attention_perturbation_mask,
)
hidden_states = hidden_states + attn_hidden_states * gate_msa
@@ -1031,6 +1053,7 @@ class LTX2VideoTransformerBlock(nn.Module):
encoder_hidden_states=None,
query_rotary_emb=audio_rotary_emb,
attention_mask=audio_self_attention_mask,
perturbation_mask=audio_self_attention_perturbation_mask,
)
audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
@@ -1123,6 +1146,8 @@ class LTX2VideoTransformerBlock(nn.Module):
key_rotary_emb=ca_audio_rotary_emb,
attention_mask=a2v_cross_attention_mask,
)
if a2v_cross_attention_perturbation_mask is not None:
a2v_attn_hidden_states = a2v_attn_hidden_states * a2v_cross_attention_perturbation_mask
hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
# 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video
@@ -1141,6 +1166,8 @@ class LTX2VideoTransformerBlock(nn.Module):
key_rotary_emb=ca_video_rotary_emb,
attention_mask=v2a_cross_attention_mask,
)
if v2a_cross_attention_perturbation_mask is not None:
v2a_attn_hidden_states = v2a_attn_hidden_states * v2a_cross_attention_perturbation_mask
audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
# 4. Feedforward
@@ -1345,7 +1372,10 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
return self.prepare_audio_coords(*args, **kwargs)
def forward(
self, coords: torch.Tensor, device: str | torch.device | None = None
self,
coords: torch.Tensor,
device: str | torch.device | None = None,
out_dtype: torch.dtype | None = None,
) -> tuple[torch.Tensor, torch.Tensor]:
device = device or coords.device
@@ -1421,6 +1451,9 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
cos_freqs = torch.swapaxes(cos_freq, 1, 2) # (B,H,T,D//2)
sin_freqs = torch.swapaxes(sin_freq, 1, 2) # (B,H,T,D//2)
if out_dtype is not None:
cos_freqs = cos_freqs.to(out_dtype)
sin_freqs = sin_freqs.to(out_dtype)
return cos_freqs, sin_freqs
@@ -1488,8 +1521,11 @@ class LTX2VideoTransformer3DModel(nn.Module):
# Shard prompt embeds across sequence
"encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
"audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
# Shard video timestep when provided as (B, seq_len)
# Keep per-token video/audio timestep embeddings aligned with
# their corresponding latent sequence shards. Official LTX
# multi-guidance supplies both tensors as (B, seq_len).
"timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False),
"audio_timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False),
},
"rope": {
0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
@@ -1856,6 +1892,12 @@ class LTX2VideoTransformer3DModel(nn.Module):
# Determine timestep for audio.
audio_timestep = audio_timestep if audio_timestep is not None else timestep
def batch_step(step: torch.Tensor) -> torch.Tensor:
return step[:, 0] if step.ndim > 1 else step
sigma = sigma if sigma is not None else batch_step(timestep)
audio_sigma = audio_sigma if audio_sigma is not None else batch_step(audio_timestep)
# convert encoder_attention_mask to a bias the same way we do for attention_mask
if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
@@ -1877,12 +1919,22 @@ class LTX2VideoTransformer3DModel(nn.Module):
batch_size, audio_num_frames, audio_hidden_states.device
)
video_rotary_emb = self.rope(video_coords, device=hidden_states.device)
audio_rotary_emb = self.audio_rope(audio_coords, device=audio_hidden_states.device)
video_rotary_emb = self.rope(video_coords, device=hidden_states.device, out_dtype=hidden_states.dtype)
audio_rotary_emb = self.audio_rope(
audio_coords,
device=audio_hidden_states.device,
out_dtype=audio_hidden_states.dtype,
)
video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :], device=hidden_states.device)
video_cross_attn_rotary_emb = self.cross_attn_rope(
video_coords[:, 0:1, :],
device=hidden_states.device,
out_dtype=hidden_states.dtype,
)
audio_cross_attn_rotary_emb = self.cross_attn_audio_rope(
audio_coords[:, 0:1, :], device=audio_hidden_states.device
audio_coords[:, 0:1, :],
device=audio_hidden_states.device,
out_dtype=audio_hidden_states.dtype,
)
# 2. Patchify input projections
@@ -1920,7 +1972,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
hidden_dtype=hidden_states.dtype,
)
video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate(
timestep.flatten() * timestep_cross_attn_gate_scale_factor,
audio_sigma.flatten() * timestep_cross_attn_gate_scale_factor,
batch_size=batch_size,
hidden_dtype=hidden_states.dtype,
)
@@ -1935,7 +1987,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
hidden_dtype=audio_hidden_states.dtype,
)
audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate(
audio_timestep.flatten() * timestep_cross_attn_gate_scale_factor,
sigma.flatten() * timestep_cross_attn_gate_scale_factor,
batch_size=batch_size,
hidden_dtype=audio_hidden_states.dtype,
)
@@ -1945,7 +1997,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
# 3.3. LTX-2.3: Compute prompt modulation from sigma
audio_sigma = audio_sigma if audio_sigma is not None else sigma
if self.prompt_modulation and sigma is not None:
temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype)
temb_prompt_audio, _ = self.audio_prompt_adaln(
@@ -1967,8 +2018,17 @@ class LTX2VideoTransformer3DModel(nn.Module):
audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
perturbation_kwargs = (attention_kwargs or {}).get("ltx_perturbation_kwargs", {})
def perturbation_mask_for(name: str, block_idx: int) -> torch.Tensor | None:
mask = perturbation_kwargs.get(f"{name}_mask")
blocks = perturbation_kwargs.get(f"{name}_blocks")
if mask is None or (blocks is not None and block_idx not in blocks):
return None
return mask
# 5. Run transformer blocks
for block in self.transformer_blocks:
for block_idx, block in enumerate(self.transformer_blocks):
block_kwargs = {
"temb": temb,
"temb_audio": temb_audio,
@@ -1984,6 +2044,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
"ca_audio_rotary_emb": audio_cross_attn_rotary_emb,
"encoder_attention_mask": encoder_attention_mask,
"audio_encoder_attention_mask": audio_encoder_attention_mask,
"video_self_attention_perturbation_mask": perturbation_mask_for("video_self_attention", block_idx),
"audio_self_attention_perturbation_mask": perturbation_mask_for("audio_self_attention", block_idx),
"a2v_cross_attention_perturbation_mask": perturbation_mask_for("a2v_cross_attention", block_idx),
"v2a_cross_attention_perturbation_mask": perturbation_mask_for("v2a_cross_attention", block_idx),
}
# NOTE: We intentionally pass the first 4 args here positionally,
# because CacheDiT expects the first 2 args to be hidden_states
@@ -2066,7 +2130,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
for param_name, weight_name, shard_id in stacked_params_mapping:
if weight_name not in name:
continue
name = name.replace(weight_name, param_name)
packed_name = name.replace(weight_name, param_name)
if packed_name not in params_dict:
continue
name = packed_name
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py
+22 -216
View File
@@ -5,250 +5,56 @@
from __future__ import annotations
from typing import Any, ClassVar
from typing import ClassVar
import torch
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.data import OmniDiffusionConfig
from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from .ltx2_components import (
LTX2_COMPONENT_PROFILE,
LTX23_COMPONENT_PROFILE,
LTXComponentProfile,
)
from .ltx2_components import LTX2_COMPONENT_PROFILE
from .ltx2_components import (
get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401
)
from .ltx2_conditioning import LTXI2VConditioningMixin
from .ltx2_guidance import (
LTX_LEGACY_VELOCITY_GUIDANCE,
LTX_OFFICIAL_X0_GUIDANCE,
from .ltx2_recipes import (
LTX2_ONE_STAGE_RECIPE,
LTX_POSITIVE_ONLY_RECIPE,
)
from .ltx2_pipeline_runtime import LTXPipelineRuntime
from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE, LTXOneStageRecipe
from .ltx2_request import LTXRequestInputs
from .ltx2_runtime import LTXRuntime
class LTXOneStagePipeline(LTXPipelineRuntime):
"""Single execution path configured by model-version and task entries."""
class LTX2Pipeline(LTXI2VConditioningMixin, LTXRuntime):
"""LTX-2 family one-stage entry, configured from checkpoint metadata."""
component_profile: ClassVar[LTXComponentProfile]
one_stage_recipe: ClassVar[LTXOneStageRecipe]
supports_request_batch = True
supports_guidance_rescale = False
def _forward_impl(
self,
req: DiffusionRequestBatch,
request_inputs: LTXRequestInputs,
*,
noise_scale: float,
sigmas: list[float] | None,
timesteps: list[int] | None,
attention_kwargs: dict[str, Any] | None,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
phase = self.run_phase(
req,
request_inputs,
noise_scale=noise_scale,
sigmas=sigmas,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
image=image,
)
return self.decode_phase(phase)
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
prompt: str | list[str] | None = None,
negative_prompt: str | list[str] | None = None,
height: int | None = None,
width: int | None = None,
num_frames: int | None = None,
frame_rate: float | None = None,
num_inference_steps: int | None = None,
sigmas: list[float] | None = None,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
*,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
del return_dict
sigmas = self._resolve_request_sigmas(req, sigmas)
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=self.one_stage_recipe.guidance_scale if guidance_scale is None else guidance_scale,
guidance_rescale=guidance_rescale,
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)
forward_kwargs = {
"noise_scale": noise_scale,
"sigmas": sigmas,
"timesteps": timesteps,
"attention_kwargs": attention_kwargs,
}
if self.support_image_input:
forward_kwargs["image"] = image
return self._forward_impl(
req,
request_inputs,
**forward_kwargs,
)
class LTX2Pipeline(LTXOneStagePipeline):
"""LTX2 one-stage text-to-video entry."""
supports_guidance_rescale = True
pipeline_kind = "one_stage"
unified_text_image_entry = True
component_profile = LTX2_COMPONENT_PROFILE
guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE
one_stage_recipe = LTX2_ONE_STAGE_RECIPE
pipeline_recipe = LTX2_ONE_STAGE_RECIPE
_dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules)
_encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules)
_vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules)
_resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules)
class LTX23Pipeline(LTXOneStagePipeline):
"""LTX2.3 one-stage text-to-video entry."""
connector_batches_cfg = True
preserve_sp_padded_audio_duration = True
reports_stage_durations = True
component_profile = LTX23_COMPONENT_PROFILE
guidance_strategy = LTX_OFFICIAL_X0_GUIDANCE
one_stage_recipe = LTX23_ONE_STAGE_RECIPE
_dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules)
_encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules)
_vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules)
_resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules)
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
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,
*,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
"""Preserve the pre-refactor LTX2.3 positional argument order."""
return super().forward(
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,
sigmas=sigmas,
timesteps=timesteps,
guidance_scale=guidance_scale,
guidance_rescale=0.0,
noise_scale=noise_scale,
num_videos_per_prompt=num_videos_per_prompt,
generator=generator,
latents=latents,
audio_latents=audio_latents,
prompt_embeds=prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
prompt_attention_mask=prompt_attention_mask,
negative_prompt_attention_mask=negative_prompt_attention_mask,
decode_timestep=decode_timestep,
decode_noise_scale=decode_noise_scale,
output_type=output_type,
return_dict=return_dict,
attention_kwargs=attention_kwargs,
max_sequence_length=max_sequence_length,
image=image,
)
class LTX2ImageToVideoPipeline(LTXI2VConditioningMixin, LTX2Pipeline):
"""LTX2 one-stage image-to-video entry."""
class LTX23ImageToVideoPipeline(LTXI2VConditioningMixin, LTX23Pipeline):
"""LTX2.3 one-stage image-to-video entry."""
supports_request_batch = True
class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline):
"""LTX2 T2V entry for FastGen DMD2-distilled models."""
pipeline_kind = "dmd2"
support_image_input = False
pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
self.__init_dmd2__()
class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline):
class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline):
"""LTX2 I2V entry for FastGen DMD2-distilled models."""
pipeline_kind = "dmd2"
unified_text_image_entry = False
pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
super().__init__(od_config=od_config, prefix=prefix)
self.__init_dmd2__()
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py
-262
View File
@@ -1,262 +0,0 @@
import os
import torch
import torch.nn as nn
from diffusers import AutoencoderKLLTX2Video
from diffusers.image_processor import PipelineImageInput
from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
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.logger import init_logger
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.distributed.utils import get_local_device
logger = init_logger(__name__)
class LTX2LatentUpsamplePipeline(nn.Module):
def __init__(
self,
od_config: OmniDiffusionConfig,
vae: AutoencoderKLLTX2Video,
latent_upsampler: LTX2LatentUpsamplerModel = None,
) -> None:
super().__init__()
if vae is None:
raise ValueError("vae must be provided")
self.vae = vae
self.device = get_local_device()
model = od_config.model
local_files_only = os.path.exists(model)
if latent_upsampler is None:
# Use cpu context to create latent upsampler. The code k[:, None] @ k[None, :] in
# diffuser's BlurDownsample is not supported on GPU as k is type of torch.Int64
with torch.device("cpu"):
latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
model,
subfolder="latent_upsampler",
torch_dtype=torch.bfloat16,
local_files_only=local_files_only,
).to(self.device)
self.latent_upsampler = latent_upsampler
self.vae_spatial_compression_ratio = (
self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32
)
self.vae_temporal_compression_ratio = (
self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8
)
self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio)
def prepare_latents(
self,
video: torch.Tensor | None = None,
batch_size: int = 1,
num_frames: int = 121,
height: int = 512,
width: int = 768,
spatial_patch_size: int = 1,
temporal_patch_size: int = 1,
dtype: torch.dtype | None = None,
device: torch.device | None = None,
generator: torch.Generator | None = None,
latents: torch.Tensor | None = None,
) -> torch.Tensor:
if latents is not None:
if latents.ndim == 3:
# Convert token seq [B, S, D] to latent video [B, C, F, H, W]
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
latents = self._unpack_latents(
latents, latent_num_frames, latent_height, latent_width, spatial_patch_size, temporal_patch_size
)
return latents.to(device=device, dtype=dtype)
video = video.to(device=device, dtype=self.vae.dtype)
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 batch"
f" size of {batch_size}. Make sure the batch size matches the length of the generators."
)
init_latents = [
retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator[i]) for i in range(batch_size)
]
else:
init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
init_latents = torch.cat(init_latents, dim=0).to(dtype)
# NOTE: latent upsampler operates on the unnormalized latents, so don't normalize here
# init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std)
return init_latents
def adain_filter_latent(self, latents: torch.Tensor, reference_latents: torch.Tensor, factor: float = 1.0):
result = latents.clone()
for i in range(latents.size(0)):
for c in range(latents.size(1)):
r_sd, r_mean = torch.std_mean(reference_latents[i, c], dim=None) # index by original dim order
i_sd, i_mean = torch.std_mean(result[i, c], dim=None)
result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean
result = torch.lerp(latents, result, factor)
return result
def tone_map_latents(self, latents: torch.Tensor, compression: float) -> torch.Tensor:
# Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot
scale_factor = compression * 0.75
abs_latents = torch.abs(latents)
# Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0
# When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect
sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0))
scales = 1.0 - 0.8 * scale_factor * sigmoid_term
filtered = latents * scales
return filtered
@staticmethod
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents
def _denormalize_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
) -> torch.Tensor:
# Denormalize latents across the channel dimension [B, C, F, H, W]
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents = latents * latents_std / scaling_factor + latents_mean
return latents
@staticmethod
# Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents
def _unpack_latents(
latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
) -> torch.Tensor:
# Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) # noqa
# are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of
# what happens in the `_pack_latents` method.
batch_size = latents.size(0)
latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
return latents
def check_inputs(self, video, height, width, latents, tone_map_compression_ratio):
if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0:
raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
if video is not None and latents is not None:
raise ValueError("Only one of `video` or `latents` can be provided.")
if video is None and latents is None:
raise ValueError("One of `video` or `latents` has to be provided.")
if not (0 <= tone_map_compression_ratio <= 1):
raise ValueError("`tone_map_compression_ratio` must be in the range [0, 1]")
def forward(
self,
video: list[PipelineImageInput] | None = None,
height: int = 512,
width: int = 768,
num_frames: int = 121,
spatial_patch_size: int = 1,
temporal_patch_size: int = 1,
latents: torch.Tensor | None = None,
latents_normalized: bool = False,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
adain_factor: float = 0.0,
tone_map_compression_ratio: float = 0.0,
generator: torch.Generator | list[torch.Generator] | None = None,
output_type: str | None = "pil",
return_dict: bool = True,
):
self.check_inputs(
video=video,
height=height,
width=width,
latents=latents,
tone_map_compression_ratio=tone_map_compression_ratio,
)
if video is not None:
# Batched video input is not yet tested/supported. TODO: take a look later
batch_size = 1
else:
batch_size = latents.shape[0]
device = self.device
if video is not None:
num_frames = len(video)
if num_frames % self.vae_temporal_compression_ratio != 1:
num_frames = num_frames // self.vae_temporal_compression_ratio * self.vae_temporal_compression_ratio + 1
video = video[:num_frames]
logger.warning(
f"Video length expected to be of the form `k * {self.vae_temporal_compression_ratio} + 1` but is {len(video)}. Truncating to {num_frames} frames." # noqa
)
video = self.video_processor.preprocess_video(video, height=height, width=width)
video = video.to(device=device, dtype=torch.float32)
latents_supplied = latents is not None
latents = self.prepare_latents(
video=video,
batch_size=batch_size,
num_frames=num_frames,
height=height,
width=width,
spatial_patch_size=spatial_patch_size,
temporal_patch_size=temporal_patch_size,
dtype=torch.float32,
device=device,
generator=generator,
latents=latents,
)
if latents_supplied and latents_normalized:
latents = self._denormalize_latents(
latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
)
latents = latents.to(self.latent_upsampler.dtype)
latents_upsampled = self.latent_upsampler(latents)
if adain_factor > 0.0:
latents = self.adain_filter_latent(latents_upsampled, latents, adain_factor)
else:
latents = latents_upsampled
if tone_map_compression_ratio > 0.0:
latents = self.tone_map_latents(latents, tone_map_compression_ratio)
if output_type == "latent":
video = latents
else:
if not self.vae.config.timestep_conditioning:
timestep = None
else:
noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
if not isinstance(decode_timestep, list):
decode_timestep = [decode_timestep] * batch_size
if decode_noise_scale is None:
decode_noise_scale = decode_timestep
elif not isinstance(decode_noise_scale, list):
decode_noise_scale = [decode_noise_scale] * batch_size
timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
:, None, None, None, None
]
latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
video = self.vae.decode(latents, timestep, return_dict=False)[0]
video = self.video_processor.postprocess_video(video, output_type=output_type)
if not return_dict:
return (video,)
return DiffusionOutput(output=(video,))
vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py
+12 -192
View File
@@ -1,207 +1,27 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Two-stage entry points for the LTX model family."""
"""Distilled two-stage entry points for the LTX model family."""
from __future__ import annotations
from vllm_omni.diffusion.data import OmniDiffusionConfig
import os
from dataclasses import replace
from typing import Any, ClassVar
import torch
from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
from .ltx2_components import (
LTX2_COMPONENT_PROFILE,
)
from .ltx2_components import LTX2_DISTILLED_COMPONENT_PROFILE
from .ltx2_components import (
get_ltx2_post_process_func as get_ltx2_post_process_func, # noqa: F401
)
from .ltx2_conditioning import LTXI2VConditioningMixin
from .ltx2_guidance import LTX_LEGACY_VELOCITY_GUIDANCE
from .ltx2_pipeline_runtime import LTXPipelineRuntime
from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE
from .ltx2_request import LTXRequestInputs
from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
from .ltx2_recipes import LTX2_DISTILLED_TWO_STAGE_RECIPE
from .pipeline_ltx2 import LTX2Pipeline
class LTX2TwoStagesPipeline(LTXPipelineRuntime):
"""Legacy distilled-only LTX2 two-stage compatibility entry."""
class LTX2DistilledPipeline(LTX2Pipeline):
"""Unified LTX-2 distilled two-stage T2V/I2V entry."""
component_profile = LTX2_COMPONENT_PROFILE
guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE
one_stage_recipe = LTX2_ONE_STAGE_RECIPE
pipeline_kind = "distilled_two_stage"
component_profile = LTX2_DISTILLED_COMPONENT_PROFILE
pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE
supports_request_batch = False
support_image_input = False
_dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules)
_encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules)
_vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules)
_resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules)
support_image_input = True
unified_text_image_entry = True
def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
model_path = od_config.model
self.distilled = "distilled" in os.path.basename(os.path.normpath(model_path))
if not self.distilled:
raise NotImplementedError(f"{model_path} is not supported for {self.__class__.__name__}.")
super().__init__(od_config=od_config, prefix=prefix)
self.upsample_pipe = LTX2LatentUpsamplePipeline(vae=self.vae, od_config=od_config)
def _run_two_stage(
self,
req: DiffusionRequestBatch,
request_inputs: LTXRequestInputs,
*,
noise_scale: float,
timesteps: list[int] | None,
attention_kwargs: dict[str, Any] | None,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
stage1 = self.run_phase(
req,
request_inputs,
noise_scale=noise_scale,
sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
image=image,
)
upscaled_video_latent = self.upsample_pipe(
latents=stage1.video,
output_type="latent",
return_dict=False,
)[0]
stage2_inputs = replace(
request_inputs,
num_inference_steps=3,
guidance_scale=1.0,
latents=upscaled_video_latent,
audio_latents=stage1.audio,
decode_timestep=0.0,
decode_noise_scale=None,
output_type="np",
)
stage2 = self.run_phase(
req,
stage2_inputs,
noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
timesteps=None,
attention_kwargs=attention_kwargs,
prompt_context=stage1.forward_context.prompt_context,
)
return self.decode_phase(stage2)
@torch.no_grad()
def forward(
self,
req: DiffusionRequestBatch,
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,
timesteps: list[int] | None = None,
guidance_scale: float = 4.0,
guidance_rescale: float = 0.0,
noise_scale: float = 0.0,
num_videos_per_prompt: int | None = 1,
generator: torch.Generator | list[torch.Generator] | None = None,
latents: torch.Tensor | None = None,
audio_latents: torch.Tensor | None = None,
prompt_embeds: torch.Tensor | None = None,
negative_prompt_embeds: torch.Tensor | None = None,
prompt_attention_mask: torch.Tensor | None = None,
negative_prompt_attention_mask: torch.Tensor | None = None,
decode_timestep: float | list[float] = 0.0,
decode_noise_scale: float | list[float] | None = None,
output_type: str = "np",
return_dict: bool = True,
attention_kwargs: dict[str, Any] | None = None,
max_sequence_length: int | None = None,
*,
image: Any | None = None,
) -> DiffusionOutput | list[DiffusionOutput]:
del return_dict
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=(
getattr(getattr(self, "one_stage_recipe", None), "guidance_scale", 4.0)
if guidance_scale is None
else guidance_scale
),
guidance_rescale=guidance_rescale,
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,
)
if self.support_image_input:
image = self._resolve_request_image(req, image, request_inputs)
return self._run_two_stage(
req,
request_inputs,
noise_scale=noise_scale,
timesteps=timesteps,
attention_kwargs=attention_kwargs,
image=image,
)
class LTX2ImageToVideoTwoStagesPipeline(LTXI2VConditioningMixin, LTX2TwoStagesPipeline):
"""LTX2 two-stage image-to-video entry."""
def forward(
self,
req: DiffusionRequestBatch,
image: Any | 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,
*args: Any,
**kwargs: Any,
) -> DiffusionOutput | list[DiffusionOutput]:
"""Preserve the legacy I2V sigma slot, which was ignored by this pipeline."""
del sigmas
return super().forward(
req,
image,
prompt,
negative_prompt,
height,
width,
num_frames,
frame_rate,
num_inference_steps,
*args,
**kwargs,
)
vllm_omni/diffusion/registry.py
+4 -28
View File
@@ -71,20 +71,10 @@ _DIFFUSION_MODELS = {
"pipeline_ltx2",
"LTX2Pipeline",
),
"LTX2ImageToVideoPipeline": (
"ltx2",
"pipeline_ltx2",
"LTX2ImageToVideoPipeline",
),
"LTX2TwoStagesPipeline": (
"LTX2DistilledPipeline": (
"ltx2",
"pipeline_ltx2_two_stage",
"LTX2TwoStagesPipeline",
),
"LTX2ImageToVideoTwoStagesPipeline": (
"ltx2",
"pipeline_ltx2_two_stage",
"LTX2ImageToVideoTwoStagesPipeline",
"LTX2DistilledPipeline",
),
"LTX2T2VDMD2Pipeline": (
"ltx2",
@@ -96,16 +86,6 @@ _DIFFUSION_MODELS = {
"pipeline_ltx2",
"LTX2I2VDMD2Pipeline",
),
"LTX23Pipeline": (
"ltx2",
"pipeline_ltx2",
"LTX23Pipeline",
),
"LTX23ImageToVideoPipeline": (
"ltx2",
"pipeline_ltx2",
"LTX23ImageToVideoPipeline",
),
"StableAudioPipeline": (
"stable_audio",
"pipeline_stable_audio",
@@ -449,7 +429,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
for attr in transformer_attrs:
if not hasattr(model, attr):
# Some pipeline like LTX2TwoStagesPipeline have recursive
# Some pipelines have recursive
# modules that have the transformer
module = find_module_with_attr(model, attr)
if module is None:
@@ -519,13 +499,9 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
"WanPipeline": "get_wan22_post_process_func",
"WanVACEPipeline": "get_wan22_vace_post_process_func",
"LTX2Pipeline": "get_ltx2_post_process_func",
"LTX2TwoStagesPipeline": "get_ltx2_post_process_func",
"LTX2ImageToVideoPipeline": "get_ltx2_post_process_func",
"LTX2ImageToVideoTwoStagesPipeline": "get_ltx2_post_process_func",
"LTX2DistilledPipeline": "get_ltx2_post_process_func",
"LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
"LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",
"LTX23Pipeline": "get_ltx2_post_process_func",
"LTX23ImageToVideoPipeline": "get_ltx2_post_process_func",
"StableAudioPipeline": "get_stable_audio_post_process_func",
"SoulXSingerPipeline": "get_soulxsinger_post_process_func",
"SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func",
vllm_omni/diffusion/worker/request_batch.py
+2 -1
View File
@@ -8,6 +8,7 @@ from collections.abc import Mapping, Sequence
from dataclasses import dataclass
from typing import Any
import numpy as np
import torch
from vllm_omni.diffusion.data import DiffusionOutput
@@ -20,7 +21,7 @@ def _slice_request_output(value: Any, start: int, stop: int) -> Any:
return tuple(_slice_request_output(item, start, stop) for item in value)
if isinstance(value, list):
return value[start:stop]
if isinstance(value, torch.Tensor):
if isinstance(value, (np.ndarray, torch.Tensor)):
return value[start:stop]
return value
vllm_omni/entrypoints/cli/serve.py
+1 -1
View File
@@ -410,7 +410,7 @@ class OmniServeCommand(CLISubcommand):
dest="model_class_name",
type=str,
default=None,
help="Override the diffusion pipeline class name (e.g. LTX2ImageToVideoPipeline).",
help="Override the diffusion pipeline class name (e.g. LTX2Pipeline).",
)
omni_config_group.add_argument(
"--diffusion-load-format",
vllm_omni/model_extras/ltx2.py
+27
View File
@@ -0,0 +1,27 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Request parameters shared by official LTX guidance pipelines."""
LTX_EXTRA_BODY_PARAMS = frozenset(
{
"video_cfg_scale",
"audio_cfg_scale",
"video_cfg_guidance_scale",
"audio_cfg_guidance_scale",
"video_stg_scale",
"audio_stg_scale",
"video_stg_guidance_scale",
"audio_stg_guidance_scale",
"video_modality_scale",
"audio_modality_scale",
"a2v_guidance_scale",
"v2a_guidance_scale",
"video_rescale_scale",
"audio_rescale_scale",
"video_stg_blocks",
"audio_stg_blocks",
}
)
LTX_EXTRA_OUTPUT_PARAMS = frozenset()
vllm_omni/model_extras/registry.py
+11
View File
@@ -37,6 +37,7 @@ from vllm_omni.model_extras.helios import (
)
from vllm_omni.model_extras.hunyuan_image3 import build_x_to_text_prompt as build_hunyuan_x_to_text_prompt
from vllm_omni.model_extras.lingbot_video import LINGBOT_VIDEO_EXTRA_BODY_PARAMS
from vllm_omni.model_extras.ltx2 import LTX_EXTRA_BODY_PARAMS, LTX_EXTRA_OUTPUT_PARAMS
from vllm_omni.model_extras.magi_human import (
MAGI_HUMAN_EXTRA_BODY_PARAMS,
MAGI_HUMAN_EXTRA_OUTPUT_PARAMS,
@@ -223,6 +224,16 @@ _EXTRA_SPECS: dict[str, dict[str, Any]] = {
"LingBotVideoPipeline": {
"extra_body_params": LINGBOT_VIDEO_EXTRA_BODY_PARAMS,
},
**{
model_class_name: {
"extra_body_params": LTX_EXTRA_BODY_PARAMS,
"extra_output_params": LTX_EXTRA_OUTPUT_PARAMS,
}
for model_class_name in (
"LTX2Pipeline",
"LTX2DistilledPipeline",
)
},
"WanVACEPipeline": {
"extra_body_params": VACE_EXTRA_BODY_PARAMS,
"extra_output_params": VACE_EXTRA_OUTPUT_PARAMS,