diff --git a/.buildkite/cuda/test-nightly.yml b/.buildkite/cuda/test-nightly.yml index 2aed12819aa..a82a4aab577 100644 --- a/.buildkite/cuda/test-nightly.yml +++ b/.buildkite/cuda/test-nightly.yml @@ -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 diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 9f99592fd1d..021df65ce79 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -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` | ✅︎ | ✅︎ | ✅︎ | | diff --git a/docs/serving/videos_api.md b/docs/serving/videos_api.md index 72d79c30427..6a94f9dab89 100644 --- a/docs/serving/videos_api.md +++ b/docs/serving/videos_api.md @@ -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 | diff --git a/docs/user_guide/examples/offline_inference/image_to_video.md b/docs/user_guide/examples/offline_inference/image_to_video.md index b3d76ac46db..db7a31ce599 100644 --- a/docs/user_guide/examples/offline_inference/image_to_video.md +++ b/docs/user_guide/examples/offline_inference/image_to_video.md @@ -4,7 +4,7 @@ Source 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 diff --git a/examples/offline_inference/text_to_video/text_to_video.md b/examples/offline_inference/text_to_video/text_to_video.md index 876b8429b44..0df1f9bd022 100644 --- a/examples/offline_inference/text_to_video/text_to_video.md +++ b/examples/offline_inference/text_to_video/text_to_video.md @@ -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). diff --git a/examples/offline_inference/text_to_video/text_to_video.py b/examples/offline_inference/text_to_video/text_to_video.py index addf7cb7a29..b0c8a2509f1 100644 --- a/examples/offline_inference/text_to_video/text_to_video.py +++ b/examples/offline_inference/text_to_video/text_to_video.py @@ -65,11 +65,42 @@ "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, diff --git a/examples/online_serving/text_to_video/README.md b/examples/online_serving/text_to_video/README.md index 71afd9e7cd0..51512fe0cdf 100644 --- a/examples/online_serving/text_to_video/README.md +++ b/examples/online_serving/text_to_video/README.md @@ -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 -``` - -#### 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 +vllm serve Lightricks/LTX-2 --omni --port 8098 ``` -#### 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 diff --git a/examples/online_serving/text_to_video/run_curl_ltx2.sh b/examples/online_serving/text_to_video/run_curl_ltx2.sh index b82f672eaab..1cf828bba09 100644 --- a/examples/online_serving/text_to_video/run_curl_ltx2.sh +++ b/examples/online_serving/text_to_video/run_curl_ltx2.sh @@ -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" ) diff --git a/examples/online_serving/text_to_video/run_server_ltx2.sh b/examples/online_serving/text_to_video/run_server_ltx2.sh index f4597d3cd28..da9351f7fcb 100644 --- a/examples/online_serving/text_to_video/run_server_ltx2.sh +++ b/examples/online_serving/text_to_video/run_server_ltx2.sh @@ -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[@]}" diff --git a/recipes/LTX/LTX-2.3.md b/recipes/LTX/LTX-2.3.md deleted file mode 100644 index 471ee7ae040..00000000000 --- a/recipes/LTX/LTX-2.3.md +++ /dev/null @@ -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: -- Diffusers-format checkpoint: -- 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. diff --git a/recipes/LTX/LTX-2.md b/recipes/LTX/LTX-2.md index df049f5c625..ca3201dc119 100644 --- a/recipes/LTX/LTX-2.md +++ b/recipes/LTX/LTX-2.md @@ -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`. -## References +## 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=` | -- Model: -- 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) +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. -## Hardware Support +## One-Stage Defaults -## GPU +| 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]` | -### 1x H200 141GB (Text-to-Video) +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`. -#### Environment +## Distilled Two-Stage Defaults -- 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 +| 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 | -#### Command +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 \ - --model-class-name LTX2Pipeline +vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600 ``` -#### Verification +```bash +vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600 +``` + +Start the distilled checkpoint with its two-stage entry: ```bash -curl http://localhost:8000/health +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: -curl -X POST http://localhost:8000/v1/videos \ +```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 "width=512" \ - -F "height=768" \ + -F "size=768x512" \ -F "num_frames=121" \ -F "fps=24" \ - -F "num_inference_steps=40" \ - -F "guidance_scale=4.0" \ - -o ltx2_t2v_h200.mp4 + -F "seed=42" \ + -o ltx_t2v.mp4 ``` -#### Notes +For I2V, provide exactly one initial image: -- 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. +```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 +``` -### 1x H200 141GB (Image-to-Video) +Use `image_reference` for a URL or JSON-safe image reference. Do not provide it +together with `input_reference`. -#### Environment +## Guidance -- 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 +LTX supports independent video/audio CFG, spatio-temporal guidance (STG), +cross-modality guidance, and rescaling: -#### Command +| 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 -vllm serve Lightricks/LTX-2 \ - --omni \ - --model-class-name LTX2ImageToVideoPipeline +# 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}' ``` -#### Verification +### Python API -```bash -curl http://localhost:8000/health +Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the +resulting `DiffusionRequestBatch` as the only positional argument: -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 +```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]) ``` -#### 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. +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 + +- +- +- +- [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) diff --git a/recipes/README.md b/recipes/README.md index cc1ac4920ad..1d89f0ef1be 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -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) | diff --git a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py index 5f8de588ede..5fe027a2007 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py +++ b/tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py @@ -8,7 +8,6 @@ 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 @@ WanT2VDMD2Pipeline: Wan22Pipeline, WanI2VDMD2Pipeline: Wan22I2VPipeline, LTX2T2VDMD2Pipeline: LTX2Pipeline, - LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline, + LTX2I2VDMD2Pipeline: LTX2Pipeline, FluxDMD2Pipeline: FluxPipeline, QwenImageDMD2Pipeline: QwenImagePipeline, } diff --git a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py index a395af95981..4df65f1fe8c 100644 --- a/tests/diffusion/models/dmd2/test_dmd2_scheduler.py +++ b/tests/diffusion/models/dmd2/test_dmd2_scheduler.py @@ -8,7 +8,6 @@ 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 @@ WanT2VDMD2Pipeline: Wan22Pipeline, WanI2VDMD2Pipeline: Wan22I2VPipeline, LTX2T2VDMD2Pipeline: LTX2Pipeline, - LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline, + LTX2I2VDMD2Pipeline: LTX2Pipeline, FluxDMD2Pipeline: FluxPipeline, QwenImageDMD2Pipeline: QwenImagePipeline, } diff --git a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py index 43737f8c102..6b4ec70ce0e 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_conditioning.py +++ b/tests/diffusion/models/ltx2/test_ltx2_conditioning.py @@ -7,6 +7,7 @@ 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 @@ def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(sel ) 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 @@ def test_denoise_timestep_kwargs_masks_video_only(self): 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 @@ def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self): "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 @@ def fake_randn_tensor(shape, generator=None, device=None, dtype=None): 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 @@ def fake_randn_tensor(shape, generator=None, device=None, dtype=None): ) 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 @@ def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypat 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 @@ def fake_randn_tensor(shape, generator=None, device=None, dtype=None): 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 @@ def fake_randn_tensor(shape, generator=None, device=None, dtype=None): 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 @@ def step(self, noise_pred, t, latents, return_dict=False): 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]]])) diff --git a/tests/diffusion/models/ltx2/test_ltx2_latents.py b/tests/diffusion/models/ltx2/test_ltx2_latents.py index 74cad09d1a9..36f7e9687d7 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_latents.py +++ b/tests/diffusion/models/ltx2/test_ltx2_latents.py @@ -4,7 +4,7 @@ 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"): diff --git a/tests/diffusion/models/ltx2/test_ltx2_parallel.py b/tests/diffusion/models/ltx2/test_ltx2_parallel.py index 0c6b12fbe48..92077a229bd 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_parallel.py +++ b/tests/diffusion/models/ltx2/test_ltx2_parallel.py @@ -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 @@ def test_combine_cfg_noise_matches_x0_space_formula(self): 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 @@ def fake_encode_prompt(**kwargs): ("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 @@ def test_forward_cfg_parallel_steps_video_and_audio_scheduler( 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 @@ def all_gather(self, tensor, separate_tensors=True): 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 @@ def to(self, device): 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 @@ def fake_encode_prompt(**kwargs): 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 @@ def fake_encode_prompt(**kwargs): 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. diff --git a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py index 3ba36ea5748..d70820f419e 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_pipeline.py +++ b/tests/diffusion/models/ltx2/test_ltx2_pipeline.py @@ -6,41 +6,52 @@ import json import os import tempfile +from dataclasses import replace from types import SimpleNamespace import pytest import torch from vllm_omni.diffusion.data import DiffusionOutput +from vllm_omni.diffusion.models.ltx2 import ltx2_components from vllm_omni.diffusion.models.ltx2.ltx2_components import ( LTX2_COMPONENT_PROFILE, + LTX2_DISTILLED_COMPONENT_PROFILE, LTX23_COMPONENT_PROFILE, + _install_connector_attention, + detect_ltx_model_version, ) from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import LTXI2VConditioningMixin from vllm_omni.diffusion.models.ltx2.ltx2_denoise import ( LTXDenoiseExecutor, LTXPhaseResult, + _official_ltx_sigmas, prepare_scheduler_stage, ) from vllm_omni.diffusion.models.ltx2.ltx2_guidance import ( - LTX_LEGACY_VELOCITY_GUIDANCE, - LTX_OFFICIAL_X0_GUIDANCE, + LTX_GUIDANCE_EXECUTOR, + LTXGuidancePlan, + LTXGuidanceSpec, + LTXModalityGuidance, + build_perturbation_kwargs, + combine_guided_x0, ) from vllm_omni.diffusion.models.ltx2.ltx2_latents import LTXAVState -from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import LTXPipelineRuntime -from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE -from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs +from vllm_omni.diffusion.models.ltx2.ltx2_recipes import ( + LTX2_DISTILLED_TWO_STAGE_RECIPE, + LTX2_ONE_STAGE_RECIPE, + LTX23_ONE_STAGE_RECIPE, + LTX_DEFAULT_NEGATIVE_PROMPT, + LTX_POSITIVE_ONLY_RECIPE, +) +from vllm_omni.diffusion.models.ltx2.ltx2_request import LTXRequestInputs, validate_pipeline_request +from vllm_omni.diffusion.models.ltx2.ltx2_runtime import LTXRuntime from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import ( - LTX2ImageToVideoPipeline, + LTX2I2VDMD2Pipeline, LTX2Pipeline, - LTX23ImageToVideoPipeline, - LTX23Pipeline, - LTXOneStagePipeline, -) -from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import ( - LTX2ImageToVideoTwoStagesPipeline, - LTX2TwoStagesPipeline, + LTX2T2VDMD2Pipeline, ) +from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline pytestmark = [pytest.mark.core_model, pytest.mark.cpu] @@ -50,21 +61,28 @@ def _make_ltx_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 -def _resolve_request_inputs_for_test(pipe, req): +def _resolve_request_inputs_for_test( + pipe, + req, + *, + guidance_scale=4.0, + negative_prompt=None, +): return pipe._resolve_request_inputs( req, prompt=None, - negative_prompt=None, + negative_prompt=negative_prompt, height=None, width=None, num_frames=None, frame_rate=None, num_inference_steps=None, - timesteps=None, - guidance_scale=4.0, + guidance_scale=guidance_scale, num_videos_per_prompt=1, generator=None, latents=None, @@ -80,38 +98,69 @@ def _resolve_request_inputs_for_test(pipe, req): ) -def test_ltx_versions_share_runtime_without_cross_version_inheritance(): +def test_ltx_public_entries_share_runtime_and_keep_recipe_boundaries(): from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import ( DistributedAutoencoderKLLTX2Video, ) - assert issubclass(LTX2Pipeline, LTXPipelineRuntime) - assert issubclass(LTX23Pipeline, LTXPipelineRuntime) - assert issubclass(LTX2TwoStagesPipeline, LTXPipelineRuntime) - assert not issubclass(LTX23Pipeline, LTX2Pipeline) + assert issubclass(LTX2Pipeline, LTXRuntime) + assert issubclass(LTX2DistilledPipeline, LTXRuntime) assert LTX2Pipeline.component_profile is LTX2_COMPONENT_PROFILE - assert LTX23Pipeline.component_profile is LTX23_COMPONENT_PROFILE - assert LTX2Pipeline.one_stage_recipe is LTX2_ONE_STAGE_RECIPE - assert LTX23Pipeline.one_stage_recipe is LTX23_ONE_STAGE_RECIPE + assert LTX2Pipeline.pipeline_recipe is LTX2_ONE_STAGE_RECIPE + assert LTX2DistilledPipeline.component_profile is LTX2_DISTILLED_COMPONENT_PROFILE assert LTX2_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video assert LTX23_COMPONENT_PROFILE.video_vae_cls is DistributedAutoencoderKLLTX2Video -def test_ltx_one_stage_variants_expose_shared_request_batch_and_distributed_decode_contracts(): - one_stage_variants = ( - LTX2Pipeline, - LTX2ImageToVideoPipeline, - LTX23Pipeline, - LTX23ImageToVideoPipeline, - ) - assert all(pipeline_cls.supports_request_batch for pipeline_cls in one_stage_variants) - assert all(pipeline_cls.distributed_video_decode for pipeline_cls in one_stage_variants) - assert not LTX2TwoStagesPipeline.supports_request_batch - assert not LTX2ImageToVideoTwoStagesPipeline.supports_request_batch +def test_ltx_checkpoint_version_detection_uses_metadata(tmp_path): + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]})) + assert detect_ltx_model_version(str(tmp_path)) == "2.3" + + (tmp_path / "model_index.json").write_text(json.dumps({"model_version": "2.3.1"})) + assert detect_ltx_model_version(str(tmp_path)) == "2.3" + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2Vocoder"]})) + assert detect_ltx_model_version(str(tmp_path)) == "2" + + +def test_ltx23_checkpoint_selects_version_specific_one_stage_profile(tmp_path, monkeypatch): + from vllm_omni.diffusion.models.ltx2 import ltx2_runtime + + (tmp_path / "model_index.json").write_text(json.dumps({"vocoder": ["ltx2", "LTX2VocoderWithBWE"]})) + + def stub_components(pipe, od_config): + pipe.od_config = od_config + pipe.vae_spatial_compression_ratio = 32 + monkeypatch.setattr(ltx2_runtime, "initialize_pipeline_components", stub_components) + monkeypatch.setattr(LTXRuntime, "setup_diffusion_pipeline_profiler", lambda *_args, **_kwargs: None) -@pytest.mark.parametrize("pipeline_cls", [LTX2Pipeline, LTX23Pipeline]) -def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cls): + pipe = LTX2Pipeline(od_config=SimpleNamespace(model=str(tmp_path), enable_diffusion_pipeline_profiler=False)) + + assert pipe.model_version == "2.3" + assert pipe.component_profile is LTX23_COMPONENT_PROFILE + assert pipe.pipeline_recipe is LTX23_ONE_STAGE_RECIPE + assert pipe.preserve_sp_padded_audio_duration + assert pipe.reports_stage_durations + + +def test_ltx_one_stage_entry_exposes_batch_image_and_distributed_decode_contracts(): + assert LTX2Pipeline.supports_request_batch + assert LTX2Pipeline.support_image_input + assert LTX2Pipeline.unified_text_image_entry + assert LTX2Pipeline.distributed_video_decode + assert not LTX2DistilledPipeline.supports_request_batch + assert LTX2DistilledPipeline.support_image_input + assert LTX2DistilledPipeline.unified_text_image_entry + + +def test_ltx_dmd2_entries_retain_their_task_boundaries(): + assert not LTX2T2VDMD2Pipeline.support_image_input + assert LTX2I2VDMD2Pipeline.support_image_input + assert not LTX2I2VDMD2Pipeline.unified_text_image_entry + + +def test_ltx_request_batch_decode_splits_video_and_audio_per_request(): from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -127,7 +176,7 @@ def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cl request_batch = DiffusionRequestBatch(requests) video = torch.tensor([[1.0], [2.0]]) audio = torch.tensor([[3.0], [4.0]]) - pipe = object.__new__(pipeline_cls) + pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) object.__setattr__(pipe, "_decode_output", lambda **_kwargs: DiffusionOutput(output=(video, audio))) forward_ctx = SimpleNamespace( @@ -159,12 +208,35 @@ def test_ltx_request_batch_decode_splits_video_and_audio_per_request(pipeline_cl torch.testing.assert_close(outputs[1].output[1], audio[1:]) -def test_ltx_one_stage_variants_share_forward_template(): - assert LTX2Pipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX2ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX23Pipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX23ImageToVideoPipeline._forward_impl is LTXOneStagePipeline._forward_impl - assert LTX2ImageToVideoPipeline.forward is LTX23ImageToVideoPipeline.forward +def test_ltx_public_entries_share_runtime_forward_template(): + assert "_forward_request" not in LTX2Pipeline.__dict__ + assert "forward" not in LTX2Pipeline.__dict__ + assert LTX2Pipeline._forward_request is LTXRuntime._forward_request + assert LTX2Pipeline._run_recipe is LTXRuntime._run_recipe + + +@pytest.mark.parametrize( + ("kwargs", "field_name"), + [ + ({"timesteps": [1]}, "timesteps"), + ({"guidance_rescale": 0.5}, "guidance_rescale"), + ({"noise_scale": 1.0}, "noise_scale"), + ({"return_dict": False}, "return_dict"), + ({"attention_kwargs": {}}, "attention_kwargs"), + ], +) +def test_ltx_public_forward_rejects_unsupported_generic_options(kwargs, field_name): + pipe = object.__new__(LTX2Pipeline) + + with pytest.raises(ValueError, match=field_name): + pipe.forward(SimpleNamespace(), **kwargs) + + +def test_ltx2_canonical_forward_requires_keyword_arguments(): + pipe = object.__new__(LTX2Pipeline) + + with pytest.raises(TypeError): + pipe.forward(SimpleNamespace(), "prompt") def test_ltx_versions_share_request_prompt_and_step_templates(): @@ -177,39 +249,315 @@ def test_ltx_versions_share_request_prompt_and_step_templates(): "_denoise_step", ) for method_name in shared_methods: - base_method = getattr(LTXPipelineRuntime, method_name) + base_method = getattr(LTXRuntime, method_name) assert getattr(LTX2Pipeline, method_name) is base_method - assert getattr(LTX23Pipeline, method_name) is base_method def test_ltx_versions_select_guidance_without_overriding_control_flow(): - assert LTX2Pipeline.guidance_strategy is LTX_LEGACY_VELOCITY_GUIDANCE - assert LTX23Pipeline.guidance_strategy is LTX_OFFICIAL_X0_GUIDANCE - assert LTX2Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step - assert LTX23Pipeline._predict_noise_for_step is LTXPipelineRuntime._predict_noise_for_step - assert LTX2Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise - assert LTX23Pipeline.combine_cfg_noise is LTXPipelineRuntime.combine_cfg_noise + assert LTX2Pipeline.guidance_executor is LTX_GUIDANCE_EXECUTOR + assert LTX2Pipeline._predict_noise_for_step is LTXRuntime._predict_noise_for_step + assert LTX2Pipeline.combine_cfg_noise is LTXRuntime.combine_cfg_noise + + +def test_ltx_official_recipes_build_all_guidance_passes(): + ltx2_plan = LTXGuidancePlan.build(LTX2_ONE_STAGE_RECIPE.request_guidance) + ltx23_plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) + + assert ltx2_plan.names == ltx23_plan.names == ("cond", "uncond", "ptb", "mod") + assert LTX2_ONE_STAGE_RECIPE.request_guidance.video.stg_blocks == (29,) + assert LTX23_ONE_STAGE_RECIPE.request_guidance.video.stg_blocks == (28,) + assert LTX23_ONE_STAGE_RECIPE.request_guidance.video.cfg_scale == 3.0 + assert LTX23_ONE_STAGE_RECIPE.request_guidance.audio.cfg_scale == 7.0 + assert LTX2_ONE_STAGE_RECIPE.phases[0].use_official_sigma_schedule + assert LTX23_ONE_STAGE_RECIPE.phases[0].use_official_sigma_schedule + assert not LTX_POSITIVE_ONLY_RECIPE.phases[0].use_official_sigma_schedule + + +def test_ltx_guidance_combines_official_x0_deltas_in_fp32(): + cond = torch.tensor([[1.0, 3.0]], dtype=torch.bfloat16) + uncond = torch.tensor([[0.0, 1.0]], dtype=torch.bfloat16) + perturbed = torch.tensor([[0.5, 2.0]], dtype=torch.bfloat16) + isolated = torch.tensor([[0.25, 2.5]], dtype=torch.bfloat16) + guidance = LTXModalityGuidance( + cfg_scale=3.0, + stg_scale=0.5, + modality_scale=2.0, + stg_blocks=(0,), + ) + + actual = combine_guided_x0( + cond=cond, + uncond_text=uncond, + uncond_perturbed=perturbed, + uncond_modality=isolated, + guidance=guidance, + ) + expected = ( + cond.float() + + 2.0 * (cond.float() - uncond.float()) + + 0.5 * (cond.float() - perturbed.float()) + + (cond.float() - isolated.float()) + ).to(cond.dtype) + + torch.testing.assert_close(actual, expected) + + +def test_ltx_guidance_skips_stg_without_perturbed_blocks(): + cond = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) + guidance = LTXModalityGuidance(stg_scale=1.0, stg_blocks=()) + + plan = LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)) + actual = combine_guided_x0( + cond=cond, + uncond_text=0.0, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + + assert plan.names == ("cond",) + torch.testing.assert_close(actual, cond) + + +def test_ltx_guidance_uses_official_close_to_disabled_semantics(): + guidance = LTXModalityGuidance( + cfg_scale=1.0 + 1e-10, + stg_scale=0.0, + modality_scale=1.0 + 1e-10, + stg_blocks=(0,), + ) + + assert LTXGuidancePlan.build(LTXGuidanceSpec(video=guidance)).names == ("cond",) + + +def test_ltx_guidance_rescale_is_invariant_to_request_batching(): + cond = torch.tensor( + [ + [[1.0, 2.0], [3.0, 4.0]], + [[10.0, 30.0], [50.0, 70.0]], + ] + ) + uncond = torch.stack((torch.zeros_like(cond[0]), cond[1])) + guidance = LTXModalityGuidance(cfg_scale=2.0, rescale_scale=1.0) + + actual = combine_guided_x0( + cond=cond, + uncond_text=uncond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + individually = torch.cat( + [ + combine_guided_x0( + cond=cond[index : index + 1], + uncond_text=uncond[index : index + 1], + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + for index in range(cond.shape[0]) + ] + ) + + # Fusing independent requests must not change either request's output. + torch.testing.assert_close(actual, individually) + + +def test_ltx_guidance_rescale_handles_zero_variance_prediction(): + cond = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]]) + guidance = LTXModalityGuidance(cfg_scale=2.0, rescale_scale=1.0) + + actual = combine_guided_x0( + cond=cond, + uncond_text=2 * cond, + uncond_perturbed=0.0, + uncond_modality=0.0, + guidance=guidance, + ) + + assert torch.isfinite(actual).all() + torch.testing.assert_close(actual, torch.zeros_like(cond)) + + +def test_ltx_guidance_builds_per_sample_stg_and_modality_masks(): + plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance) + perturbations = build_perturbation_kwargs(plan, batch_size=2, reference=torch.ones(8, 1, 1)) + + assert perturbations["video_self_attention_blocks"] == (28,) + assert perturbations["audio_self_attention_blocks"] == (28,) + assert perturbations["video_self_attention_mask"].flatten().tolist() == [1, 1, 1, 1, 0, 0, 1, 1] + assert perturbations["a2v_cross_attention_mask"].flatten().tolist() == [1, 1, 1, 1, 1, 1, 0, 0] + torch.testing.assert_close( + perturbations["a2v_cross_attention_mask"], + perturbations["v2a_cross_attention_mask"], + ) + + +def test_ltx_default_sigma_schedule_matches_official_shift_and_stretch(): + scheduler = SimpleNamespace( + config={ + "base_image_seq_len": 1024, + "max_image_seq_len": 4096, + "base_shift": 0.95, + "max_shift": 2.05, + "shift_terminal": 0.1, + } + ) + + sigmas = _official_ltx_sigmas(scheduler, steps=4, device=torch.device("cpu")) + + torch.testing.assert_close( + sigmas, + torch.tensor([1.0, 0.8671, 0.6316, 0.1, 0.0]), + rtol=1e-4, + atol=1e-4, + ) + + +def test_ltx_sigma_schedule_does_not_stretch_when_terminal_config_is_missing(): + scheduler = SimpleNamespace( + config={ + "base_image_seq_len": 1024, + "max_image_seq_len": 4096, + "base_shift": 0.95, + "max_shift": 2.05, + } + ) + + sigmas = _official_ltx_sigmas(scheduler, steps=4, device=torch.device("cpu")) + + torch.testing.assert_close( + sigmas, + torch.tensor([1.0, 0.9589, 0.8859, 0.7214, 0.0]), + rtol=1e-4, + atol=1e-4, + ) + + +def test_ltx_sigma_schedule_does_not_stretch_when_terminal_is_disabled(): + scheduler = SimpleNamespace( + config={ + "base_image_seq_len": 1024, + "max_image_seq_len": 4096, + "base_shift": 0.95, + "max_shift": 2.05, + "shift_terminal": None, + } + ) + + sigmas = _official_ltx_sigmas(scheduler, steps=4, device=torch.device("cpu")) + + torch.testing.assert_close( + sigmas, + torch.tensor([1.0, 0.9589, 0.8859, 0.7214, 0.0]), + rtol=1e-4, + atol=1e-4, + ) + + +def test_ltx_component_loader_installs_omni_connector_attention(monkeypatch): + installed = [] + kernels = [] + + class OmniAttention: + def __init__(self, **kwargs): + kernels.append(kwargs) + + class Attention: + heads = 32 + head_dim = 128 + inner_kv_dim = 4096 + + def set_processor(self, processor): + installed.append(processor) + monkeypatch.setattr(ltx2_components, "OmniAttention", OmniAttention) + connectors = SimpleNamespace( + video_connector=SimpleNamespace( + learnable_registers=object(), + transformer_blocks=[SimpleNamespace(attn1=Attention())], + ), + audio_connector=SimpleNamespace( + learnable_registers=object(), + transformer_blocks=[SimpleNamespace(attn1=Attention())], + ), + ) + + _install_connector_attention(connectors) + + assert len(installed) == 2 + assert all(processor.__class__.__name__ == "_LTXConnectorAttnProcessor" for processor in installed) + assert all(processor.has_learned_registers for processor in installed) + assert all(kernel["role"] == "ltx2.connector" for kernel in kernels) + assert all(kernel["role_category"] == "self" for kernel in kernels) + assert all(kernel["skip_sequence_parallel"] for kernel in kernels) + assert all(kernel["disable_kv_quant"] for kernel in kernels) + + +@pytest.mark.parametrize("has_learned_registers", [False, True]) +def test_ltx_connector_attention_dispatches_through_omni_kernel(has_learned_registers): + calls = [] + + class Backend: + @staticmethod + def get_name(): + return "FLASH_ATTN" + + class OmniAttention: + attn_backend = Backend + + def __call__(self, query, key, value, metadata): + calls.append((query, key, value, metadata)) + return query + + attention = SimpleNamespace( + heads=2, + head_dim=4, + inner_kv_dim=8, + norm_q=torch.nn.Identity(), + norm_k=torch.nn.Identity(), + to_q=torch.nn.Identity(), + to_k=torch.nn.Identity(), + to_v=torch.nn.Identity(), + to_out=(torch.nn.Identity(), torch.nn.Identity()), + to_gate_logits=None, + rope_type="split", + omni_attention=OmniAttention(), + ) + hidden_states = torch.randn(2, 3, 8) + additive_mask = torch.zeros(2, 1, 3, 3) + + output = ltx2_components._LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers)( + attention, + hidden_states, + attention_mask=additive_mask, + ) + + query, key, value, metadata = calls[0] + assert query.shape == key.shape == value.shape == (2, 3, 2, 4) + if has_learned_registers: + assert metadata is None + else: + assert metadata.attn_mask.shape == (2, 3) + assert metadata.attn_mask.dtype == torch.bool + assert metadata.attn_mask.all() + torch.testing.assert_close(output, hidden_states) -def test_ltx2_two_stage_variants_share_stage_orchestration(): - assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTX2TwoStagesPipeline) - assert issubclass(LTX2ImageToVideoTwoStagesPipeline, LTXI2VConditioningMixin) - assert LTX2ImageToVideoTwoStagesPipeline._run_two_stage is LTX2TwoStagesPipeline._run_two_stage + +def test_ltx2_distilled_pipeline_shares_recipe_runtime(): + assert issubclass(LTX2DistilledPipeline, LTXI2VConditioningMixin) + assert issubclass(LTX2DistilledPipeline, LTX2Pipeline) + assert LTX2DistilledPipeline._run_recipe is LTXRuntime._run_recipe + assert "_forward_request" not in LTX2DistilledPipeline.__dict__ def test_ltx_variants_share_denoise_loop_and_i2v_conditioning(): assert "_denoise_loop" not in LTX2Pipeline.__dict__ - assert "_denoise_loop" not in LTX23Pipeline.__dict__ - assert issubclass(LTX2ImageToVideoPipeline, LTXI2VConditioningMixin) - assert issubclass(LTX23ImageToVideoPipeline, LTXI2VConditioningMixin) - assert LTX2ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents - assert LTX23ImageToVideoPipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents - assert LTX2Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents - assert LTX23Pipeline.prepare_latents is LTXPipelineRuntime.prepare_latents - assert LTX2Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents - assert LTX23Pipeline.prepare_audio_latents is LTXPipelineRuntime.prepare_audio_latents - assert LTX2Pipeline.decode_phase is LTXPipelineRuntime.decode_phase - assert LTX23Pipeline.decode_phase is LTXPipelineRuntime.decode_phase + assert issubclass(LTX2Pipeline, LTXI2VConditioningMixin) + assert LTX2Pipeline.prepare_latents is LTXI2VConditioningMixin.prepare_latents + assert LTX2Pipeline.prepare_audio_latents is LTXRuntime.prepare_audio_latents + assert LTX2Pipeline.decode_phase is LTXRuntime.decode_phase def test_denoise_executor_owns_progress_and_interrupt(): @@ -246,17 +594,16 @@ def step(index, timestep, state): torch.testing.assert_close(state.audio, torch.tensor(11.0)) -def test_ltx2_two_stage_reuses_prompt_context_between_phases(): +def test_ltx2_distilled_two_stage_executes_declarative_phase_plan(): request_inputs = LTXRequestInputs( prompt="prompt", - negative_prompt="negative", - height=32, - width=32, + negative_prompt="", + height=64, + width=64, num_frames=1, frame_rate=24.0, - num_inference_steps=4, - guidance_scale=4.0, - guidance_rescale=0.0, + num_inference_steps=8, + guidance=LTXGuidanceSpec.positive_only(), num_videos_per_prompt=1, generator=None, latents=None, @@ -277,55 +624,308 @@ def resolve_request_inputs(req, **kwargs): return request_inputs def run_phase(req, inputs, *, prompt_context=None, **kwargs): - phase_calls.append((inputs, prompt_context)) + phase_recipe = kwargs["phase_recipe"] + assert kwargs["image"] is source_image + phase_calls.append((phase_recipe, inputs, prompt_context)) if len(phase_calls) == 1: assert prompt_context is None assert inputs.decode_timestep == 0.25 assert inputs.decode_noise_scale == 0.5 context = prompt_context_sentinel - video = torch.tensor([1.0]) - audio = torch.tensor([2.0]) + assert phase_recipe.name == "generate_lowres" + assert (inputs.height, inputs.width) == (32, 32) + assert inputs.num_inference_steps == 8 + assert kwargs["noise_scale"] == 1.0 + video = torch.ones(1, 128, 1, 1, 1) + audio = torch.full((1, 8, 1, 2), 2.0) + audio_for_next_phase = torch.full((1, 8, 1, 2), 5.0) else: assert prompt_context is prompt_context_sentinel - torch.testing.assert_close(inputs.latents, torch.tensor([11.0])) - torch.testing.assert_close(inputs.audio_latents, torch.tensor([2.0])) + assert phase_recipe.name == "refine" + assert (inputs.height, inputs.width) == (64, 64) + assert inputs.latents.shape == (1, 128, 1, 2, 2) + torch.testing.assert_close(inputs.audio_latents, torch.full((1, 8, 1, 2), 5.0)) + assert inputs.audio_latents_normalized assert inputs.guidance_scale == 1.0 assert inputs.num_inference_steps == 3 assert inputs.decode_timestep == 0.0 assert inputs.decode_noise_scale is None context = prompt_context - video = torch.tensor([3.0]) - audio = torch.tensor([4.0]) + video = torch.full((1, 128, 1, 2, 2), 3.0) + audio = torch.full((1, 8, 1, 2), 4.0) + audio_for_next_phase = None return LTXPhaseResult( forward_context=SimpleNamespace(prompt_context=context), video=video, audio=audio, + audio_for_next_phase=audio_for_next_phase, ) def decode_phase(phase): return DiffusionOutput(output=(phase.video, phase.audio)) class FakeUpsampler(torch.nn.Module): - def forward(self, *, latents, output_type, return_dict): - assert output_type == "latent" - assert not return_dict - return (latents + 10,) + dtype = torch.float32 + + def forward(self, latents): + return latents.repeat_interleave(2, dim=-2).repeat_interleave(2, dim=-1) prompt_context_sentinel = prompt_context - pipeline = object.__new__(LTX2TwoStagesPipeline) + pipeline = object.__new__(LTX2DistilledPipeline) torch.nn.Module.__init__(pipeline) - pipeline.distilled = True - pipeline.upsample_pipe = FakeUpsampler() + pipeline.device = torch.device("cpu") + pipeline.vae_spatial_compression_ratio = 32 + pipeline.vae_temporal_compression_ratio = 8 + pipeline.latent_upsampler = FakeUpsampler() object.__setattr__(pipeline, "_resolve_request_inputs", resolve_request_inputs) object.__setattr__(pipeline, "run_phase", run_phase) object.__setattr__(pipeline, "decode_phase", decode_phase) - output = pipeline.forward(SimpleNamespace()) + source_image = object() + output = pipeline.forward(SimpleNamespace(sampling_params_list=[], prompts=[]), image=source_image) assert len(phase_calls) == 2 - assert phase_calls[1][1] is prompt_context_sentinel - torch.testing.assert_close(output.output[0], torch.tensor([3.0])) - torch.testing.assert_close(output.output[1], torch.tensor([4.0])) + assert phase_calls[1][2] is prompt_context_sentinel + torch.testing.assert_close(output.output[0], torch.full((1, 128, 1, 2, 2), 3.0)) + torch.testing.assert_close(output.output[1], torch.full((1, 8, 1, 2), 4.0)) + + +def test_ltx2_distilled_stage2_reapplies_final_resolution_i2v_conditioning(): + from vllm_omni.diffusion.models.ltx2 import ltx2_latents + + pipeline = object.__new__(LTX2DistilledPipeline) + torch.nn.Module.__init__(pipeline) + pipeline.vae_spatial_compression_ratio = 32 + pipeline.vae_temporal_compression_ratio = 8 + pipeline.transformer_spatial_patch_size = 1 + pipeline.transformer_temporal_patch_size = 1 + pipeline.vae = SimpleNamespace( + latents_mean=torch.zeros(2), + latents_std=torch.ones(2), + config=SimpleNamespace(scaling_factor=1.0), + ) + object.__setattr__( + pipeline, + "_encode_i2v_image_latents", + lambda image, **_kwargs: torch.full((image.shape[0], 2, 1, 2, 2), 7.0), + ) + + packed, conditioning_mask = pipeline.prepare_latents( + image=torch.zeros(1, 3, 64, 64), + batch_size=1, + num_channels_latents=2, + height=64, + width=64, + num_frames=9, + noise_scale=0.0, + dtype=torch.float32, + device=torch.device("cpu"), + latents=torch.ones(1, 2, 2, 2, 2), + ) + unpacked = ltx2_latents.unpack_latents(packed, 2, 2, 2) + + torch.testing.assert_close(unpacked[:, :, :1], torch.full((1, 2, 1, 2, 2), 7.0)) + torch.testing.assert_close(unpacked[:, :, 1:], torch.ones(1, 2, 1, 2, 2)) + assert conditioning_mask.shape == (1, 8) + assert conditioning_mask[:, :4].all() + assert not conditioning_mask[:, 4:].any() + + +def test_ltx2_distilled_two_stage_recipe_is_fixed_positive_only(): + stage1, stage2 = LTX2_DISTILLED_TWO_STAGE_RECIPE.phases + + assert stage1.spatial_downscale == 2 + assert stage1.guidance == stage2.guidance == LTXGuidanceSpec.positive_only() + assert LTXGuidancePlan.build(stage1.guidance).names == ("cond",) + assert LTXGuidancePlan.build(stage2.guidance).names == ("cond",) + assert not stage1.allow_guidance_override + assert not stage2.allow_guidance_override + assert stage1.noise_scale == 1.0 + assert stage1.sigmas == (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0) + assert stage2.sigmas == (0.909375, 0.725, 0.421875, 0.0) + assert not LTX2_DISTILLED_TWO_STAGE_RECIPE.allow_request_sigmas + assert not LTX2_DISTILLED_TWO_STAGE_RECIPE.allow_request_latents + assert not LTX2_DISTILLED_TWO_STAGE_RECIPE.allow_negative_prompt + assert LTX2_DISTILLED_TWO_STAGE_RECIPE.fixed_num_inference_steps + assert (LTX2_DISTILLED_TWO_STAGE_RECIPE.height, LTX2_DISTILLED_TWO_STAGE_RECIPE.width) == (1024, 1536) + + +@pytest.mark.parametrize( + ("request_inputs", "error"), + [ + (lambda request: replace(request, height=32), "divisible by 64"), + (lambda request: replace(request, num_frames=8), "8 \\* k \\+ 1"), + (lambda request: replace(request, num_inference_steps=7), "uses 8 fixed Stage 1"), + (lambda request: replace(request, latents=torch.empty(1)), "request-provided"), + (lambda request: replace(request, audio_latents=torch.empty(1)), "request-provided"), + (lambda request: replace(request, negative_prompt="negative"), "negative conditioning"), + (lambda request: replace(request, negative_prompt_embeds=torch.empty(1)), "negative conditioning"), + (lambda request: replace(request, guidance=LTX2_ONE_STAGE_RECIPE.request_guidance), "positive-only"), + ], +) +def test_ltx_two_stage_recipe_rejects_unsupported_request_values(request_inputs, error): + request = LTXRequestInputs( + prompt="prompt", + negative_prompt="", + height=64, + width=64, + num_frames=1, + frame_rate=24.0, + num_inference_steps=8, + guidance=LTXGuidanceSpec.positive_only(), + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=16, + ) + + with pytest.raises(ValueError, match=error): + validate_pipeline_request( + request_inputs(request), + pipeline_recipe=LTX2_DISTILLED_TWO_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2DistilledPipeline", + ) + + +def test_ltx_two_stage_recipe_rejects_sigmas(): + request = LTXRequestInputs( + prompt="prompt", + negative_prompt="", + height=64, + width=64, + num_frames=1, + frame_rate=24.0, + num_inference_steps=8, + guidance=LTXGuidanceSpec.positive_only(), + num_videos_per_prompt=1, + generator=None, + latents=None, + audio_latents=None, + prompt_embeds=None, + negative_prompt_embeds=None, + prompt_attention_mask=None, + negative_prompt_attention_mask=None, + decode_timestep=0.0, + decode_noise_scale=None, + output_type="np", + max_sequence_length=16, + ) + + with pytest.raises(ValueError, match="fixed phase sigma schedules"): + validate_pipeline_request( + request, + pipeline_recipe=LTX2_DISTILLED_TWO_STAGE_RECIPE, + vae_spatial_compression_ratio=32, + vae_temporal_compression_ratio=8, + pipeline_name="LTX2DistilledPipeline", + request_sigmas=[1.0, 0.0], + ) + + +@pytest.mark.parametrize( + ("direct_kwargs", "sampling_kwargs", "error"), + [ + ({"sigmas": [1.0, 0.0]}, {}, "fixed phase sigma schedules"), + ({}, {"sigmas": [1.0, 0.0]}, "fixed phase sigma schedules"), + ({"latents": torch.empty(1)}, {}, "request-provided video or audio latents"), + ({}, {"latents": torch.empty(1)}, "request-provided video or audio latents"), + ({"audio_latents": torch.empty(1)}, {}, "request-provided video or audio latents"), + ( + {}, + {"extra_args": {"audio_latents": torch.empty(1)}}, + "request-provided video or audio latents", + ), + ({"negative_prompt": "negative"}, {}, "negative conditioning"), + ({"negative_prompt_embeds": torch.empty(1)}, {}, "negative conditioning"), + ({}, {"num_inference_steps": 7}, "uses 8 fixed Stage 1"), + ({"guidance_scale": 4.0}, {}, "fixed positive-only guidance"), + ({}, {"guidance_scale": 4.0}, "fixed positive-only guidance"), + ], + ids=( + "direct-sigmas", + "request-sigmas", + "direct-video-latents", + "request-video-latents", + "direct-audio-latents", + "request-audio-latents", + "direct-negative-prompt", + "direct-negative-embeds", + "request-num-steps", + "direct-guidance", + "request-guidance", + ), +) +def test_ltx_distilled_forward_rejects_fixed_recipe_overrides(direct_kwargs, sampling_kwargs, error): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + sampling_values = { + "height": 64, + "width": 64, + "num_frames": 1, + "num_inference_steps": 8, + **sampling_kwargs, + } + sampling = OmniDiffusionSamplingParams(**sampling_values) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=sampling, + request_id="ltx-distilled-fixed-inputs", + ) + ] + ) + pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline) + + with pytest.raises(ValueError, match=error): + pipeline.forward(req, **direct_kwargs) + + +def test_ltx_distilled_dummy_run_uses_fixed_recipe_values(): + from vllm_omni.diffusion.request import DUMMY_DIFFUSION_REQUEST_ID, OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="dummy run", + sampling_params=OmniDiffusionSamplingParams( + height=512, + width=512, + num_frames=LTX2DistilledPipeline.dummy_run_num_frames, + num_inference_steps=1, + guidance_scale=0.0, + ), + request_id=DUMMY_DIFFUSION_REQUEST_ID, + ) + ] + ) + pipeline = _make_ltx_request_pipe(LTX2DistilledPipeline) + seen = {} + + def fake_run_recipe(_req, request_inputs, **_kwargs): + seen["request_inputs"] = request_inputs + return "dummy-output" + + object.__setattr__(pipeline, "_run_recipe", fake_run_recipe) + + assert pipeline.forward(req) == "dummy-output" + assert seen["request_inputs"].num_frames == 1 + assert seen["request_inputs"].num_inference_steps == 8 + assert seen["request_inputs"].guidance == LTXGuidanceSpec.positive_only() def test_ltx_custom_sigmas_bypass_scheduler_shifting(): @@ -337,10 +937,7 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): max_shift=2.05, shift_terminal=0.1, ) - pipeline = SimpleNamespace( - scheduler=scheduler, - _make_video_audio_scheduler=lambda *args, **kwargs: object(), - ) + pipeline = SimpleNamespace(scheduler=scheduler) sigmas = [1.0, 0.75, 0.25, 0.0] audio_scheduler, _, timesteps = prepare_scheduler_stage( @@ -352,6 +949,7 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): latent_num_frames=1, latent_height=1, latent_width=1, + use_official_sigma_schedule=False, ) expected = torch.tensor(sigmas, dtype=torch.float32) @@ -361,14 +959,56 @@ def test_ltx_custom_sigmas_bypass_scheduler_shifting(): class TestLTXRequestParsing: - @pytest.mark.parametrize( - ("t2v_cls", "i2v_cls"), - [ - (LTX2Pipeline, LTX2ImageToVideoPipeline), - (LTX23Pipeline, LTX23ImageToVideoPipeline), - ], - ) - def test_t2v_and_i2v_share_request_input_resolution(self, t2v_cls, i2v_cls): + def test_request_resolves_official_and_explicit_negative_prompts_per_sample(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt=prompt, + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2), + request_id=f"ltx-negative-prompt-{index}", + ) + for index, prompt in enumerate( + ( + "use the official default", + {"prompt": "disable the default", "negative_prompt": ""}, + {"prompt": "custom negative", "negative_prompt": "custom"}, + ) + ) + ] + ) + + resolved = _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req) + + assert resolved.negative_prompt == [LTX_DEFAULT_NEGATIVE_PROMPT, "", "custom"] + + def test_forward_negative_prompt_fallback_is_not_discarded(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=2), + request_id="ltx-forward-negative-prompt", + ) + ] + ) + + resolved = _resolve_request_inputs_for_test( + _make_ltx_request_pipe(LTX2Pipeline), + req, + negative_prompt="forward fallback", + ) + + assert resolved.negative_prompt == ["forward fallback"] + + def test_unified_t2v_i2v_entry_resolves_request_inputs(self): from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -415,58 +1055,164 @@ def test_t2v_and_i2v_share_request_input_resolution(self, t2v_cls, i2v_cls): ] ) - resolved_t2v = _resolve_request_inputs_for_test( - _make_ltx_request_pipe(t2v_cls), - req, + resolved = _resolve_request_inputs_for_test(_make_ltx_request_pipe(LTX2Pipeline), req) + + assert resolved.prompt is None + assert resolved.negative_prompt is None + assert resolved.height == 384 + assert resolved.width == 512 + assert resolved.num_frames == 25 + assert resolved.frame_rate == 12.5 + assert resolved.num_inference_steps == 2 + assert resolved.guidance_scale == 6.0 + assert resolved.num_videos_per_prompt == 2 + assert isinstance(resolved.generator, list) + assert [gen.initial_seed() for gen in resolved.generator] == [123, 123] + assert resolved.decode_timestep == [0.1] + assert resolved.decode_noise_scale == [0.2] + assert resolved.output_type == "latent" + assert resolved.max_sequence_length == 17 + torch.testing.assert_close(resolved.latents, video_latents) + torch.testing.assert_close(resolved.audio_latents, audio_latents) + torch.testing.assert_close(resolved.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) + torch.testing.assert_close(resolved.prompt_attention_mask, torch.stack([prompt_attention_mask])) + torch.testing.assert_close(resolved.negative_prompt_attention_mask, torch.stack([negative_attention_mask])) + + def test_nonbatch_distilled_request_uses_precomputed_positive_embeddings(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + prompt_embeds = torch.tensor([[1.0, 2.0]]) + prompt_attention_mask = torch.tensor([True, False]) + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={ + "prompt": "must not be forwarded with embeddings", + "additional_information": { + "prompt_embeds": prompt_embeds, + "attention_mask": prompt_attention_mask, + }, + }, + sampling_params=OmniDiffusionSamplingParams(num_inference_steps=8), + request_id="ltx-distilled-prompt-embeds", + ) + ] ) - resolved_i2v = _resolve_request_inputs_for_test( - _make_ltx_request_pipe(i2v_cls), + + resolved = _resolve_request_inputs_for_test( + _make_ltx_request_pipe(LTX2DistilledPipeline), req, + guidance_scale=None, ) - assert resolved_i2v.prompt == resolved_t2v.prompt is None - assert resolved_i2v.negative_prompt == resolved_t2v.negative_prompt is None - assert resolved_i2v.height == resolved_t2v.height == 384 - assert resolved_i2v.width == resolved_t2v.width == 512 - assert resolved_i2v.num_frames == resolved_t2v.num_frames == 25 - assert resolved_i2v.frame_rate == resolved_t2v.frame_rate == 12.5 - assert resolved_i2v.num_inference_steps == resolved_t2v.num_inference_steps == 2 - assert resolved_i2v.guidance_scale == resolved_t2v.guidance_scale == 6.0 - assert resolved_i2v.num_videos_per_prompt == resolved_t2v.num_videos_per_prompt == 2 - assert isinstance(resolved_i2v.generator, list) - assert isinstance(resolved_t2v.generator, list) - assert [gen.initial_seed() for gen in resolved_i2v.generator] == [123, 123] - assert [gen.initial_seed() for gen in resolved_t2v.generator] == [123, 123] - assert resolved_i2v.decode_timestep == resolved_t2v.decode_timestep == [0.1] - assert resolved_i2v.decode_noise_scale == resolved_t2v.decode_noise_scale == [0.2] - assert resolved_i2v.output_type == resolved_t2v.output_type == "latent" - assert resolved_i2v.max_sequence_length == resolved_t2v.max_sequence_length == 17 - torch.testing.assert_close(resolved_i2v.latents, video_latents) - torch.testing.assert_close(resolved_t2v.latents, video_latents) - torch.testing.assert_close(resolved_i2v.audio_latents, audio_latents) - torch.testing.assert_close(resolved_t2v.audio_latents, audio_latents) - torch.testing.assert_close(resolved_i2v.prompt_embeds, torch.stack([prompt_embeds])) - torch.testing.assert_close(resolved_t2v.prompt_embeds, torch.stack([prompt_embeds])) - torch.testing.assert_close(resolved_i2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) - torch.testing.assert_close(resolved_t2v.negative_prompt_embeds, torch.stack([negative_prompt_embeds])) - torch.testing.assert_close(resolved_i2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) - torch.testing.assert_close(resolved_t2v.prompt_attention_mask, torch.stack([prompt_attention_mask])) - torch.testing.assert_close( - resolved_i2v.negative_prompt_attention_mask, - torch.stack([negative_attention_mask]), + assert resolved.prompt is None + assert resolved.negative_prompt == [""] + torch.testing.assert_close(resolved.prompt_embeds, torch.stack([prompt_embeds])) + torch.testing.assert_close(resolved.prompt_attention_mask, torch.stack([prompt_attention_mask])) + + def test_request_resolves_per_modality_guidance_and_common_cfg_override(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt={"prompt": "prompt", "negative_prompt": "negative"}, + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args={ + "video_cfg_scale": 4.0, + "audio_cfg_guidance_scale": 6.0, + "video_stg_guidance_scale": 0.5, + "audio_stg_scale": 0.25, + "a2v_guidance_scale": 2.0, + "v2a_guidance_scale": 4.0, + "video_rescale_scale": 0.2, + "audio_rescale_scale": 0.3, + "video_stg_blocks": [3, 4], + "audio_stg_blocks": 5, + }, + ), + request_id="ltx-guidance-overrides", + ) + ] ) - torch.testing.assert_close( - resolved_t2v.negative_prompt_attention_mask, - torch.stack([negative_attention_mask]), + pipe = _make_ltx_request_pipe(LTX2Pipeline) + + resolved = _resolve_request_inputs_for_test(pipe, req, guidance_scale=None) + + assert resolved.guidance.video == LTXModalityGuidance(4.0, 0.5, 2.0, 0.2, (3, 4)) + assert resolved.guidance.audio == LTXModalityGuidance(6.0, 0.25, 4.0, 0.3, (5,)) + + common = _resolve_request_inputs_for_test(pipe, req, guidance_scale=2.0) + assert common.guidance.video.cfg_scale == common.guidance.audio.cfg_scale == 2.0 + # LTX exposes independent video/audio rescale controls, not the + # model-agnostic sampling field. + assert common.guidance.video.rescale_scale == 0.2 + assert common.guidance.audio.rescale_scale == 0.3 + + def test_request_rejects_unsupported_generic_scheduler_and_rescale_controls(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + cases = ( + (OmniDiffusionSamplingParams(timesteps=torch.tensor([1.0])), "timesteps"), + (OmniDiffusionSamplingParams(guidance_rescale=0.5), "guidance_rescale"), + (OmniDiffusionSamplingParams(extra_args={"flow_shift": 3.0}), "flow_shift"), + ) + pipe = _make_ltx_request_pipe(LTX2Pipeline) + + for index, (sampling, field_name) in enumerate(cases): + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt", + sampling_params=sampling, + request_id=f"ltx-unsupported-{index}", + ) + ] + ) + with pytest.raises(ValueError, match=field_name): + _resolve_request_inputs_for_test(pipe, req) + + def test_request_batch_rejects_mixed_guidance_specs(self): + from vllm_omni.diffusion.request import OmniDiffusionRequest + from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch + from vllm_omni.inputs.data import OmniDiffusionSamplingParams + + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="prompt-0", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + extra_args={"video_cfg_scale": scale}, + ), + request_id=f"ltx-mixed-guidance-{index}", + ) + for index, scale in enumerate((3.0, 4.0)) + ] ) + with pytest.raises(ValueError, match="identical video/audio guidance"): + _resolve_request_inputs_for_test( + _make_ltx_request_pipe(LTX2Pipeline), + req, + guidance_scale=None, + ) + def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + 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_ltx_request_pipe(LTX23Pipeline) + pipe = _make_ltx_request_pipe(LTX2Pipeline) req = DiffusionRequestBatch( [ OmniDiffusionRequest( @@ -499,11 +1245,40 @@ def test_request_input_resolution_rejects_mixed_precomputed_prompt_fields(self): _resolve_request_inputs_for_test(pipe, req) -class TestLTX23ForwardStages: +class TestLTXForwardStages: + def test_encode_prompt_batches_positive_and_negative_gemma_inputs(self): + pipe = _make_ltx_request_pipe(LTX2Pipeline) + calls = [] + + def fake_get_gemma_prompt_embeds(**kwargs): + calls.append(kwargs) + count = len(kwargs["prompt"]) + embeds = torch.arange(count, dtype=torch.float32).view(count, 1, 1) + mask = torch.arange(count).view(count, 1) + return embeds, mask + + object.__setattr__(pipe, "_get_gemma_prompt_embeds", fake_get_gemma_prompt_embeds) + + prompt_embeds, prompt_mask, negative_embeds, negative_mask = pipe.encode_prompt( + prompt=["positive-0", "positive-1"], + negative_prompt=["negative-0", "negative-1"], + do_classifier_free_guidance=True, + num_videos_per_prompt=2, + device=torch.device("cpu"), + ) + + assert len(calls) == 1 + assert calls[0]["prompt"] == ["positive-0", "positive-1", "negative-0", "negative-1"] + assert calls[0]["num_videos_per_prompt"] == 1 + torch.testing.assert_close(prompt_embeds.flatten(), torch.tensor([0.0, 0.0, 1.0, 1.0])) + torch.testing.assert_close(negative_embeds.flatten(), torch.tensor([2.0, 2.0, 3.0, 3.0])) + torch.testing.assert_close(prompt_mask.flatten(), torch.tensor([0, 0, 1, 1])) + torch.testing.assert_close(negative_mask.flatten(), torch.tensor([2, 2, 3, 3])) + def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline - pipe = _make_ltx_request_pipe(LTX23Pipeline) + pipe = _make_ltx_request_pipe(LTX2Pipeline) prompt_embeds = torch.tensor([[[1.0], [2.0]], [[3.0], [4.0]]]) negative_prompt_embeds = torch.tensor([[[-1.0], [-2.0]], [[-3.0], [-4.0]]]) prompt_attention_mask = torch.tensor([[True, False], [False, True]]) @@ -542,13 +1317,13 @@ def test_encode_prompt_repeats_precomputed_embeds_for_num_outputs(self): torch.tensor([[False, False], [False, False], [True, True], [True, True]]), ) - def test_t2v_forward_delegates_to_shared_forward_impl(self): - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + def test_t2v_forward_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_ltx_request_pipe(LTX23Pipeline) + pipe = _make_ltx_request_pipe(LTX2Pipeline) request_sigmas = [1.0, 0.4] req = DiffusionRequestBatch( [ @@ -569,25 +1344,17 @@ def test_t2v_forward_delegates_to_shared_forward_impl(self): ] ) fallback_sigmas = [1.0, 0.5] - timesteps = [1000, 500] - attention_kwargs = {"scale": 1.0} 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 ["delegated"] - object.__setattr__(pipe, "_forward_impl", fake_forward_impl) + object.__setattr__(pipe, "_run_recipe", fake_run_recipe) - output = pipe.forward( - req, - sigmas=fallback_sigmas, - timesteps=timesteps, - noise_scale=0.25, - attention_kwargs=attention_kwargs, - ) + output = pipe.forward(req, sigmas=fallback_sigmas) assert output == ["delegated"] assert seen["req"] is req @@ -598,20 +1365,52 @@ def fake_forward_impl(req_arg, request_inputs, **kwargs): assert seen["request_inputs"].guidance_scale == 4.5 assert seen["request_inputs"].output_type == "latent" assert seen["kwargs"] == { - "noise_scale": 0.25, - "sigmas": request_sigmas, - "timesteps": timesteps, - "attention_kwargs": attention_kwargs, + "request_sigmas": request_sigmas, + "image": None, } + def test_t2v_forward_prefers_request_sigmas(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 + + request_sigmas = [1.0, 0.25] + req = DiffusionRequestBatch( + [ + OmniDiffusionRequest( + prompt="custom sigma schedule", + sampling_params=OmniDiffusionSamplingParams( + num_inference_steps=2, + sigmas=request_sigmas, + ), + request_id="ltx23-request-sigmas", + ) + ] + ) + pipe = _make_ltx_request_pipe(LTX2Pipeline) + seen = {} + + def fake_run_recipe(req_arg, request_inputs, **kwargs): + del req_arg, request_inputs + seen.update(kwargs) + return ["delegated"] + + object.__setattr__(pipe, "_run_recipe", fake_run_recipe) + + output = pipe.forward(req, sigmas=[1.0, 0.5]) + + assert output == ["delegated"] + assert seen["request_sigmas"] == request_sigmas + class TestPipelineComponents: - def test_ltx23_pipeline_declares_offload_components(self): - """LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery.""" - from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline + def test_ltx_pipeline_declares_offload_components(self): + """LTX2Pipeline exposes all one-stage modules to offload discovery.""" + from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery - pipe = object.__new__(LTX23Pipeline) + pipe = object.__new__(LTX2Pipeline) torch.nn.Module.__init__(pipe) pipe.transformer = torch.nn.Linear(1, 1) pipe.text_encoder = torch.nn.Linear(1, 1) @@ -622,19 +1421,22 @@ def test_ltx23_pipeline_declares_offload_components(self): modules = ModuleDiscovery.discover(pipe) - assert LTX23Pipeline._dit_modules == ["transformer"] - assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"] - assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"] - assert LTX23Pipeline._resident_modules == ["vocoder"] + assert LTX2Pipeline._dit_modules == ["transformer"] + assert LTX2Pipeline._encoder_modules == ["text_encoder", "connectors"] + assert LTX2Pipeline._vae_modules == ["vae", "audio_vae"] + assert LTX2Pipeline._resident_modules == ["vocoder"] assert modules.dit_names == ["transformer"] assert modules.encoder_names == ["text_encoder", "connectors"] assert modules.resident_names == ["vocoder"] assert len(modules.vaes) == 2 + def test_distilled_profile_keeps_upsampler_in_managed_resident_components(self): + assert LTX2_DISTILLED_COMPONENT_PROFILE.resident_modules == ("vocoder", "latent_upsampler") + class TestLTX23DecodeConditioning: def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self): - from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value + from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _expand_per_prompt_decode_value assert _expand_per_prompt_decode_value( [0.1, 0.2], @@ -656,7 +1458,7 @@ def test_decode_conditioning_expands_per_prompt_values_to_effective_batch(self): ) == [0.1, 0.2, 0.3, 0.4] def test_decode_conditioning_rejects_ambiguous_lengths(self): - from vllm_omni.diffusion.models.ltx2.ltx2_pipeline_runtime import _expand_per_prompt_decode_value + from vllm_omni.diffusion.models.ltx2.ltx2_runtime import _expand_per_prompt_decode_value with pytest.raises(ValueError, match="decode_timestep"): _expand_per_prompt_decode_value( @@ -668,19 +1470,26 @@ def test_decode_conditioning_rejects_ambiguous_lengths(self): class TestRegistryIntegration: - """Verify all LTX-2.3 pipeline variants are registered.""" + """Verify the unified one-stage and existing two-stage entries.""" def test_pipeline_module_paths(self): """Registry entries must point to the correct modules.""" from vllm_omni.diffusion.registry import _DIFFUSION_MODELS - assert _DIFFUSION_MODELS["LTX23Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX23Pipeline") - - assert _DIFFUSION_MODELS["LTX23ImageToVideoPipeline"] == ( + assert _DIFFUSION_MODELS["LTX2Pipeline"] == ("ltx2", "pipeline_ltx2", "LTX2Pipeline") + assert _DIFFUSION_MODELS["LTX2DistilledPipeline"] == ( "ltx2", - "pipeline_ltx2", - "LTX23ImageToVideoPipeline", + "pipeline_ltx2_two_stage", + "LTX2DistilledPipeline", ) + removed_entries = { + "LTX2ImageToVideoPipeline", + "LTX23Pipeline", + "LTX23ImageToVideoPipeline", + "LTX2TwoStagesPipeline", + "LTX2ImageToVideoTwoStagesPipeline", + } + assert removed_entries.isdisjoint(_DIFFUSION_MODELS) def test_post_process_funcs_registered(self): """Pipeline variants must map to get_ltx2_post_process_func.""" @@ -688,13 +1497,9 @@ def test_post_process_funcs_registered(self): expected = [ "LTX2Pipeline", - "LTX2ImageToVideoPipeline", - "LTX2TwoStagesPipeline", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", ] for name in expected: assert name in _DIFFUSION_POST_PROCESS_FUNCS, f"{name} not in _DIFFUSION_POST_PROCESS_FUNCS" @@ -704,13 +1509,9 @@ def test_post_process_funcs_registered(self): "model_class_name", [ "LTX2Pipeline", - "LTX2ImageToVideoPipeline", - "LTX2TwoStagesPipeline", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", ], ) def test_post_process_func_resolves_from_every_entry_module(self, model_class_name): @@ -735,10 +1536,7 @@ def test_cache_dit_for_ltx2_does_not_have_custom_enablers_registered(self): # NOTE: We used to have custom enablers for this model, but refactored to handle # it more generically. Now we only need to ensure it has git cache adapter config. - expected = [ - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", - ] + expected = ["LTX2Pipeline"] for name in expected: assert name not in CUSTOM_DIT_ENABLERS, f"{name} not in CUSTOM_DIT_ENABLERS" diff --git a/tests/diffusion/models/ltx2/test_ltx2_transformer.py b/tests/diffusion/models/ltx2/test_ltx2_transformer.py index e5ccfd013dc..de037fef867 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_transformer.py +++ b/tests/diffusion/models/ltx2/test_ltx2_transformer.py @@ -10,11 +10,12 @@ 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 diff --git a/tests/diffusion/models/ltx2/test_ltx2_vae.py b/tests/diffusion/models/ltx2/test_ltx2_vae.py index 4ae7a0721d3..108bb960d92 100644 --- a/tests/diffusion/models/ltx2/test_ltx2_vae.py +++ b/tests/diffusion/models/ltx2/test_ltx2_vae.py @@ -14,7 +14,7 @@ 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 @@ def test_non_output_rank_only_enters_collective_vae_decode( 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 @@ def decode(self, *_args, **_kwargs): 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 diff --git a/tests/e2e/accuracy/ltx/__init__.py b/tests/e2e/accuracy/ltx/__init__.py new file mode 100644 index 00000000000..08143263b9e --- /dev/null +++ b/tests/e2e/accuracy/ltx/__init__.py @@ -0,0 +1 @@ +"""LTX family official-reference accuracy tests.""" diff --git a/tests/e2e/accuracy/ltx/run_ltx_reference.py b/tests/e2e/accuracy/ltx/run_ltx_reference.py new file mode 100644 index 00000000000..a4a180047d3 --- /dev/null +++ b/tests/e2e/accuracy/ltx/run_ltx_reference.py @@ -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() diff --git a/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py new file mode 100644 index 00000000000..647ece7a74e --- /dev/null +++ b/tests/e2e/accuracy/ltx/test_ltx_official_similarity.py @@ -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 diff --git a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py b/tests/e2e/accuracy/test_ltx2_3_video_similarity.py deleted file mode 100644 index a169817a20f..00000000000 --- a/tests/e2e/accuracy/test_ltx2_3_video_similarity.py +++ /dev/null @@ -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, - ) diff --git a/tests/e2e/online_serving/test_ltx2_expansion.py b/tests/e2e/online_serving/test_ltx2_expansion.py index 82855f62fc8..4d0eb6c9027 100644 --- a/tests/e2e/online_serving/test_ltx2_expansion.py +++ b/tests/e2e/online_serving/test_ltx2_expansion.py @@ -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 @@ 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, } diff --git a/tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py b/tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py deleted file mode 100644 index 757d906ec3b..00000000000 --- a/tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py +++ /dev/null @@ -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") diff --git a/tests/model_extras/test_model_extras.py b/tests/model_extras/test_model_extras.py index 47c10823532..c73dbdf1e80 100644 --- a/tests/model_extras/test_model_extras.py +++ b/tests/model_extras/test_model_extras.py @@ -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: diff --git a/tests/model_tests/diffusion/test_alignment.py b/tests/model_tests/diffusion/test_alignment.py index 9d0a4911a6d..29cfc6f2172 100644 --- a/tests/model_tests/diffusion/test_alignment.py +++ b/tests/model_tests/diffusion/test_alignment.py @@ -30,13 +30,9 @@ "OvisImagePipeline", "WanPipeline", "WanVACEPipeline", - "LTX2ImageToVideoPipeline", - "LTX2TwoStagesPipeline", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", "LTX2T2VDMD2Pipeline", "LTX2I2VDMD2Pipeline", - "LTX23Pipeline", - "LTX23ImageToVideoPipeline", "StableAudioPipeline", "WanImageToVideoPipeline", "WanS2VPipeline", diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 9864d8dbfaf..d7770c266b1 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -634,7 +634,7 @@ def _load_model_with_hsdp( # 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) diff --git a/vllm_omni/diffusion/models/ltx2/__init__.py b/vllm_omni/diffusion/models/ltx2/__init__.py index ca34e0390c2..300295bd844 100644 --- a/vllm_omni/diffusion/models/ltx2/__init__.py +++ b/vllm_omni/diffusion/models/ltx2/__init__.py @@ -9,28 +9,16 @@ 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", diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_components.py b/vllm_omni/diffusion/models/ltx2/ltx2_components.py index cfe990a4e09..54c866eb7ce 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_components.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_components.py @@ -7,6 +7,7 @@ import inspect import json +import logging import os from dataclasses import dataclass from typing import TYPE_CHECKING, Any @@ -14,11 +15,14 @@ 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 @@ 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 @@ "audio_vae", "vocoder", "scheduler", + "latent_upsampler", ) +logger = logging.getLogger(__name__) @dataclass(frozen=True) @@ -81,6 +92,164 @@ class 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( diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py index 48d319f6dcf..6bd759057bd 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py @@ -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 @@ def encode_prompt( 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 @@ def _prepare_prompt_context( 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 @@ def __init__(self, *args, **kwargs): 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 @@ def _resolve_request_image( 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 @@ def _resolve_request_image( 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,10 +394,77 @@ def _check_forward_inputs( 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 _encode_i2v_image_latents( + self, + 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.") + if batch_size % image_batch_size != 0: + raise ValueError( + f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) " + "for image-to-video outputs." + ) + num_videos_per_prompt = batch_size // image_batch_size + + if isinstance(generator, list): + if len(generator) != batch_size: + raise ValueError( + f"You have passed a list of generators of length {len(generator)}, but requested an effective" + f" batch size of {batch_size}. Make sure the batch size matches the length of the generators." + ) + image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)] + init_latents = [ + retrieve_latents( + self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), + image_generators[i], + "argmax", + ) + for i in range(image_batch_size) + ] + else: + init_latents = [ + retrieve_latents( + self.vae.encode(img.unsqueeze(0).unsqueeze(2)), + generator, + "argmax", + ) + for img in image + ] + + init_latents = torch.cat(init_latents, dim=0).to(dtype=dtype) + if num_videos_per_prompt > 1: + init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0) + 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, @@ -328,7 +478,21 @@ def prepare_latents( device: torch.device | None = None, generator: torch.Generator | list[torch.Generator] | None = None, latents: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.Tensor]: + ) -> 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, @@ -336,13 +500,16 @@ def prepare_latents( 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: + unpacked_latents = latents.ndim == 5 if latents.ndim == 5: - batch_size, _, num_frames, height, width = latents.shape - mask_shape = (batch_size, 1, num_frames, height, width) + 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( @@ -351,11 +518,16 @@ def prepare_latents( self.vae.latents_std, self.vae.config.scaling_factor, ) - latents = latent_ops.create_noised_state( - latents, - noise_scale * (1 - conditioning_mask), - generator, - ) + 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, @@ -375,78 +547,68 @@ def prepare_latents( "Provided `latents` tensor has shape" f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}." ) - return latents.to(device=device, dtype=dtype), conditioning_mask - - if image is None: - raise ValueError("`image` must be provided when `latents` is None.") - - image_batch_size = image.shape[0] - if image_batch_size == 0: - raise ValueError("`image` batch is empty.") - if batch_size % image_batch_size != 0: - raise ValueError( - f"`batch_size` ({batch_size}) must be divisible by image batch size ({image_batch_size}) " - "for image-to-video outputs." - ) - num_videos_per_prompt = batch_size // image_batch_size - - if isinstance(generator, list): - if len(generator) != batch_size: - raise ValueError( - f"You have passed a list of generators of length {len(generator)}, but requested an effective" - f" batch size of {batch_size}. Make sure the batch size matches the length of the generators." + 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, ) - image_generators = [generator[i * num_videos_per_prompt] for i in range(image_batch_size)] - init_latents = [ - retrieve_latents( - self.vae.encode(image[i].unsqueeze(0).unsqueeze(2)), - image_generators[i], - "argmax", + elif image is not None: + image_latents = self._encode_i2v_image_latents( + image, + batch_size=batch_size, + generator=generator, + dtype=dtype, ) - for i in range(image_batch_size) - ] - else: - init_latents = [ - retrieve_latents( - self.vae.encode(img.unsqueeze(0).unsqueeze(2)), - generator, - "argmax", + unpacked = latent_ops.unpack_latents( + latents, + num_frames, + height, + width, + self.transformer_spatial_patch_size, + self.transformer_temporal_patch_size, ) - for img in image - ] + 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 - init_latents = torch.cat(init_latents, dim=0).to(dtype) - if num_videos_per_prompt > 1: - init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0) - init_latents = latent_ops.normalize_latents( - init_latents, - self.vae.latents_mean, - self.vae.latents_std, + 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 @@ def _step_video_latents_i2v( ) 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 @@ def _prepare_video_latents_stage( 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 @@ def _prepare_video_latents_stage( ) 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 @@ def _prepare_denoise_context_for_cfg( ) -> 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 @@ def _denoise_timestep_kwargs( 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) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py index d0dfa3e54e2..262c9334c45 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_denoise.py @@ -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 @@ def run( 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 @@ def run( 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 @@ def run( 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 @@ def run( 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 @@ def run( 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, + ) diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py index 395f8e7af73..0b3d2dc2bf5 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_guidance.py @@ -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 @@ 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]: ... + @property + def do_modality_guidance(self) -> bool: + return not math.isclose(self.modality_scale, 1.0) - def predict_noise( - self, - pipeline: Any, - index: int, - timestep: torch.Tensor, - state: LTXAVState, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> tuple[torch.Tensor, torch.Tensor]: ... +@dataclass(frozen=True) +class LTXGuidanceSpec: + """Video and audio guidance requested for one denoise phase.""" -class _LTXGuidanceStrategyBase: - def validate_cfg_world_size(self, cfg_world_size: int) -> None: - del cfg_world_size + video: LTXModalityGuidance = LTXModalityGuidance() + audio: LTXModalityGuidance = LTXModalityGuidance() - def prepare_denoise_context( - self, - pipeline: Any, - forward_ctx: LTXForwardContext, - denoise_ctx: LTXDenoiseContext, - ) -> LTXDenoiseContext: - del pipeline, forward_ctx - return denoise_ctx + @classmethod + def positive_only(cls) -> LTXGuidanceSpec: + return cls() - 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} + @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 -def combine_velocity_via_x0( - sample: torch.Tensor, - positive_velocity: torch.Tensor, - negative_velocity: torch.Tensor, - sigma: torch.Tensor, - guidance_scale: float, -) -> 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 + @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 -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 +@dataclass(frozen=True) +class LTXDenoisePass: + """One conditioned Transformer evaluation in a guidance batch.""" - 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, - ) + name: str + negative_video_context: bool = False + negative_audio_context: bool = False -class LTXOfficialX0Guidance(_LTXGuidanceStrategyBase): - """Official LTX x0-space CFG behavior currently used by LTX2.3.""" +@dataclass(frozen=True) +class LTXGuidancePlan: + """Concrete Transformer passes derived from a guidance specification.""" - 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}." + 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, + velocity: torch.Tensor, + sigmas: torch.Tensor, + step_index: int, +) -> torch.Tensor: + 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) + + +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) + + +def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor: + return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1)) + + +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} - - 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, - ) - audio_combined = combine_velocity_via_x0( - context["audio_latents"], - audio_pos, - audio_neg, - context["audio_sigma"], - true_cfg_scale, + 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, + } + + @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, ) - 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, + plan: LTXGuidancePlan, + index: int, + timestep: torch.Tensor, + state: LTXAVState, + forward_ctx: LTXForwardContext, + denoise_ctx: LTXDenoiseContext, ) -> 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, + 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=video_context, + audio_encoder_hidden_states=audio_context, + encoder_attention_mask=None, + audio_encoder_attention_mask=None, + ts=timestep.expand(video_input.shape[0]), + ) + 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[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[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]: - 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, + 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 ) - 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], + audio_context = ( + prompt.negative_connector_audio_prompt_embeds + if denoise_pass.negative_audio_context + else prompt.positive_connector_audio_prompt_embeds ) - - 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( + 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=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=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("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( + 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_cond, - video_uncond, + video_splits, pipeline.scheduler.sigmas[index], - guidance_scale, - ) - audio_uncond, audio_cond = audio_pred.chunk(2) - audio_pred = combine_velocity_via_x0( + plan.spec.video, + model_sigma=pipeline._video_guidance_model_sigma( + pipeline.scheduler.sigmas[index], + denoise_ctx, + ), + ), + guide_modality( state.audio, - audio_cond, - audio_uncond, - audio_scheduler.sigmas[index], - guidance_scale, - ) - return video_pred, audio_pred + audio_splits, + forward_ctx.audio_scheduler.sigmas[index], + plan.spec.audio, + ), + ) + +LTX_GUIDANCE_EXECUTOR = LTXGuidanceExecutor() -LTX_LEGACY_VELOCITY_GUIDANCE = LTXLegacyVelocityGuidance() -LTX_OFFICIAL_X0_GUIDANCE = LTXOfficialX0Guidance() +# Kept as a narrow compatibility helper for CFG-parallel tests and extensions. +combine_velocity_via_x0 = LTXGuidanceExecutor.combine_cfg_velocity diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py index 4da3edf9205..bb0155e794d 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_latents.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_latents.py @@ -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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py index 6491b4a5f29..a70f349ac9e 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_recipes.py @@ -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 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 LTXOneStageRecipe: +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) + + +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, +} -LTX2_ONE_STAGE_RECIPE = LTXOneStageRecipe() -LTX23_ONE_STAGE_RECIPE = LTXOneStageRecipe() +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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_request.py b/vllm_omni/diffusion/models/ltx2/ltx2_request.py index d8a4356b354..e2eef19df82 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_request.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_request.py @@ -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 @@ def _resolve_request_inputs( 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 @@ def _resolve_request_inputs( 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 + 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.") + 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 - 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 - ] + negative_prompt = _resolve_negative_prompts( + req.prompts, + negative_prompt if negative_prompt is not None else self.pipeline_recipe.negative_prompt, + ) - 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 + 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.one_stage_recipe.num_inference_steps + sampling.num_inference_steps or num_inference_steps or self.pipeline_recipe.num_inference_steps ) - 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.") + 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 @@ def _resolve_request_inputs( ] 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 @@ def _resolve_request_inputs( 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, diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py similarity index 59% rename from vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py rename to vllm_omni/diffusion/models/ltx2/ltx2_runtime.py index 08bf2033039..b5bb18a9b1e 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_pipeline_runtime.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_runtime.py @@ -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.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 @@ 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]: return latent_ops.prepare_audio_latents( self, @@ -172,19 +312,16 @@ def prepare_audio_latents( 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 @@ def combine_cfg_noise( ): 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 @@ def _synchronize_cfg_parallel_step_output( # 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 @@ def _decode_output( 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 @@ def _prepare_video_latents_stage( 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 @@ def _prepare_audio_latents_stage( 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 @@ def _resolve_audio_latent_length( 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 @@ 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]: return build_transformer_kwargs( self, @@ -487,6 +651,7 @@ def _build_transformer_kwargs( 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 @@ def _predict_noise_for_step( 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 @@ def _unpack_and_denormalize_stage( ) 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 @@ def run_phase( 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 @@ def run_phase( 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 diff --git a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py index 3d1a2d66324..95724bd19a0 100644 --- a/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py +++ b/vllm_omni/diffusion/models/ltx2/ltx2_transformer.py @@ -57,6 +57,9 @@ 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 @@ def forward(self, x: torch.Tensor) -> torch.Tensor: 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 @@ def __init__(self): "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 @@ def _prepare_attention_mask( # 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 @@ def __call__( 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 @@ def __call__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def __init__( 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 @@ def forward( 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 @@ def forward( 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 @@ def forward( 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 @@ def forward( 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 @@ def forward( 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 @@ def prepare_coords(self, *args, **kwargs): 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 @@ def forward( 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 @@ def _build_sp_plan(rope_type: str) -> dict[str, Any]: # 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 @@ def forward( # 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 @@ def forward( 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 @@ def forward( 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 @@ def forward( 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 @@ def forward( 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 @@ def forward( 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 @@ def forward( "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 @@ def _maybe_shard_weight(weight: torch.Tensor, param: torch.Tensor) -> torch.Tens 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) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py index 5de1f4862d1..ff66af5e17b 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py @@ -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.""" - component_profile: ClassVar[LTXComponentProfile] - one_stage_recipe: ClassVar[LTXOneStageRecipe] +class LTX2Pipeline(LTXI2VConditioningMixin, LTXRuntime): + """LTX-2 family one-stage entry, configured from checkpoint metadata.""" - 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 - _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 + 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) - - @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__() diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py deleted file mode 100644 index 0c72a41d5ea..00000000000 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py +++ /dev/null @@ -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,)) diff --git a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py index 8b786b8175e..0e5bbc5583e 100644 --- a/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py +++ b/vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py @@ -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 - -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 vllm_omni.diffusion.data import OmniDiffusionConfig -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, - ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index a0b1e9a7b19..aaedc7b00fc 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -71,20 +71,10 @@ "pipeline_ltx2", "LTX2Pipeline", ), - "LTX2ImageToVideoPipeline": ( - "ltx2", - "pipeline_ltx2", - "LTX2ImageToVideoPipeline", - ), - "LTX2TwoStagesPipeline": ( - "ltx2", - "pipeline_ltx2_two_stage", - "LTX2TwoStagesPipeline", - ), - "LTX2ImageToVideoTwoStagesPipeline": ( + "LTX2DistilledPipeline": ( "ltx2", "pipeline_ltx2_two_stage", - "LTX2ImageToVideoTwoStagesPipeline", + "LTX2DistilledPipeline", ), "LTX2T2VDMD2Pipeline": ( "ltx2", @@ -96,16 +86,6 @@ "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 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) - "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", diff --git a/vllm_omni/diffusion/worker/request_batch.py b/vllm_omni/diffusion/worker/request_batch.py index 488d1eec2e0..079a213dbf2 100644 --- a/vllm_omni/diffusion/worker/request_batch.py +++ b/vllm_omni/diffusion/worker/request_batch.py @@ -8,6 +8,7 @@ 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 diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index faec5f0b1b1..13ee808be06 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -410,7 +410,7 @@ def subparser_init(self, subparsers: argparse._SubParsersAction) -> TrackingArgu 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", diff --git a/vllm_omni/model_extras/ltx2.py b/vllm_omni/model_extras/ltx2.py new file mode 100644 index 00000000000..5bac584690f --- /dev/null +++ b/vllm_omni/model_extras/ltx2.py @@ -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() diff --git a/vllm_omni/model_extras/registry.py b/vllm_omni/model_extras/registry.py index edf28101caa..dab5b45cf64 100644 --- a/vllm_omni/model_extras/registry.py +++ b/vllm_omni/model_extras/registry.py @@ -37,6 +37,7 @@ ) 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 @@ def default_image_to_video_prompt( "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,