Unskip test for invalid logprobs handling

Remove skip marker from test for invalid logprobs.

Signed-off-by: Gao Han <hgaoaf@connect.ust.hk>

.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
 

.buildkite/npu/test-npu-nightly.yml

@@ -1,25 +1,59 @@
 steps:
-  - group: ":card_index_dividers: Diffusion X2V Model Test"
-    key: nightly-diffusion-x2v-group
+  - group: ":card_index_dividers: Omni Model Test"
+    key: nightly-omni-test-group
     depends_on: upload-nightly-pipeline
     if: >-
       build.env("NIGHTLY") == "1" ||
       build.pull_request.labels includes "nightly-test" ||
-      build.pull_request.labels includes "diffusion-x2v-test"
+      build.pull_request.labels includes "omni-test"
     steps:
-      - label: ":full_moon: Diffusion X2V · Function Test"
-        mirror_hardwares: a2b3_npu_4
+      - label: ":full_moon: Omni · Perf Test · No Async Chunk"
+        key: nightly-omni-performance-no-async-chunk
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
         env:
           VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
           HF_TOKEN: "${HF_TOKEN}"
         commands:
-          - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -k npu_i2v --run-level "full_model"
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+      - label: ":full_moon: Omni · Perf Test · Async Chunk"
+        key: nightly-omni-performance-async-chunk
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+        commands:
+          - export BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+            exit $$EXIT
+      - label: ":robot: Omni · MiniCPM-o 4.5 Test"
+        mirror_hardwares: a3_npu_2
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+        commands:
+          - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5_expansion.py -m 'full_model' --run-level 'full_model'
   - group: ":card_index_dividers: TTS Model Test"
     key: nightly-tts-test-group
     depends_on: upload-nightly-pipeline
     if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
     steps:
       - label: ":full_moon: TTS · Perf Test"
+        timeout_in_minutes: 180
         mirror_hardwares: a3_npu_2
         env:
           VLLM_WORKER_MULTIPROC_METHOD: spawn
@@ -34,25 +68,56 @@ steps:
             EXIT=$$?
             buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
             exit $$EXIT
-
-  - group: ":card_index_dividers: Omni Model Test"
-    key: nightly-omni-group
+  - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
+    key: nightly-diffusion-x2iat-group
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test"
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2iat-test"
     steps:
-      - label: ":robot: Omni · MiniCPM-o 4.5 Test"
-        agents:
-          queue: "ascend-a3"
-          resource_class: "npu-2"
-        image: "swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/vllm-omni-ci-npu:${BUILDKITE_COMMIT}-a3"
-        plugins:
-          - kubernetes:
-              podSpecPatch:
-                imagePullSecrets:
-                  - name: swr-secret
+      - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · Qwen-Image-Edit-2511"
+        key: nightly-diffusion-x2iat-performance-qwen-image-edit-2511
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_8
         env:
           VLLM_WORKER_MULTIPROC_METHOD: spawn
           VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
           HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
         commands:
-          - pytest -s -v tests/e2e/offline_inference/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5.py tests/e2e/online_serving/test_minicpmo_4_5_expansion.py -m 'full_model' --run-level 'full_model'
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - export CACHE_DIT_VERSION=1.3.0
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT
+  - group: ":card_index_dividers: Diffusion X2V Model Test"
+    key: nightly-diffusion-x2v-group
+    depends_on: upload-nightly-pipeline
+    if: >-
+      build.env("NIGHTLY") == "1" ||
+      build.pull_request.labels includes "nightly-test" ||
+      build.pull_request.labels includes "diffusion-x2v-test"
+    steps:
+      - label: ":full_moon: Diffusion X2V · Perf Test"
+        key: nightly-diffusion-x2v-performance
+        timeout_in_minutes: 180
+        mirror_hardwares: a3_npu_4
+        env:
+          VLLM_WORKER_MULTIPROC_METHOD: spawn
+          VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
+          HF_TOKEN: "${HF_TOKEN}"
+          DIFFUSION_ATTENTION_BACKEND: TORCH_SDPA
+        commands:
+          - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+          - |
+            set +e
+            pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py -m npu --test-config-file tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json
+            EXIT=$$?
+            buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
+            buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+            exit $$EXIT

.buildkite/npu/test-npu-ready.yml

@@ -31,4 +31,4 @@ steps:
           VLLM_ALLOW_LONG_MAX_MODEL_LEN: "1"
           HF_TOKEN: "${HF_TOKEN}"
         commands:
-          - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
+          - pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model and npu' --run-level 'core_model'

docker/Dockerfile.npu.ci

@@ -37,6 +37,10 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "pyttsx3>=2.99" \
       "opencv-python>=4.12.0.88" \
       "imageio[ffmpeg]>=0.6.0" \
+      "jiwer>=3.0.0" \
+      "zhon>=2.0.0" \
+      "zhconv>=1.4.2" \
+      "funasr>=1.0.0" \
       "hyperpyyaml>=1.2.0" \
       "onnx>=1.16.0" \
       "librosa>=0.10.0" \

docker/Dockerfile.npu.ci.a3

@@ -37,6 +37,10 @@ RUN source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
       "pyttsx3>=2.99" \
       "opencv-python>=4.12.0.88" \
       "imageio[ffmpeg]>=0.6.0" \
+      "jiwer>=3.0.0" \
+      "zhon>=2.0.0" \
+      "zhconv>=1.4.2" \
+      "funasr>=1.0.0" \
       "hyperpyyaml>=1.2.0" \
       "onnx>=1.16.0" \
       "librosa>=0.10.0" \

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` | ✅︎ | ✅︎ | ✅︎ | |

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 |

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -4,7 +4,7 @@ Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inf
 
 
 This example demonstrates how to generate videos from images using Wan2.2,
-LTX-2/LTX-2.3, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
+LTX-2, and HunyuanVideo-1.5 Image-to-Video models with vLLM-Omni's
 offline inference API.
 
 ## Local CLI Usage
@@ -53,34 +53,22 @@ python image_to_video.py \
   --output i2v_output.mp4
 ```
 
-### LTX-2.3
+### LTX-2
 
 ```bash
 python image_to_video.py \
-  --model diffusers/LTX-2.3-Diffusers \
-  --model-class-name LTX23ImageToVideoPipeline \
+  --model Lightricks/LTX-2 \
   --image cherry_blossom.jpg \
   --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
-  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
-  --height 384 \
-  --width 512 \
-  --num-frames 25 \
-  --guidance-scale 4.0 \
-  --num-inference-steps 20 \
-  --frame-rate 24 \
-  --fps 24 \
-  --output ltx23_i2v_output.mp4
+  --output ltx2_i2v_output.mp4
 ```
 
-Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
-upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
-this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
-the LTX-2.3 image-to-video pipeline.
+See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
+pipeline selection, T2V, defaults, and advanced options.
 
 Key arguments:
 
-- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or
-  LTX-2/LTX-2.3).
+- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, or LTX-2).
 - `--image`: Path to input image (required).
 - `--prompt`: Text description of desired motion/animation.
 - `--height/--width`: Output resolution (auto-calculated from image if not set).

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -12,7 +12,7 @@ For backend selection and SageAttention usage, see the [Diffusion Attention Back
 | Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
-| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
+| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 
@@ -35,45 +35,17 @@ python text_to_video.py \
   --output t2v_out.mp4
 ```
 
-LTX2 example:
+### LTX-2
 
 ```bash
 python text_to_video.py \
-  --model "Lightricks/LTX-2" \
-  --prompt "A cinematic close-up of ocean waves at golden hour." \
-  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
-  --height 512 \
-  --width 768 \
-  --num-frames 121 \
-  --num-inference-steps 40 \
-  --guidance-scale 4.0 \
-  --frame-rate 24 \
-  --output ltx2_out.mp4
-```
-
-### LTX-2.3
-
-```bash
-python text_to_video.py \
-  --model diffusers/LTX-2.3-Diffusers \
-  --model-class-name LTX23Pipeline \
+  --model Lightricks/LTX-2 \
   --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
-  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
-  --height 384 \
-  --width 512 \
-  --num-frames 25 \
-  --num-inference-steps 20 \
-  --guidance-scale 4.0 \
-  --frame-rate 24 \
-  --fps 24 \
-  --audio-sample-rate 48000 \
-  --output ltx23_t2v_output.mp4
+  --output ltx2_output.mp4
 ```
 
-Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
-upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
-this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
-text-to-video pipeline explicitly.
+See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
+pipeline selection, I2V, defaults, and advanced options.
 
 ### HunyuanVideo-1.5 (480p)
 
@@ -135,8 +107,7 @@ python text_to_video.py \
 ### Common
 
 - `--model`: Diffusers model ID or local path.
-- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
-  for LTX-2.3 text-to-video.
+- `--model-class-name`: Optional explicit pipeline override.
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution. Default depends on model.
 - `--num-frames`: number of frames. Default depends on model.
@@ -151,8 +122,7 @@ python text_to_video.py \
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
-- `--audio-sample-rate`: audio sample rate for embedded audio (when the
-  pipeline returns audio; LTX-2.3 outputs 48kHz audio).
+- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
 - `--flow-shift`: scheduler flow_shift parameter.
 

docs/user_guide/examples/online_serving/text_to_video.md

@@ -13,7 +13,6 @@ This example demonstrates how to deploy text-to-video models for online video ge
 | Wan2.1 T2V (14B) | `Wan-AI/Wan2.1-T2V-14B-Diffusers` |
 | Wan2.2 T2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` |
 | LTX-2 | `Lightricks/LTX-2` |
-| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` |
 
 ## Wan2.2 T2V
 
@@ -192,7 +191,7 @@ curl -X POST http://localhost:8091/v1/videos \
 | `negative_prompt`     | str    | None    | Negative prompt                                  |
 | `width`               | int    | None    | Video width in pixels                            |
 | `height`              | int    | None    | Video height in pixels                           |
-| `num_frames`          | int    | None    | Number of frames to generate                     |
+| `num_frames`          | int    | 1       | Number of frames to generate; set explicitly for video generation |
 | `fps`                 | int    | None    | Frames per second for output video               |
 | `num_inference_steps` | int    | None    | Number of denoising steps                        |
 | `guidance_scale`      | float  | None    | CFG guidance scale (low-noise stage)             |
@@ -290,101 +289,12 @@ done
 
 ## LTX-2
 
-### Start Server
-
-#### Basic Start
-
 ```bash
-vllm serve Lightricks/LTX-2 --omni --port 8098 \
-    --enforce-eager --boundary-ratio 1.0
+vllm serve Lightricks/LTX-2 --omni --port 8098
 ```
 
-For multi-GPU memory reduction, you can enable HSDP:
-
-```bash
-vllm serve Lightricks/LTX-2 --omni --port 8098 \
-    --enforce-eager --boundary-ratio 1.0 \
-    --use-hsdp --hsdp-shard-size 2
-```
-
-#### Start with Optimization Presets
-
-Use the LTX-2 startup script with built-in optimization presets:
-
-```bash
-# Baseline (1 GPU, eager)
-bash run_server_ltx2.sh baseline
-
-# 4-GPU Ulysses sequence parallelism (lossless)
-bash run_server_ltx2.sh ulysses4
-
-# Cache-DiT lossy acceleration (1 GPU, ~1.4× speedup)
-bash run_server_ltx2.sh cache-dit
-
-# Best combo: 4-GPU Ulysses SP + Cache-DiT (~2.2× speedup)
-bash run_server_ltx2.sh best-combo
-```
-
-#### Optimization Benchmarks
-
-Benchmarked on H800, online serving (480×768, 41 frames, 20 steps, `seed=42`).
-"Inference" is the server-reported inference time; excludes HTTP/poll overhead.
-
-| Preset | Server Command | Inference (s) | Speedup | Type |
-|--------|---------------|---------------|---------|------|
-| `baseline` | `--enforce-eager` | 10.3 | 1.00× | — |
-| `compile` | *(default, no --enforce-eager)* | ~10.3 (warm) | ~1.00× | Lossless |
-| `ulysses4` | `--enforce-eager --usp 4` | ~10.3 | ~1.00× | Lossless |
-| `cache-dit` | `--enforce-eager --cache-backend cache_dit` | 7.4 avg | ~1.4× | Lossy |
-| `best-combo` | `--enforce-eager --usp 4 --cache-backend cache_dit` | 4.7 avg | **~2.2×** | Lossless + Lossy |
-
-**Observations**:
-- **torch.compile**: On H800, warm-request inference time matches the eager baseline (~10.3s).
-  The first request pays ~6s compilation overhead. Benefit depends on model architecture and GPU.
-- **Ulysses SP (4 GPU)**: No measurable speedup alone for 41-frame generation at this resolution.
-  Communication overhead outweighs gains at this sequence length.
-- **Cache-DiT**: Inference varies per request (6–10s) due to dynamic caching decisions.
-  Average is ~7.4s (~1.4× speedup) with slight quality tradeoff.
-- **Best combo**: 4-GPU Ulysses SP + Cache-DiT synergize well — Cache-DiT reduces per-step
-  computation, making the communication overhead of Ulysses SP worthwhile. Average ~4.7s
-  (~2.2× speedup).
-- **FP8 quantization**: Reduces VRAM but does not speed up LTX-2 on H800 (compute-bound).
-
-**Deployment Recommendations**:
-- For **production with quality priority**: use `baseline` with `--enforce-eager`
-- For **maximum throughput** (4 GPUs, quality tradeoff): use `best-combo` (~2.2× speedup)
-- For **single-GPU throughput**: use `cache-dit` (~1.4× speedup)
-- `--enforce-eager` is recommended to avoid torch.compile warmup latency on first request
-
-### Send Requests (curl)
-
-```bash
-# Using the provided script
-bash run_curl_ltx2.sh
-
-# Or directly
-curl -sS -X POST http://localhost:8098/v1/videos \
-  -H "Accept: application/json" \
-  -F "prompt=A serene lakeside sunrise with mist over the water." \
-  -F "width=768" \
-  -F "height=480" \
-  -F "num_frames=41" \
-  -F "fps=24" \
-  -F "num_inference_steps=20" \
-  -F "guidance_scale=3.0" \
-  -F "seed=42"
-```
-
-## LTX-2.3
-
-LTX-2.3 uses a Diffusers-format checkpoint and a dedicated pipeline class:
-
-```bash
-vllm serve diffusers/LTX-2.3-Diffusers --omni --port 8098 \
-    --model-class-name LTX23Pipeline
-```
-
-Send requests through the same `/v1/videos` API shown in the LTX-2 example.
+See the [LTX-2 recipe](../../../../recipes/LTX/LTX-2.md) for all checkpoints,
+pipeline selection, requests, defaults, and advanced options.
 
 ## Example materials
 

examples/offline_inference/audex/README.md

@@ -0,0 +1,114 @@
+# Audex (Nemotron-Labs-Audex-2B) offline inference
+
+One checkpoint, four tasks — each script builds its own `Omni` engine with
+the right deploy yaml, so they run as plain `python <script>` with no extra
+flags. Pass the HF repo ROOT as `--model` (default
+`nvidia/Nemotron-Labs-Audex-2B`); per-stage subfolders and download subsets
+resolve automatically.
+
+| script | pipeline (`vllm_omni/deploy/<name>.yaml`) | audio in | text out | speech out | general audio out |
+|---|---|---|---|---|---|
+| `text_to_speech.py` | `audex_tts` | ❌ | ❌ | ✅ | ❌ |
+| `text_to_audio.py` | `audex_tta` | ❌ | ❌ | ❌ | ✅ |
+| `audio_qa.py` | `audex_thinker_only` | ✅ | ✅ | ❌ | ❌ |
+| `speech_to_speech.py` | `audex_s2s` | ✅ | ✅ | ✅ | ❌ |
+
+Shared plumbing (corpus/WAV IO, tokenizer loading, the CFG cond/uncond
+pair contract) lives in `common.py`; each task script keeps only its own
+prompt recipe and flow. Note `text_to_audio.py` here drives the Audex
+two-stage AR pipeline — for diffusion text-to-audio models (Stable Audio
+Open, AudioX) use the generic
+`examples/offline_inference/text_to_audio/text_to_audio.py` instead.
+
+## text_to_speech.py — text → speech
+
+Thinker generates `<speechcodec_N>` tokens; the streaming causal decoder
+turns them into 16 kHz WAVs (one file per prompt, slugified filename).
+
+```bash
+# three built-in demo sentences -> results/audex_wavs/
+python examples/offline_inference/audex/text_to_speech.py
+
+# your own texts, with classifier-free guidance (official quality setting)
+python examples/offline_inference/audex/text_to_speech.py \
+    --texts "Hello there." "Nice to meet you." --cfg-scale 1.5 \
+    --output-dir results/my_tts
+```
+
+Key flags: `--texts` / `--texts-file` (TSV `utt_id<TAB>text`),
+`--cfg-scale` (default 1.0 = off; 1.5 recommended — guided requests run
+one at a time, each paired with a length-matched null prompt),
+`--temperature` (overrides the deploy yaml's stage-0 default).
+
+## text_to_audio.py — caption → general audio (sound effects)
+
+Same thinker over the interleaved 4-codebook `<audiocodec_N>` RVQ block,
+decoded by the external XCodec1 checkpoint. CFG is effectively mandatory
+(default scale 3.0). Clips are capped at `--codec-cap 4000` codec tokens
+(10 s).
+
+```bash
+# three built-in demo captions -> results/audex_tta_wavs/
+python examples/offline_inference/audex/text_to_audio.py
+
+python examples/offline_inference/audex/text_to_audio.py \
+    --captions "Thunder rolling across a valley." --output-dir results/my_tta
+```
+
+XCodec1 resolves from `--xcodec1-path`, the `XCODEC1_PATH` env var, or the
+default `hf-audio/xcodec-hubert-general-balanced` repo (downloaded on
+first use).
+
+## audio_qa.py — audio (+ instruction) → text
+
+Single-stage audio understanding on the full checkpoint (NV-Whisper
+encoder + projector + LM). The default question is ASR
+("Transcribe the input speech."); without `--audio-files`, vLLM's public
+`mary_had_lamb` asset is transcribed.
+
+```bash
+# transcribe the built-in demo asset
+python examples/offline_inference/audex/audio_qa.py
+
+# your own clips / free-form QA
+python examples/offline_inference/audex/audio_qa.py \
+    --audio-files a.wav b.wav --question "What language is being spoken?"
+```
+
+## speech_to_speech.py — spoken question → spoken answer
+
+The official three-pass cascade over ONE `audex_s2s` deployment: ASR
+(audio → transcript, text-final) → chat (transcript → answer, text-final)
+→ TTS (answer → speech through the streaming decoder). Only the TTS pass
+touches the codec path. Without `--audio-file`, the `mary_had_lamb` asset
+is used as the spoken input.
+
+```bash
+# full cascade on the built-in demo asset -> results/audex_s2s_answer.wav
+python examples/offline_inference/audex/speech_to_speech.py
+
+python examples/offline_inference/audex/speech_to_speech.py \
+    --audio-file question.wav --output results/answer.wav --cfg-scale 1.5
+```
+
+`--cfg-scale` (default 1.5) applies to the TTS pass only; 1.0 disables it.
+
+Online-serving counterparts (server + HTTP clients) live in
+`examples/online_serving/audex/`.
+
+## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
+
+The same four scripts serve the 30B MoE checkpoint (hybrid Mamba + MoE
+NemotronH thinker; identical token space and decoder). The 30B REQUIRES an
+explicit 30B deploy yaml — the model root's default resolution lands on the
+2B-tuned configs:
+
+    python examples/offline_inference/audex/text_to_speech.py \
+        --model nvidia/Nemotron-Labs-Audex-30B-A3B \
+        --deploy-config vllm_omni/deploy/audex_tts_30b.yaml
+
+Per-mode yamls: `audex_{tts,tta,thinker_only,s2s}_30b.yaml`. Defaults fit a
+single H100 80 GB (thinker + decoder share the card; prefix caching is off —
+hybrid Mamba); if long sequences OOM, set `tensor_parallel_size: 2` on
+stage 0. First run downloads ~60 GB. Verified on a single H100 80 GB: ~61 GiB
+weights + healthy KV headroom; TP2 was not needed.

examples/offline_inference/audex/audio_qa.py

@@ -0,0 +1,93 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) audio understanding example.
+
+Speech (or general audio) + text instruction → text, through the
+single-stage ``audex_thinker_only`` pipeline: NV-Whisper
+encoder + relu2 projector + 2B dense LM on ``checkpoint_folder_full``.
+
+The prompt format mirrors the official audioqa script: a ChatML user turn
+holding the ``<so_embedding>`` placeholder (expanded by the processor to
+750 embedding positions per 30 s clip) plus the instruction, with a closed
+``<think></think>`` priming.
+
+Examples (without --audio-files, vLLM's public ``mary_had_lamb`` asset is
+transcribed):
+
+    # Transcribe WAVs (ASR):
+    python examples/offline_inference/audex/audio_qa.py \\
+        --audio-files a.wav b.wav
+
+    # Free-form audio QA:
+    python examples/offline_inference/audex/audio_qa.py \\
+        --audio-files a.wav --question "What language is being spoken?"
+"""
+
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from common import default_deploy_config, load_audio_file, request_index
+from vllm.assets.audio import AudioAsset
+
+from vllm_omni import Omni
+
+ASR_QUESTION = "Transcribe the input speech."
+
+
+def build_prompt(question: str) -> str:
+    return f"<|im_start|>user\n<so_embedding>\n{question}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex audio understanding")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument(
+        "--audio-files",
+        type=str,
+        nargs="+",
+        default=None,
+        help="Input clips. Defaults to vLLM's mary_had_lamb audio asset.",
+    )
+    parser.add_argument("--question", type=str, default=ASR_QUESTION)
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS pipeline; audio
+        # understanding needs the single-stage thinker-only pipeline.
+        default=default_deploy_config("audex_thinker_only.yaml"),
+        help="Deploy yaml (defaults to the audex_thinker_only pipeline).",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    prompt_text = build_prompt(args.question)
+    if args.audio_files:
+        clips = [load_audio_file(path) for path in args.audio_files]
+    else:
+        clips = [AudioAsset("mary_had_lamb").audio_and_sample_rate]
+    prompts = []
+    for audio, sr in clips:
+        prompts.append(
+            {
+                "prompt": prompt_text,
+                "multi_modal_data": {"audio": (audio, sr)},
+            }
+        )
+
+    outputs = engine.generate(prompts)
+
+    ordered = sorted(outputs, key=request_index)
+    names = [Path(p).name for p in args.audio_files] if args.audio_files else ["mary_had_lamb"]
+    for name, req_output in zip(names, ordered):
+        text = req_output.outputs[0].text if req_output.outputs else ""
+        print(f"{name}: {text.strip()}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/common.py

@@ -0,0 +1,149 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Shared helpers for the offline Audex examples in this directory.
+
+Import style follows the repo's example convention (same-directory module,
+see ``text_to_speech/qwen3_tts/tts_common.py``): scripts run as files, so
+``from common import ...`` resolves via the script's own directory.
+"""
+
+from __future__ import annotations
+
+import copy
+import re
+from collections.abc import Callable
+from pathlib import Path
+from typing import Any
+
+import numpy as np
+import soundfile as sf
+import torch
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
+
+SAMPLE_RATE = 16_000
+_DEPLOY_DIR = Path(__file__).resolve().parents[3] / "vllm_omni" / "deploy"
+
+# Guided decoding sharpens the distribution, so the unguided temperature
+# (0.1) adds excess sampling noise under CFG. Measured on the en-24 gate
+# corpus at cfg 1.5: temp 0.1 -> CER 7.31%, temp 0.05 -> CER 6.87% (vs the
+# unguided baseline 7.24%).
+GUIDED_TEMPERATURE = 0.05
+
+
+def default_deploy_config(name: str) -> str:
+    """Absolute path of a deploy yaml bundled with vllm_omni."""
+    return str(_DEPLOY_DIR / name)
+
+
+def slugify(text: str, fallback: str = "prompt") -> str:
+    slug = re.sub(r"\s+", "_", text.strip().lower())
+    slug = re.sub(r"[^a-z0-9_]+", "", slug)
+    return slug[:48] or fallback
+
+
+def load_corpus(
+    corpus_file: str | None,
+    texts: list[str] | None,
+    default_texts: tuple[str, ...],
+) -> list[tuple[str, str]]:
+    """(utt_id, text) pairs from a TSV corpus file, CLI texts, or defaults."""
+    if corpus_file:
+        corpus = []
+        for line in Path(corpus_file).read_text().splitlines():
+            if not line.strip():
+                continue
+            utt, text = line.split("\t", 1)
+            corpus.append((utt, text.strip()))
+        return corpus
+    resolved = texts if texts else list(default_texts)
+    return [(slugify(t), t) for t in resolved]
+
+
+def load_audio_file(path: str) -> tuple[np.ndarray, int]:
+    """Mono float32 samples + native sample rate (processors resample)."""
+    audio, sr = sf.read(path, dtype="float32")
+    if audio.ndim == 2:
+        audio = audio.mean(axis=1)
+    return audio, sr
+
+
+def extract_pcm(multimodal_output: dict) -> torch.Tensor:
+    """Flatten one request's audio payload to a 1-D float32 CPU tensor."""
+    audio = multimodal_output.get("model_outputs")
+    if audio is None:
+        audio = multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
+    if isinstance(audio, list):
+        valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
+        if not valid:
+            raise ValueError("audio list is empty")
+        return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
+    return torch.as_tensor(audio).float().cpu().reshape(-1)
+
+
+def write_wav(path: Path, pcm: torch.Tensor, sample_rate: int = SAMPLE_RATE) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    arr = (np.clip(pcm.numpy(), -1.0, 1.0) * 32767.0).astype(np.int16)
+    sf.write(str(path), arr, sample_rate, format="WAV", subtype="PCM_16")
+
+
+def load_audex_tokenizer(model: str, profile: str = "tts", folder: str = "checkpoint_folder_audiogen"):
+    """Thinker tokenizer for CFG/TTA prompt building (resolves the snapshot)."""
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+
+    root = ensure_audex_snapshot(model, profile=profile)
+    return AutoTokenizer.from_pretrained(str(Path(root) / folder), trust_remote_code=True)
+
+
+def request_index(req_output) -> int:
+    """Submission index from the numeric request-id prefix.
+
+    ``Omni.generate`` may return outputs out of submission order, and a
+    lexicographic sort would misplace "10_..." before "2_...".
+    """
+    match = re.search(r"(\d+)", str(req_output.request_id))
+    return int(match.group(1)) if match else 0
+
+
+def guided_stage_params(
+    engine: Omni,
+    cond_prompt: str,
+    tokenizer,
+    cfg_scale: float,
+    pair_id: str,
+    *,
+    temperature: float | None = None,
+    extra_args: dict[str, Any] | None = None,
+    null_prompt_builder: Callable[[str, Any], str] = build_null_prompt,
+) -> list:
+    """Stage sampling params carrying the CFG pair contract for ONE request.
+
+    Clones the engine's resolved defaults (never mutates the shared list),
+    applies ``extra_args`` unconditionally (e.g. the TTA RVQ contract), and
+    attaches the cond/uncond pair contract when ``cfg_scale > 1.0``. Guided
+    requests must go one at a time: each needs its own pair id and a
+    length-matched null prompt built from its own text.
+    """
+    params = copy.deepcopy(engine.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    if temperature is not None:
+        stage0.temperature = temperature
+    if stage0.extra_args is None:
+        stage0.extra_args = {}
+    if extra_args:
+        stage0.extra_args.update(extra_args)
+    if cfg_scale > 1.0:
+        stage0.extra_args.update(
+            {
+                "cfg_scale": float(cfg_scale),
+                "cfg_role": "cond",
+                "cfg_pair_id": pair_id,
+                "cfg_null_prompt": null_prompt_builder(cond_prompt, tokenizer),
+            }
+        )
+    return params

examples/offline_inference/audex/speech_to_speech.py

@@ -0,0 +1,277 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) cascaded speech-to-speech example.
+
+The official three-pass cascade over ONE ``audex_s2s``
+deployment:
+
+  1. ASR pass  (audio + "Transcribe the input speech." → transcript, text)
+  2. chat pass (transcript → answer, text)
+  3. TTS pass  (answer → speech, audio via the streaming causal decoder)
+
+Text passes carry ``modalities: ["text"]`` and finish at stage 0; only the
+TTS pass (``modalities: ["audio"]``) streams codec frames into code2wav.
+
+Example (without --audio-file, vLLM's public ``mary_had_lamb`` asset is
+used as the spoken input):
+
+    python examples/offline_inference/audex/speech_to_speech.py \\
+        --audio-file question.wav --output results/answer.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import copy
+import re
+from pathlib import Path
+
+import torch
+from common import (
+    SAMPLE_RATE,
+    default_deploy_config,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_audio_file,
+    write_wav,
+)
+from vllm.assets.audio import AudioAsset
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+ASR_QUESTION = "Transcribe the input speech."
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
+# prefix, open <think> priming, temp 1.0/top_p 0.95, seed retry) applies to
+# the 30B-A3B ONLY (see is_30b()/_chat_prompt below and
+# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
+# split). It is load-bearing on the 30B: without the formatting-instruction
+# prefix a bare transcript turn is treated as a TTS request and the model
+# emits literal <speechcodec_N> tokens instead of a text answer (the system
+# prompt alone does not prevent this). The 2B must NOT use this recipe: it
+# is long-verified stable on a bare prompt + engine-default sampling, while
+# the official temp-1.0 seed-0 draw hallucinates off-language text (yet
+# closes </think> cleanly, so the seed retry never fires) and fails the
+# round-trip WER gate — measured, reproduced byte-identically.
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+
+
+def is_30b(model: str) -> bool:
+    # Users pass arbitrary repo IDs or local snapshot paths, so this is a
+    # heuristic (not an exact match against the official 30B-A3B repo ID)
+    # — good enough for an example script; keep in sync with the same
+    # branch in examples/online_serving/audex/client.py and
+    # tests/e2e/offline_inference/test_audex_s2s.py.
+    return "30b" in model.lower()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex cascaded S2S")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument(
+        "--audio-file",
+        type=str,
+        default=None,
+        help="Spoken question (WAV). Defaults to vLLM's mary_had_lamb audio asset.",
+    )
+    parser.add_argument("--output", type=str, default="results/audex_s2s_answer.wav")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS-only pipeline; the
+        # cascade needs the audio-capable 2-stage full pipeline.
+        default=default_deploy_config("audex_s2s.yaml"),
+        help="Deploy yaml (defaults to the audex_s2s pipeline).",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=1.5,
+        help="CFG strength for the TTS pass (1.0 disables; official setting 1.5).",
+    )
+    return parser.parse_args()
+
+
+def _chat_prompt(user_text: str) -> str:
+    # Official chat-pass recipe: system prompt + formatting instruction
+    # prefixed to the user text, thinking ENABLED via an open <think> block
+    # (the official demo's web UI ships "Enable reasoning" checked, i.e.
+    # build_chat_prompt(..., enable_thinking=True); chat_template.jinja then
+    # renders "<|im_start|>assistant\n<think>\n"). The instruction invites
+    # the model to reason before answering — priming the block open keeps
+    # that reasoning inside <think>...</think> so _extract_final_response
+    # below can drop it; with a closed block the 30B leaks the reasoning as
+    # plain text in front of the answer.
+    return (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+        f"<|im_start|>user\n{TEXT_INSTRUCTION}\n\n{user_text}<|im_end|>\n"
+        "<|im_start|>assistant\n<think>\n"
+    )
+
+
+def _asr_prompt() -> str:
+    return f"<|im_start|>user\n<so_embedding>\n{ASR_QUESTION}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer
+    # (cascaded_s2s_web_server.py): keep only what follows the LAST closing
+    # think tag, falling back to the full text when the model skipped
+    # reasoning or never closed the block. Benign no-op for the 2B, whose
+    # answers carry no think block.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _clean(text: str) -> str:
+    # The full checkpoint prefixes transcripts with a language tag sentence
+    # and quotes the content; extract the quoted payload when present. Span
+    # first-to-last quote (not split on every apostrophe) so contractions
+    # inside the transcript (e.g. "don't") survive intact.
+    text = text.strip()
+    match = re.search(r"'(.*)'", text, flags=re.S)
+    if match and match.group(1).strip():
+        return match.group(1).strip()
+    return text
+
+
+def _text_of(outputs) -> str:
+    (req_output,) = outputs
+    return (req_output.outputs[0].text or "").strip()
+
+
+CHAT_MAX_TOKENS = 2048  # official --text-max-tokens default
+# The official demo samples this pass unseeded; some trajectories deliberate
+# past the token budget without ever closing the think block, so retry a few
+# deterministic seeds until one closes (see the chat pass below).
+CHAT_SEEDS = (0, 1, 2)
+
+
+def _chat_params(engine: Omni, seed: int):
+    # Official demo text-pass sampling (cascaded_s2s_web_server.py defaults:
+    # --text-temperature 1.0, --text-top-p 0.95, --text-max-tokens 2048).
+    # The yaml stage-0 defaults (temperature 0.1, top_k 80) are TTS-thinker
+    # settings; under near-greedy decoding the 30B loops inside its
+    # reasoning block until the token budget runs out and never emits
+    # </think>, so the extracted "answer" would be truncated reasoning.
+    params = copy.deepcopy(engine.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    stage0.temperature = 1.0
+    stage0.top_p = 0.95
+    stage0.top_k = -1
+    stage0.max_tokens = CHAT_MAX_TOKENS
+    stage0.seed = seed
+    return params
+
+
+def _chat_once(engine: Omni, prompt: str, seed: int) -> tuple[str, bool]:
+    (req_output,) = engine.generate([{"prompt": prompt, "modalities": ["text"]}], _chat_params(engine, seed))
+    out = req_output.outputs[0]
+    text = (out.text or "").strip()
+    truncated = getattr(out, "finish_reason", None) == "length" or len(out.token_ids or []) >= CHAT_MAX_TOKENS
+    return text, truncated
+
+
+def main():
+    args = parse_args()
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    tokenizer = None
+    if args.cfg_scale > 1.0:
+        tokenizer = load_audex_tokenizer(args.model, profile="full", folder="checkpoint_folder_full")
+
+    if args.audio_file:
+        audio, sr = load_audio_file(args.audio_file)
+    else:
+        audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    # Pass 1 — ASR (text-final; never touches the speech decoder).
+    transcript = _clean(
+        _text_of(
+            engine.generate(
+                [
+                    {
+                        "prompt": _asr_prompt(),
+                        "multi_modal_data": {"audio": (audio, sr)},
+                        "modalities": ["text"],
+                    }
+                ]
+            )
+        )
+    )
+    print(f"[1/3] transcript : {transcript}")
+    if not transcript:
+        raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
+
+    # Pass 2 — chat answer (text-final). The recipe is size-conditional (see
+    # the TEXT_INSTRUCTION comment above): the 30B requires the official
+    # recipe, the 2B regresses under it.
+    if is_30b(args.model):
+        # Official recipe. The raw completion is
+        # <reasoning...></think><answer>; speak only the final response.
+        # Retry with the next seed only when the budget ran out before the
+        # reasoning block closed (a natural stop without </think> is a
+        # complete answer); if every seed spins, fall back to the official
+        # demo's behavior of keeping the full text.
+        raw = ""
+        for seed in CHAT_SEEDS:
+            raw, truncated = _chat_once(engine, _chat_prompt(transcript), seed)
+            if "</think>" in raw or not truncated:
+                break
+            print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare prompt with a closed think block, engine-default sampling
+        # (yaml near-greedy), no retry; the answer is used as-is.
+        chat_prompt = f"<|im_start|>user\n{transcript}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+        answer = _text_of(engine.generate([{"prompt": chat_prompt, "modalities": ["text"]}]))
+    print(f"[2/3] answer     : {answer}")
+    if not answer:
+        raise SystemExit("Chat pass produced an empty answer; aborting cascade")
+
+    # Pass 3 — TTS (audio-final; streams through code2wav). Both final
+    # stages of the pipeline may emit an output record for the request;
+    # collect the audio-bearing one.
+    # The formatting instruction puts every sentence on its own line;
+    # flatten to plain prose for the TTS thinker.
+    tts_prompt = build_cond_prompt(re.sub(r"\s+", " ", answer).strip())
+    outputs = engine.generate(
+        [{"prompt": tts_prompt, "modalities": ["audio"]}],
+        guided_stage_params(engine, tts_prompt, tokenizer, args.cfg_scale, "s2s-tts-pass"),
+    )
+    chunks: list[torch.Tensor] = []
+    for req_output in outputs:
+        mm = getattr(req_output, "multimodal_output", None) or {}
+        if not mm and req_output.outputs:
+            mm = getattr(req_output.outputs[0], "multimodal_output", None) or {}
+        audio_val = mm.get("model_outputs", mm.get("audio"))
+        if audio_val is None:
+            continue
+        if isinstance(audio_val, list):
+            chunks.extend(torch.as_tensor(a).float().cpu().reshape(-1) for a in audio_val if a is not None)
+        else:
+            chunks.append(torch.as_tensor(audio_val).float().cpu().reshape(-1))
+    pcm = torch.cat(chunks) if chunks else torch.empty(0)
+    if pcm.numel() == 0:
+        raise SystemExit("TTS pass produced empty audio")
+
+    out_path = Path(args.output)
+    write_wav(out_path, pcm)
+    print(f"[3/3] speech     : {pcm.numel() / SAMPLE_RATE:.2f}s -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/text_to_audio.py

@@ -0,0 +1,157 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) text-to-audio inference example.
+
+Caption → general audio through the vLLM-Omni TTA pipeline: the audiogen
+thinker generates interleaved 4-codebook <audiocodec_N> RVQ tokens under an
+RVQ phase mask, and the external XCodec1 checkpoint decodes them to a 16 kHz
+mono WAV.
+
+(For DIFFUSION text-to-audio models — Stable Audio Open, AudioX — use
+``examples/offline_inference/text_to_audio/text_to_audio.py`` instead; Audex
+TTA is a two-stage autoregressive pipeline with a different interface.)
+
+Classifier-free guidance is effectively mandatory for TTA quality (official
+default scale 3.0), so every request submits a cond/uncond pair; requests
+run one at a time.
+
+Pass the HF repo ROOT as --model. XCodec1 resolves from --xcodec1-path /
+XCODEC1_PATH / the default hf-audio repo.
+
+Example:
+
+    python examples/offline_inference/audex/text_to_audio.py \\
+        --captions "Heavy rain falling on a tin roof." \\
+        --output-dir results/audex_tta_wavs
+"""
+
+from __future__ import annotations
+
+import argparse
+import os
+import time
+from pathlib import Path
+
+from common import (
+    SAMPLE_RATE,
+    default_deploy_config,
+    extract_pcm,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_corpus,
+    write_wav,
+)
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt, build_tta_null_prompt
+from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
+
+DEFAULT_CAPTIONS = (
+    "Heavy rain falling on a tin roof.",
+    "A dog barking in the distance while birds chirp.",
+    "Ocean waves crashing on a rocky shore.",
+)
+DEFAULT_CODEC_CAP = 4000
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex TTA inference")
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument("--captions", type=str, nargs="+", default=None, help="Audio captions.")
+    parser.add_argument(
+        "--captions-file",
+        type=str,
+        default=None,
+        help="TSV corpus: one 'utt_id<TAB>caption' per line (overrides --captions).",
+    )
+    parser.add_argument("--output-dir", type=str, default="results/audex_tta_wavs")
+    parser.add_argument(
+        "--deploy-config",
+        type=str,
+        # The model root's default deploy yaml is the TTS pipeline; TTA
+        # prompts and RVQ sampling params require the dedicated TTA pipeline.
+        default=default_deploy_config("audex_tta.yaml"),
+        help="Deploy yaml (defaults to the audex_tta pipeline).",
+    )
+    parser.add_argument(
+        "--xcodec1-path",
+        type=str,
+        default=None,
+        help="Local XCodec1 checkpoint (defaults to $XCODEC1_PATH or the hf-audio repo).",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=3.0,
+        help="Classifier-free guidance strength (official TTA default 3.0; 1.0 disables).",
+    )
+    parser.add_argument(
+        "--codec-cap",
+        type=int,
+        default=DEFAULT_CODEC_CAP,
+        help="Max generated codec tokens before the phase mask forces <audiogen_end>.",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    if args.xcodec1_path:
+        os.environ["XCODEC1_PATH"] = args.xcodec1_path
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    corpus = load_corpus(args.captions_file, args.captions, DEFAULT_CAPTIONS)
+
+    tokenizer = load_audex_tokenizer(args.model, profile="tta")
+    phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
+    tta_rvq = {
+        "phase_token_ids": phase_token_ids,
+        "start_tid": start_tid,
+        "end_tid": end_tid,
+        "codec_cap": args.codec_cap,
+        # The prompt ends with <audiogen_start>, so generation begins at
+        # RVQ phase 0 immediately.
+        "start_in_prompt": True,
+    }
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    print(f"Model       : {args.model}")
+    print(f"Captions    : {len(corpus)}")
+    print(f"CFG scale   : {args.cfg_scale}")
+    print(f"Codec cap   : {args.codec_cap}")
+    print(f"Output dir  : {output_dir}")
+
+    total_elapsed = 0.0
+    total_dur = 0.0
+    for utt, caption in corpus:
+        cond_prompt = build_tta_cond_prompt(caption)
+        sampling = guided_stage_params(
+            engine,
+            cond_prompt,
+            tokenizer,
+            args.cfg_scale,
+            f"tta-{utt}",
+            extra_args={"tta_rvq": tta_rvq},
+            null_prompt_builder=build_tta_null_prompt,
+        )
+        t_start = time.perf_counter()
+        outputs = engine.generate([cond_prompt], sampling)
+        total_elapsed += time.perf_counter() - t_start
+
+        (req_output,) = outputs
+        pcm = extract_pcm(req_output.outputs[0].multimodal_output)
+        dur = pcm.numel() / SAMPLE_RATE
+        if dur <= 0:
+            raise RuntimeError(f"empty audio for {utt}")
+        out_path = output_dir / f"{utt}.wav"
+        write_wav(out_path, pcm)
+        total_dur += dur
+        print(f"  {utt:<48} dur={dur:6.2f}s  -> {out_path}")
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/audex/text_to_speech.py

@@ -0,0 +1,162 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline Audex (Nemotron-Labs-Audex-2B) TTS inference example.
+
+Runs Stage 0 (thinker) + Stage 1 (streaming causal speech decoder) end-to-end
+through the vLLM-Omni engine and saves a 16 kHz mono WAV per prompt.
+
+Pass the HF repo ROOT as --model; per-stage subfolders resolve automatically.
+
+Examples:
+
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts "The weather is so good today." \\
+        --output-dir results/audex_wavs
+
+    # TSV corpus (utt_id<TAB>text per line), all requests in flight at once:
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts-file corpus.tsv --output-dir results/audex_wavs
+
+    # One request at a time:
+    python examples/offline_inference/audex/text_to_speech.py \\
+        --texts-file corpus.tsv --sequential --output-dir results/audex_wavs
+"""
+
+from __future__ import annotations
+
+import argparse
+import time
+from pathlib import Path
+
+from common import (
+    GUIDED_TEMPERATURE,
+    SAMPLE_RATE,
+    extract_pcm,
+    guided_stage_params,
+    load_audex_tokenizer,
+    load_corpus,
+    request_index,
+    write_wav,
+)
+
+from vllm_omni import Omni
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+DEFAULT_TEXTS = (
+    "Hello world.",
+    "The quick brown fox jumps over the lazy dog.",
+    "Today is a beautiful day for a walk in the park.",
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Offline Audex TTS inference")
+    parser.add_argument(
+        "--model",
+        type=str,
+        default="nvidia/Nemotron-Labs-Audex-2B",
+        help="HF repo root (or local snapshot root).",
+    )
+    parser.add_argument("--texts", type=str, nargs="+", default=None, help="Plain-text prompts.")
+    parser.add_argument(
+        "--texts-file",
+        type=str,
+        default=None,
+        help="TSV corpus: one 'utt_id<TAB>text' per line (overrides --texts).",
+    )
+    parser.add_argument("--output-dir", type=str, default="results/audex_wavs")
+    parser.add_argument("--deploy-config", type=str, default=None, help="Override the deploy config path.")
+    parser.add_argument(
+        "--sequential",
+        action="store_true",
+        help="Submit one request at a time instead of all at once.",
+    )
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=1.0,
+        help=(
+            "Classifier-free guidance strength. 1.0 disables guidance (default, "
+            "matches the v1 baseline); 1.5 is the official quality setting. "
+            "Guided runs submit one request at a time (each carries its own "
+            "length-matched null prompt and pair id)."
+        ),
+    )
+    parser.add_argument(
+        "--temperature",
+        type=float,
+        default=None,
+        help="Override the deploy yaml's stage-0 sampling temperature.",
+    )
+    return parser.parse_args()
+
+
+def _write_outputs(outputs, batch: list[tuple[str, str]], output_dir: Path) -> float:
+    ordered = sorted(outputs, key=request_index)
+    if len(ordered) != len(batch):
+        raise RuntimeError(f"expected {len(batch)} outputs, got {len(ordered)}")
+
+    total_dur = 0.0
+    for (utt, _text), req_output in zip(batch, ordered):
+        pcm = extract_pcm(req_output.outputs[0].multimodal_output)
+        dur = pcm.numel() / SAMPLE_RATE
+        if dur <= 0:
+            raise RuntimeError(f"empty audio for {utt}")
+        out_path = output_dir / f"{utt}.wav"
+        write_wav(out_path, pcm)
+        total_dur += dur
+        print(f"  {utt:<40} dur={dur:6.2f}s  -> {out_path}")
+    return total_dur
+
+
+def main():
+    args = parse_args()
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    corpus = load_corpus(args.texts_file, args.texts, DEFAULT_TEXTS)
+
+    engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
+
+    cfg_enabled = args.cfg_scale > 1.0
+    tokenizer = load_audex_tokenizer(args.model) if cfg_enabled else None
+
+    print(f"Model       : {args.model}")
+    print(f"Prompts     : {len(corpus)}")
+    print(f"Mode        : {'sequential' if args.sequential or cfg_enabled else 'batched'}")
+    print(f"CFG scale   : {args.cfg_scale}" + (" (guidance off)" if not cfg_enabled else ""))
+    print(f"Output dir  : {output_dir}")
+
+    total_elapsed = 0.0
+    total_dur = 0.0
+    # CFG runs one request at a time: each request needs its own pair id and
+    # text-specific null prompt in the stage-0 sampling params.
+    batches = [[item] for item in corpus] if (args.sequential or cfg_enabled) else [corpus]
+    for batch in batches:
+        prompts = [build_cond_prompt(text) for _utt, text in batch]
+        sampling_params_list = None
+        if cfg_enabled:
+            sampling_params_list = guided_stage_params(
+                engine,
+                prompts[0],
+                tokenizer,
+                args.cfg_scale,
+                f"cfg-{batch[0][0]}",
+                temperature=GUIDED_TEMPERATURE,
+            )
+        if args.temperature is not None:
+            import copy
+
+            if sampling_params_list is None:
+                sampling_params_list = copy.deepcopy(engine.resolve_sampling_params_list(None))
+            sampling_params_list[0].temperature = args.temperature
+        t_start = time.perf_counter()
+        outputs = engine.generate(prompts, sampling_params_list)
+        total_elapsed += time.perf_counter() - t_start
+        total_dur += _write_outputs(outputs, batch, output_dir)
+
+    rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
+    print(f"Total infer : {total_elapsed:.2f}s  total audio: {total_dur:.2f}s  RTF: {rtf:.3f}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/image_to_video/README.md

@@ -1,7 +1,7 @@
 # Image-To-Video
 
-This shared example generates videos from images with VACE, Wan2.2,
-LTX-2/LTX-2.3, HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
+This shared example generates videos from images with VACE, Wan2.2, LTX-2,
+HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
 
 - `image_to_video.py`: command-line script for single video generation with advanced options.
 
@@ -26,7 +26,7 @@ This folder provides a unified CLI script for image-to-video generation using vL
 | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | 480 x 832 | 81 | 50 | 5.0 | Around 60 GiB BF16 for basic single-card usage |
 | `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | 480 x 832 | 81 | 50 | 4.0 | Around 20–25 GiB BF16, smallest I2V model |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v` | 480 x 832 | 121 | 50 | 6.0 | Around 100 GiB at default settings; the example enables `--enable-cpu-offload` + VAE tiling/slicing to fit an 80 GiB card |
-| LTX2 (local path + `--model-class-name LTX2ImageToVideoPipeline`) | 512 x 768 | 121 | 40 | 4.0 | Memory use depends on frame count and tensor parallelism |
+| `Lightricks/LTX-2` | 512 x 768 | 121 | 40 | video 3.0 / audio 7.0 | Memory use depends on frame count and tensor parallelism |
 
 !!! info
     Peak VRAM: based on basic single-card usage, batch size = 1, without any acceleration/optimization features. Some model weights cannot fit into one card with 80 GiB VRAM, which may need to use CPU offloading.
@@ -105,7 +105,7 @@ python image_to_video.py \
 | Argument | Type | Default | Description |
 | -------- | ---- | ------- | ----------- |
 | `--model` | str | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | Diffusers I2V model ID or local path |
-| `--model-class-name` | str | `None` | Override model class name (e.g., `LTX2ImageToVideoPipeline`) |
+| `--model-class-name` | str | `None` | Optional pipeline override |
 | `--image` | str | (required) | Path to input image |
 | `--prompt` | str | `""` | Text description of desired motion/animation |
 | `--negative-prompt` | str | `""` | Optional list of artifacts to suppress |
@@ -195,24 +195,19 @@ python image_to_video.py \
   --output hunyuan_i2v.mp4
 ```
 
-### LTX2 Image-to-Video
+### LTX-2
 
 ```bash
 python image_to_video.py \
-  --model /path/to/LTX-2 \
-  --model-class-name LTX2ImageToVideoPipeline \
+  --model Lightricks/LTX-2 \
   --image cherry_blossom.jpg \
   --prompt "A cinematic dolly shot of cherry blossoms" \
-  --height 512 \
-  --width 768 \
-  --num-frames 121 \
-  --num-inference-steps 40 \
-  --guidance-scale 4.0 \
-  --frame-rate 24 \
-  --fps 24 \
   --output ltx2_i2v.mp4
 ```
 
+See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
+pipeline selection, T2V, defaults, and advanced options.
+
 ## Advanced Features
 
 ### CFG Parallel
@@ -233,30 +228,6 @@ python image_to_video.py \
   --output i2v_cached.mp4
 ```
 
-### LTX-2.3
-
-```bash
-python image_to_video.py \
-  --model diffusers/LTX-2.3-Diffusers \
-  --model-class-name LTX23ImageToVideoPipeline \
-  --image cherry_blossom.jpg \
-  --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
-  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
-  --height 384 \
-  --width 512 \
-  --num-frames 25 \
-  --guidance-scale 4.0 \
-  --num-inference-steps 20 \
-  --frame-rate 24 \
-  --fps 24 \
-  --output ltx23_i2v_output.mp4
-```
-
-Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
-upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
-this pipeline. Pass `--model-class-name LTX23ImageToVideoPipeline` to select
-the LTX-2.3 image-to-video pipeline.
-
 ### Cosmos3
 
 ```bash
@@ -274,8 +245,8 @@ python image_to_video.py \
 
 Key arguments:
 
-- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V,
-  LTX-2/LTX-2.3, Cosmos3, or VACE).
+- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V, LTX-2,
+  Cosmos3, or VACE).
 - `--image`: Path to the first-frame or source image.
 - `--last-image`: Optional last-frame condition for models such as VACE.
 - `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
@@ -290,7 +261,7 @@ Key arguments:
 - `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
 - `--negative-prompt`: Optional list of artifacts to suppress.
 - `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
-- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
+- `--flow-shift`: Scheduler flow shift (default: Wan 5.0, Cosmos3 10.0). LTX does not consume this option.
 - `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
 - `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).

examples/offline_inference/image_to_video/image_to_video.py

@@ -23,16 +23,14 @@ Usage:
 
     # LTX2 image-to-video
     python image_to_video.py --model /path/to/LTX-2 \
-        --model-class-name LTX2ImageToVideoPipeline \
         --image input.jpg --prompt "A cinematic dolly shot of a boat" \
-        --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
+        --num-frames 121 --num-inference-steps 40 \
         --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
 
     # LTX-2.3 image-to-video
     python image_to_video.py --model diffusers/LTX-2.3-Diffusers \
-        --model-class-name LTX23ImageToVideoPipeline \
         --image input.jpg --prompt "A cinematic dolly shot of a boat" \
-        --height 384 --width 512 --num-frames 25 --num-inference-steps 20 \
+        --height 512 --width 768 --num-frames 121 --num-inference-steps 30 \
         --frame-rate 24 --fps 24 --output ltx23_i2v.mp4
 
     # HunyuanVideo-1.5 I2V (480p)
@@ -100,7 +98,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--model-class-name",
         default=None,
-        help="Override model class name (e.g., LTX2ImageToVideoPipeline or LTX23ImageToVideoPipeline).",
+        help="Override model class name (LTX checkpoints default to LTX2Pipeline).",
     )
     parser.add_argument(
         "--deploy-config",
@@ -117,7 +115,7 @@ def parse_args() -> argparse.Namespace:
         help="Path to a reference image. Repeat to provide multiple references.",
     )
     parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
     parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
     parser.add_argument(
@@ -331,7 +329,10 @@ def main():
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
     model_name = str(args.model).lower() if args.model is not None else ""
     model_class_name = args.model_class_name
-    is_ltx2 = model_class_name in {"LTX2ImageToVideoPipeline", "LTX23ImageToVideoPipeline"}
+    model_class_name_lower = (model_class_name or "").lower()
+    is_ltx2_distilled = "distilled" in model_class_name_lower or "distilled" in model_name
+    is_ltx23 = "ltx23" in model_class_name_lower or "ltx-2.3" in model_name
+    is_ltx2 = is_ltx2_distilled or is_ltx23 or "ltx2" in model_class_name_lower or "ltx-2" in model_name
     is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
 
     image = PIL.Image.open(args.image).convert("RGB") if args.image else None
@@ -356,13 +357,23 @@ def main():
             1280 * 720,
             16,
         )
+    elif is_ltx2_distilled:
+        d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
+            24,
+            None,
+            121,
+            8,
+            None,
+            1024 * 1536,
+            64,
+        )
     elif is_ltx2:
         d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
             24,
-            4.0,
+            None,
             121,
-            40,
-            5.0,
+            30 if is_ltx23 else 40,
+            None,
             512 * 768,
             32,
         )
@@ -445,7 +456,6 @@ def main():
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
-        flow_shift=flow_shift,
         diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
         diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
         diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
@@ -460,6 +470,8 @@ def main():
     )
     if args.deploy_config:
         omni_kwargs["deploy_config"] = args.deploy_config
+    if flow_shift is not None:
+        omni_kwargs["flow_shift"] = flow_shift
     if args.quantization is not None:
         omni_kwargs["quantization"] = args.quantization
     # Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
@@ -492,10 +504,14 @@ def main():
     print(f"  Video size: {width}x{height}")
     print(f"{'=' * 60}\n")
 
+    negative_prompt = args.negative_prompt
+    if negative_prompt is None and not is_ltx2:
+        # Preserve the historical empty-prompt behavior for non-LTX examples.
+        negative_prompt = ""
     prompt_dict = build_image_to_video_prompt(
         model_class_name=model_class_name,
         prompt=args.prompt,
-        negative_prompt=args.negative_prompt,
+        negative_prompt=negative_prompt,
         media_inputs=media_inputs,
         height=height,
         width=width,
@@ -511,11 +527,10 @@ def main():
         num_inference_steps=num_inference_steps,
         num_frames=num_frames,
         frame_rate=frame_rate,
-        extra_args={
-            "sample_solver": args.sample_solver,
-            "flow_shift": flow_shift,
-        },
+        extra_args={"sample_solver": args.sample_solver},
     )
+    if flow_shift is not None:
+        sampling_params.extra_args["flow_shift"] = flow_shift
 
     # Route model-specific knobs through extra_body, filtered against the model's
     # declared extra_body_params. Models without a declaration only forward explicit

examples/offline_inference/text_to_video/text_to_video.md

@@ -8,7 +8,7 @@ A unified script for text-to-video generation. Supports multiple models with mod
 |---|---|---|---|---|---|
 | `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
-| `diffusers/LTX-2.3-Diffusers` | 384x512 | 25 | 20 | 4.0 | 96GB-class GPU |
+| `Lightricks/LTX-2` | 512x768 | 121 | 40 | video 3.0 / audio 7.0 | Model-dependent |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
 | `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
@@ -55,45 +55,17 @@ python text_to_video.py \
   --output vace_t2v_output.mp4
 ```
 
-LTX2 example:
+### LTX-2
 
 ```bash
 python text_to_video.py \
-  --model "Lightricks/LTX-2" \
-  --prompt "A cinematic close-up of ocean waves at golden hour." \
-  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
-  --height 512 \
-  --width 768 \
-  --num-frames 121 \
-  --num-inference-steps 40 \
-  --guidance-scale 4.0 \
-  --frame-rate 24 \
-  --output ltx2_out.mp4
-```
-
-### LTX-2.3
-
-```bash
-python text_to_video.py \
-  --model diffusers/LTX-2.3-Diffusers \
-  --model-class-name LTX23Pipeline \
+  --model Lightricks/LTX-2 \
   --prompt "Cherry blossoms swaying gently in the breeze with synchronized ambient sound" \
-  --negative-prompt "worst quality, inconsistent motion, blurry, jittery, distorted" \
-  --height 384 \
-  --width 512 \
-  --num-frames 25 \
-  --num-inference-steps 20 \
-  --guidance-scale 4.0 \
-  --frame-rate 24 \
-  --fps 24 \
-  --audio-sample-rate 48000 \
-  --output ltx23_t2v_output.mp4
+  --output ltx2_output.mp4
 ```
 
-Use the Diffusers-format checkpoint `diffusers/LTX-2.3-Diffusers`; the
-upstream `Lightricks/LTX-2.3` raw safetensors repo is not directly loadable by
-this pipeline. Pass `--model-class-name LTX23Pipeline` to select the LTX-2.3
-text-to-video pipeline explicitly.
+See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
+pipeline selection, I2V, defaults, and advanced options.
 
 ### HunyuanVideo-1.5 (480p)
 
@@ -212,8 +184,7 @@ python text_to_video.py \
 ### Common
 
 - `--model`: Diffusers model ID or local path.
-- `--model-class-name`: Optional explicit pipeline class. Use `LTX23Pipeline`
-  for LTX-2.3 text-to-video.
+- `--model-class-name`: Optional explicit pipeline override.
 - `--prompt`: text description (string).
 - `--height/--width`: output resolution. Default depends on model.
 - `--num-frames`: number of frames. Default depends on model.
@@ -229,8 +200,7 @@ python text_to_video.py \
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
 - `--enable-layerwise-offload`: enable layerwise offloading on DiT modules.
 - `--frame-rate`: generation FPS for pipelines that require it (e.g., LTX2).
-- `--audio-sample-rate`: audio sample rate for embedded audio (when the
-  pipeline returns audio; LTX-2.3 outputs 48kHz audio).
+- `--audio-sample-rate`: fallback audio sample rate when the pipeline returns audio.
 - `--quantization`: quantization method (such as `fp8` for FP8).
 - `--flow-shift`: scheduler flow_shift parameter.
 - `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).

examples/offline_inference/text_to_video/text_to_video.py

@@ -65,11 +65,42 @@ _MODEL_PRESETS = {
         "fps": 16,
         "output": "helios_output.mp4",
     },
+    "ltx2": {
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "num_inference_steps": 40,
+        "fps": 24,
+        "output": "ltx2_output.mp4",
+    },
+    "ltx2_distilled": {
+        "height": 1024,
+        "width": 1536,
+        "num_frames": 121,
+        "num_inference_steps": 8,
+        "fps": 24,
+        "output": "ltx2_distilled_output.mp4",
+    },
+    "ltx23": {
+        "height": 512,
+        "width": 768,
+        "num_frames": 121,
+        "num_inference_steps": 30,
+        "fps": 24,
+        "output": "ltx23_output.mp4",
+    },
 }
 
 
-def _detect_preset(model: str) -> dict:
+def _detect_preset(model: str, model_class_name: str | None = None) -> dict:
     model_lower = model.lower()
+    class_lower = (model_class_name or "").lower()
+    if "distilled" in class_lower or "distilled" in model_lower:
+        return _MODEL_PRESETS["ltx2_distilled"]
+    if "ltx23" in class_lower or "ltx-2.3" in model_lower or "ltx_2.3" in model_lower:
+        return _MODEL_PRESETS["ltx23"]
+    if "ltx2" in class_lower or "ltx-2" in model_lower or "ltx_2" in model_lower:
+        return _MODEL_PRESETS["ltx2"]
     if "vace" in model_lower:
         return _MODEL_PRESETS["vace"]
     if "cosmos" in model_lower:
@@ -122,7 +153,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--model-class-name",
         default=None,
-        help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
+        help="Override model class name (e.g., LTX2Pipeline).",
     )
     parser.add_argument(
         "--deploy-config",
@@ -130,7 +161,7 @@ def parse_args() -> argparse.Namespace:
         help="Optional deploy config YAML to use for pipeline-backed runs.",
     )
     parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument("--negative-prompt", default=None, help="Negative prompt. Default: model-specific.")
     parser.add_argument(
         "--extra-body",
         type=parse_extra_body,
@@ -331,7 +362,7 @@ def main():
     args = parse_args()
     model_class_name = args.model_class_name
 
-    preset = _detect_preset(args.model)
+    preset = _detect_preset(args.model, model_class_name)
     for key, default_val in preset.items():
         if getattr(args, key.replace("-", "_"), None) is None:
             setattr(args, key.replace("-", "_"), default_val)
@@ -422,8 +453,11 @@ def main():
     print(f"{'=' * 60}\n")
 
     prompt_dict = {"prompt": args.prompt}
-    if args.negative_prompt:
+    if args.negative_prompt is not None:
         prompt_dict["negative_prompt"] = args.negative_prompt
+    elif preset not in (_MODEL_PRESETS["ltx2"], _MODEL_PRESETS["ltx23"]):
+        # Preserve the historical empty-prompt behavior for non-LTX examples.
+        prompt_dict["negative_prompt"] = ""
 
     sampling_kwargs = dict(
         height=args.height,

examples/online_serving/audex/README.md

@@ -0,0 +1,89 @@
+# Audex (Nemotron-Labs-Audex-2B) online serving
+
+One checkpoint, four deployment modes. `run_server.sh` starts the server
+for a given `MODE`; `client.py --mode <mode>` tests it. The capability
+matrix:
+
+| mode (`vllm_omni/deploy/audex_<mode>.yaml`) | audio in | text out | speech out | general audio out | endpoint |
+|---|---|---|---|---|---|
+| `tts` | ❌ | ❌ | ✅ | ❌ | `/v1/audio/speech` |
+| `tta` | ❌ | ❌ | ❌ | ✅ | `/v1/audio/speech` |
+| `thinker_only` | ✅ | ✅ | ❌ | ❌ | `/v1/chat/completions` |
+| `s2s` | ✅ | ✅ | ✅ | ❌ | both |
+
+Audio-input modes of `client.py` fall back to vLLM's public
+`mary_had_lamb` asset when `--audio-file` is omitted; CFG defaults follow
+the official settings per mode (tts/s2s 1.5, tta 3.0; `--cfg-scale 1.0`
+disables).
+
+## tts — text → speech
+
+Text-only thinker (`checkpoint_folder_audiogen`) emits `<speechcodec_N>`
+tokens; the streaming causal decoder produces 16 kHz WAVs. This is also
+the DEFAULT pipeline when the repo root is served without a deploy config.
+
+```bash
+./run_server.sh                       # MODE=tts, port 8097
+
+python client.py --mode tts --text "Hello there." --output hello.wav
+# or raw curl:
+curl -s http://localhost:8097/v1/audio/speech \
+    -H 'Content-Type: application/json' \
+    -d '{"model": "nvidia/Nemotron-Labs-Audex-2B", "input": "Hello there.", "response_format": "wav", "extra_params": {"cfg_scale": 1.5}}' \
+    -o hello.wav
+```
+
+## tta — caption → general audio
+
+Same thinker over the 4-codebook `<audiocodec_N>` RVQ block, decoded by
+the external XCodec1 checkpoint (auto-downloaded; override with
+`XCODEC1_PATH`). Clips are capped at 10 s.
+
+```bash
+MODE=tta ./run_server.sh
+
+python client.py --mode tta --caption "Heavy rain falling on a tin roof." \
+    --output rain.wav
+```
+
+## thinker_only — audio (+ instruction) → text
+
+Single-stage audio understanding on the full checkpoint (NV-Whisper
+encoder + projector + LM). Send the clip as `input_audio` chat content.
+
+```bash
+MODE=thinker_only ./run_server.sh
+
+python client.py --mode thinker_only                      # transcribe the demo asset
+python client.py --mode thinker_only --audio-file a.wav \
+    --question "What language is being spoken?"
+```
+
+## s2s — spoken question → spoken answer
+
+The audio-capable thinker plus the speech decoder in one deployment.
+`client.py --mode s2s` runs the official three passes: ASR (chat with
+`input_audio`) → chat on the transcript → TTS on the answer. Text-final
+passes carry `"modalities": ["text"]`; only the TTS pass streams through
+the decoder.
+
+```bash
+MODE=s2s PORT=8098 ./run_server.sh
+
+python client.py --mode s2s --port 8098 \
+    --audio-file question.wav --output answer.wav
+```
+
+Offline counterparts (no HTTP, one `Omni` engine per script) live in
+`examples/offline_inference/audex/`.
+
+## 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B)
+
+`SIZE=30b` switches run_server.sh to the 30B yamls and model id:
+
+    SIZE=30b ./run_server.sh                 # 30B tts on port 8097
+    SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
+
+Pass `--model nvidia/Nemotron-Labs-Audex-30B-A3B` to the client — the server
+validates the model id and 404s on a mismatch. Single-H100 defaults; see the
+yaml comments for the TP2 fallback. First launch downloads ~60 GB.

examples/online_serving/audex/client.py

@@ -0,0 +1,319 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Audex online client for all four deployment modes.
+
+Pick ``--mode`` to match the server's deploy yaml (see run_server.sh):
+
+  tts           /v1/audio/speech: text -> speech WAV
+  tta           /v1/audio/speech: caption -> general-audio WAV
+  thinker_only  /v1/chat/completions: audio (+ question) -> text
+  s2s           the official three-pass cascade against one audex_s2s
+                server: ASR (chat, input_audio) -> chat -> TTS
+
+Audio-input modes fall back to vLLM's public ``mary_had_lamb`` asset when
+``--audio-file`` is omitted.
+
+Examples:
+
+    python examples/online_serving/audex/client.py --mode tts \\
+        --text "Hello there." --output hello.wav
+    python examples/online_serving/audex/client.py --mode tta \\
+        --caption "Heavy rain falling on a tin roof." --output rain.wav
+    python examples/online_serving/audex/client.py --mode thinker_only
+    python examples/online_serving/audex/client.py --mode s2s \\
+        --audio-file question.wav --output answer.wav
+"""
+
+from __future__ import annotations
+
+import argparse
+import base64
+import re
+from pathlib import Path
+
+import requests
+
+ASR_QUESTION = "Transcribe the input speech."
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+# The official chat-pass recipe below (SYSTEM_PROMPT + TEXT_INSTRUCTION
+# prefix, reasoning enabled, temp 1.0/top_p 0.95, seed retry) applies to the
+# 30B-A3B ONLY (see is_30b() below and
+# tests/e2e/offline_inference/test_audex_s2s.py, which mirrors this same
+# split). It is load-bearing on the 30B: without the formatting-instruction
+# prefix a bare transcript turn is treated as a TTS request and the model
+# emits literal <speechcodec_N> tokens instead of a text answer (the system
+# prompt alone does not prevent this). The 2B must NOT use this recipe: the
+# official temp-1.0 seed-0 draw measurably hallucinates off-language text on
+# the 2B (yet closes </think> cleanly, so the seed retry never fires) and
+# fails the round-trip WER gate — measured, reproduced byte-identically. The
+# 2B instead keeps its original bare-message / enable_thinking=False /
+# temperature=0.0 recipe.
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+DEFAULT_TEXT = "Hello there! This is the Audex text to speech pipeline."
+DEFAULT_CAPTION = "Heavy rain falling on a tin roof."
+# Official quality settings; 1.0 disables guidance.
+DEFAULT_CFG = {"tts": 1.5, "tta": 3.0, "s2s": 1.5}
+
+
+def is_30b(model: str) -> bool:
+    # Users pass arbitrary repo IDs or local snapshot paths, so this is a
+    # heuristic (not an exact match against the official 30B-A3B repo ID)
+    # — good enough for an example script; keep in sync with the same
+    # branch in examples/offline_inference/audex/speech_to_speech.py and
+    # tests/e2e/offline_inference/test_audex_s2s.py.
+    return "30b" in model.lower()
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Audex online client (tts / tta / thinker_only / s2s)")
+    parser.add_argument("--mode", choices=("tts", "tta", "thinker_only", "s2s"), default="s2s")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8097)
+    parser.add_argument("--model", type=str, default="nvidia/Nemotron-Labs-Audex-2B")
+    parser.add_argument("--text", type=str, default=DEFAULT_TEXT, help="TTS input text (mode=tts).")
+    parser.add_argument("--caption", type=str, default=DEFAULT_CAPTION, help="TTA caption (mode=tta).")
+    parser.add_argument(
+        "--audio-file",
+        type=str,
+        default=None,
+        help="Input WAV for thinker_only/s2s. Defaults to vLLM's mary_had_lamb asset.",
+    )
+    parser.add_argument("--question", type=str, default=ASR_QUESTION, help="Instruction (mode=thinker_only).")
+    parser.add_argument("--output", type=str, default="results/audex_online.wav")
+    parser.add_argument(
+        "--cfg-scale",
+        type=float,
+        default=None,
+        help="CFG strength; defaults per mode (tts/s2s 1.5, tta 3.0). 1.0 disables.",
+    )
+    parser.add_argument("--max-tokens", type=int, default=512)
+    return parser.parse_args()
+
+
+def _audio_b64(audio_file: str | None) -> str:
+    if audio_file:
+        return base64.b64encode(Path(audio_file).read_bytes()).decode()
+    # Fall back to vLLM's public asset, re-encoded as WAV for the chat API.
+    import io
+
+    import numpy as np
+    import soundfile as sf
+    from vllm.assets.audio import AudioAsset
+
+    audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    buf = io.BytesIO()
+    sf.write(buf, np.asarray(audio, dtype="float32"), int(sr), format="WAV")
+    return base64.b64encode(buf.getvalue()).decode()
+
+
+def _chat(
+    base_url: str,
+    model: str,
+    messages: list[dict],
+    max_tokens: int,
+    enable_thinking: bool = False,
+    temperature: float = 0.0,
+    top_p: float = 1.0,
+    seed: int = 0,
+) -> tuple[str, str]:
+    resp = requests.post(
+        f"{base_url}/v1/chat/completions",
+        json={
+            "model": model,
+            "messages": messages,
+            "max_tokens": max_tokens,
+            "temperature": temperature,
+            "top_p": top_p,
+            "seed": seed,
+            # Text-final passes must say so explicitly: without "modalities"
+            # the server falls back to the deployment's configured output
+            # modalities, which for the s2s pipeline include the audio final
+            # stage — routing ASR/chat through code2wav instead of stopping
+            # at stage 0.
+            "modalities": ["text"],
+            # The checkpoint's chat template thinking toggle. The official
+            # cascaded demo transcribes with enable_thinking=False but runs
+            # the chat pass with reasoning ENABLED (its web UI ships the
+            # "Enable reasoning" checkbox checked), priming the generation
+            # prompt with an open "<think>\n" so the model's reasoning stays
+            # inside the think block; _extract_final_response then drops it.
+            # With thinking disabled the 30B leaks its reasoning as plain
+            # text in front of the final response.
+            "chat_template_kwargs": {"enable_thinking": enable_thinking},
+        },
+        timeout=300,
+    )
+    resp.raise_for_status()
+    choice = resp.json()["choices"][0]
+    return choice["message"]["content"].strip(), choice.get("finish_reason") or ""
+
+
+def _speech(base_url: str, model: str, text: str, cfg_scale: float, out_path: Path) -> None:
+    payload: dict = {"model": model, "input": text, "response_format": "wav"}
+    if cfg_scale > 1.0:
+        payload["extra_params"] = {"cfg_scale": cfg_scale}
+    resp = requests.post(f"{base_url}/v1/audio/speech", json=payload, timeout=600)
+    resp.raise_for_status()
+    out_path.parent.mkdir(parents=True, exist_ok=True)
+    out_path.write_bytes(resp.content)
+
+
+def _audio_question(audio_b64: str, question: str) -> list[dict]:
+    return [
+        {
+            "role": "user",
+            "content": [
+                {"type": "input_audio", "input_audio": {"data": audio_b64, "format": "wav"}},
+                {"type": "text", "text": question},
+            ],
+        }
+    ]
+
+
+def _tts_text(answer: str) -> str:
+    # The formatting instruction's CRITICAL rule puts each final-response
+    # sentence on its own line, and the official demo TTS-es those sentence
+    # segments (split_sentence_segments). Mirror that per-line segmentation
+    # but keep a single /v1/audio/speech request: speak the leading whole
+    # sentences of the CLEANED final response up to ~200 chars, never
+    # cutting mid-sentence. Markdown markup is stripped as before (it makes
+    # the TTS thinker stop almost immediately).
+    sentences = []
+    for line in answer.splitlines():
+        line = re.sub(r"[*#_`]+", "", line)
+        line = re.sub(r"\s+", " ", line).strip()
+        if line:
+            sentences.append(line)
+    picked: list[str] = []
+    total = 0
+    for sentence in sentences:
+        if picked and total + len(sentence) > 200:
+            break
+        picked.append(sentence)
+        total += len(sentence) + 1
+    return " ".join(picked)
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer
+    # (cascaded_s2s_web_server.py): the chat pass reasons inside the primed
+    # <think> block, so keep only what follows the LAST closing tag; fall
+    # back to the full text when the model skipped reasoning or never
+    # closed the block. Benign no-op for the 2B's clean answers.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _clean_transcript(text: str) -> str:
+    # The full checkpoint prefixes transcripts with a language-tag sentence
+    # and quotes the content; extract the quoted payload when present. Span
+    # first-to-last quote (not split on every apostrophe) so contractions
+    # inside the transcript (e.g. "don't") survive intact.
+    text = text.strip()
+    match = re.search(r"'(.*)'", text, flags=re.S)
+    if match and match.group(1).strip():
+        return match.group(1).strip()
+    return text
+
+
+def main():
+    args = parse_args()
+    base_url = f"http://{args.host}:{args.port}"
+    cfg_scale = args.cfg_scale if args.cfg_scale is not None else DEFAULT_CFG.get(args.mode, 1.0)
+    out_path = Path(args.output)
+
+    if args.mode == "tts":
+        _speech(base_url, args.model, args.text, cfg_scale, out_path)
+        print(f"speech: {out_path.stat().st_size} bytes -> {out_path}")
+        return
+
+    if args.mode == "tta":
+        _speech(base_url, args.model, args.caption, cfg_scale, out_path)
+        print(f"audio : {out_path.stat().st_size} bytes -> {out_path}")
+        return
+
+    audio_b64 = _audio_b64(args.audio_file)
+
+    if args.mode == "thinker_only":
+        content, _ = _chat(base_url, args.model, _audio_question(audio_b64, args.question), args.max_tokens)
+        print(f"answer: {_extract_final_response(content)}")
+        return
+
+    # s2s — the official three-pass cascade.
+    asr_content, _ = _chat(base_url, args.model, _audio_question(audio_b64, ASR_QUESTION), args.max_tokens)
+    transcript = _clean_transcript(asr_content)
+    print(f"[1/3] transcript : {transcript}")
+    if not transcript:
+        raise SystemExit("ASR pass produced an empty transcript; aborting cascade")
+
+    # The chat-pass recipe is size-conditional (see the TEXT_INSTRUCTION
+    # comment above): the 30B requires the official recipe, the 2B keeps its
+    # original bare-message recipe.
+    if is_30b(args.model):
+        # Official recipe: system prompt + formatting instruction prefixed
+        # to the transcript, reasoning enabled, and the official demo's
+        # text-pass sampling (its defaults: --text-temperature 1.0,
+        # --text-top-p 0.95, --text-max-tokens 2048). Near-greedy decoding
+        # makes the 30B loop inside its reasoning block until the token
+        # budget runs out without ever closing </think>. Even at the
+        # official temperature some trajectories spin past the budget, so
+        # retry a few deterministic seeds; a natural stop without </think>
+        # is a complete answer, and if every seed spins we keep the full
+        # text like the official demo does.
+        chat_messages = [
+            {"role": "system", "content": SYSTEM_PROMPT},
+            {"role": "user", "content": f"{TEXT_INSTRUCTION}\n\n{transcript}"},
+        ]
+        raw = ""
+        for seed in (0, 1, 2):
+            raw, finish_reason = _chat(
+                base_url,
+                args.model,
+                chat_messages,
+                max(args.max_tokens, 2048),
+                enable_thinking=True,
+                temperature=1.0,
+                top_p=0.95,
+                seed=seed,
+            )
+            if "</think>" in raw or finish_reason != "length":
+                break
+            print(f"[2/3] chat seed {seed} hit the token budget mid-reasoning; retrying")
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare user message (no system/instruction prefix), reasoning
+        # disabled, temperature 0.0, no retry — its original recipe. The
+        # official temp-1.0 seed-0 recipe measurably makes the 2B
+        # hallucinate off-language answers (see the module-level comment).
+        chat_messages = [{"role": "user", "content": transcript}]
+        raw, _ = _chat(
+            base_url,
+            args.model,
+            chat_messages,
+            args.max_tokens,
+            enable_thinking=False,
+            temperature=0.0,
+        )
+        answer = _extract_final_response(raw)
+    print(f"[2/3] answer     : {answer}")
+    if not answer:
+        raise SystemExit("Chat pass produced an empty answer; aborting cascade")
+
+    _speech(base_url, args.model, _tts_text(answer), cfg_scale, out_path)
+    print(f"[3/3] speech     : {out_path.stat().st_size} bytes -> {out_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/audex/run_server.sh

@@ -0,0 +1,57 @@
+#!/bin/bash
+# Launch the vLLM-Omni server for Audex (Nemotron-Labs-Audex-2B).
+#
+# MODE picks the deployment (see README.md for the capability matrix):
+#   tts           text -> speech            /v1/audio/speech   (default)
+#   tta           caption -> general audio  /v1/audio/speech
+#   thinker_only  audio -> text             /v1/chat/completions
+#   s2s           cascaded speech-to-speech (both endpoints)
+#
+# Pass the HF repo ROOT as MODEL — per-stage subfolders resolve
+# automatically.
+#
+# Usage:
+#   ./run_server.sh                       # 2B tts on port 8097, GPU 0
+#   SIZE=30b ./run_server.sh              # 30B-A3B tts (explicit 30B yaml)
+#   SIZE=30b MODE=s2s PORT=8098 ./run_server.sh
+#   MODE=tta ./run_server.sh              # needs XCodec1 (auto-downloaded)
+#   MODEL=/path/to/local/snapshot ./run_server.sh
+
+set -e
+
+MODE="${MODE:-${1:-tts}}"
+SIZE="${SIZE:-2b}"
+if [ "$SIZE" = "30b" ]; then
+    MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-30B-A3B}"
+    YAML_SUFFIX="_30b"
+elif [ "$SIZE" = "2b" ]; then
+    MODEL="${MODEL:-nvidia/Nemotron-Labs-Audex-2B}"
+    YAML_SUFFIX=""
+else
+    echo "Unknown SIZE '$SIZE'; expected 2b|30b" >&2
+    exit 1
+fi
+PORT="${PORT:-8097}"
+GPUS="${GPUS:-0}"
+
+REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../../.." && pwd)"
+DEPLOY_YAML="$REPO_ROOT/vllm_omni/deploy/audex_${MODE}${YAML_SUFFIX}.yaml"
+if [ ! -f "$DEPLOY_YAML" ]; then
+    echo "Unknown MODE '$MODE' / SIZE '$SIZE' (no $DEPLOY_YAML); expected MODE=tts|tta|thinker_only|s2s, SIZE=2b|30b" >&2
+    exit 1
+fi
+
+echo "Starting Audex server"
+echo "  MODE=$MODE ($DEPLOY_YAML)"
+echo "  SIZE=$SIZE"
+echo "  MODEL=$MODEL"
+echo "  PORT=$PORT"
+echo "  CUDA_VISIBLE_DEVICES=$GPUS"
+
+CUDA_VISIBLE_DEVICES="$GPUS" \
+vllm-omni serve "$MODEL" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --trust-remote-code \
+    --stage-configs-path "$DEPLOY_YAML" \
+    --omni

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
+vllm serve Lightricks/LTX-2 --omni --port 8098
 ```
 
-#### Start with Optimization Presets
-
-Use the LTX-2 startup script with built-in optimization presets:
-
-```bash
-# Baseline (1 GPU, eager)
-bash run_server_ltx2.sh baseline
-
-# 4-GPU Ulysses sequence parallelism (lossless)
-bash run_server_ltx2.sh ulysses4
-
-# Cache-DiT lossy acceleration (1 GPU, ~1.4× speedup)
-bash run_server_ltx2.sh cache-dit
-
-# Best combo: 4-GPU Ulysses SP + Cache-DiT (~2.2× speedup)
-bash run_server_ltx2.sh best-combo
-```
-
-#### Optimization Benchmarks
-
-Benchmarked on H800, online serving (480×768, 41 frames, 20 steps, `seed=42`).
-"Inference" is the server-reported inference time; excludes HTTP/poll overhead.
-
-| Preset | Server Command | Inference (s) | Speedup | Type |
-|--------|---------------|---------------|---------|------|
-| `baseline` | `--enforce-eager` | 10.3 | 1.00× | — |
-| `compile` | *(default, no --enforce-eager)* | ~10.3 (warm) | ~1.00× | Lossless |
-| `ulysses4` | `--enforce-eager --usp 4` | ~10.3 | ~1.00× | Lossless |
-| `cache-dit` | `--enforce-eager --cache-backend cache_dit` | 7.4 avg | ~1.4× | Lossy |
-| `best-combo` | `--enforce-eager --usp 4 --cache-backend cache_dit` | 4.7 avg | **~2.2×** | Lossless + Lossy |
-
-**Observations**:
-- **torch.compile**: On H800, warm-request inference time matches the eager baseline (~10.3s).
-  The first request pays ~6s compilation overhead. Benefit depends on model architecture and GPU.
-- **Ulysses SP (4 GPU)**: No measurable speedup alone for 41-frame generation at this resolution.
-  Communication overhead outweighs gains at this sequence length.
-- **Cache-DiT**: Inference varies per request (6–10s) due to dynamic caching decisions.
-  Average is ~7.4s (~1.4× speedup) with slight quality tradeoff.
-- **Best combo**: 4-GPU Ulysses SP + Cache-DiT synergize well — Cache-DiT reduces per-step
-  computation, making the communication overhead of Ulysses SP worthwhile. Average ~4.7s
-  (~2.2× speedup).
-- **FP8 quantization**: Reduces VRAM but does not speed up LTX-2 on H800 (compute-bound).
-
-**Deployment Recommendations**:
-- For **production with quality priority**: use `baseline` with `--enforce-eager`
-- For **maximum throughput** (4 GPUs, quality tradeoff): use `best-combo` (~2.2× speedup)
-- For **single-GPU throughput**: use `cache-dit` (~1.4× speedup)
-- `--enforce-eager` is recommended to avoid torch.compile warmup latency on first request
-
-### Send Requests (curl)
-
-```bash
-# Using the provided script
-bash run_curl_ltx2.sh
-
-# Or directly
-curl -sS -X POST http://localhost:8098/v1/videos \
-  -H "Accept: application/json" \
-  -F "prompt=A serene lakeside sunrise with mist over the water." \
-  -F "width=768" \
-  -F "height=480" \
-  -F "num_frames=41" \
-  -F "fps=24" \
-  -F "num_inference_steps=20" \
-  -F "guidance_scale=3.0" \
-  -F "seed=42"
-```
+See the [LTX-2 recipe](../../../recipes/LTX/LTX-2.md) for all checkpoints,
+pipeline selection, requests, defaults, and advanced options.
 
 ## Helios
 

examples/online_serving/text_to_video/run_curl_ltx2.sh

@@ -22,7 +22,6 @@ create_response=$(
     -F "num_frames=41" \
     -F "fps=24" \
     -F "num_inference_steps=20" \
-    -F "guidance_scale=3.0" \
     -F "seed=42"
 )
 

examples/online_serving/text_to_video/run_server_ltx2.sh

@@ -21,8 +21,6 @@ set -euo pipefail
 
 MODEL="${MODEL:-Lightricks/LTX-2}"
 PORT="${PORT:-8098}"
-FLOW_SHIFT="${FLOW_SHIFT:-1.0}"
-BOUNDARY_RATIO="${BOUNDARY_RATIO:-1.0}"
 
 PRESET="${1:-baseline}"
 
@@ -66,19 +64,13 @@ case "$PRESET" in
         echo "Environment variables:"
         echo "  MODEL           - Model path (default: Lightricks/LTX-2)"
         echo "  PORT            - Server port (default: 8098)"
-        echo "  FLOW_SHIFT      - Scheduler flow shift (default: 1.0)"
-        echo "  BOUNDARY_RATIO  - Boundary ratio (default: 1.0)"
         exit 1
         ;;
 esac
 
 echo "Model: $MODEL"
 echo "Port: $PORT"
-echo "Flow shift: $FLOW_SHIFT"
-echo "Boundary ratio: $BOUNDARY_RATIO"
 
 vllm serve "$MODEL" --omni \
     --port "$PORT" \
-    --flow-shift "$FLOW_SHIFT" \
-    --boundary-ratio "$BOUNDARY_RATIO" \
     "${EXTRA_ARGS[@]}"

recipes/LTX/LTX-2.3.md

@@ -1,154 +0,0 @@
-# LTX-2.3
-
-> 22B parameter text-to-video + audio generation model served via vLLM-Omni
-
-## Summary
-
-- Vendor: Lightricks
-- Model: `diffusers/LTX-2.3-Diffusers`
-- Task: Text-to-video and image-to-video with synchronized audio generation
-- Mode: Online serving (pure diffusion)
-- Maintainer: @oglok
-
-## When to use this recipe
-
-Use this recipe when you want to serve LTX-2.3 for text-to-video or
-image-to-video generation with audio. The model generates videos up to 20+
-seconds at 768x512 resolution with 48kHz audio from text prompts, and can
-condition video generation on an initial image. Start with a 96GB-class GPU for
-validation because the model combines a 22B parameter transformer with text
-encoder, VAE, and vocoder components.
-
-## References
-
-- Upstream raw checkpoints: <https://huggingface.co/Lightricks/LTX-2.3>
-- Diffusers-format checkpoint: <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
-- Requires `diffusers >= 0.38.0` (install from git: `pip install git+https://github.com/huggingface/diffusers.git`)
-
-## Serving
-
-### Command
-
-```bash
-vllm serve diffusers/LTX-2.3-Diffusers \
-  --omni \
-  --model-class-name LTX23Pipeline \
-  --stage-init-timeout 600
-```
-
-For image-to-video:
-
-```bash
-vllm serve diffusers/LTX-2.3-Diffusers \
-  --omni \
-  --model-class-name LTX23ImageToVideoPipeline \
-  --stage-init-timeout 600
-```
-
-### T2V Verification
-
-```bash
-# Health check
-curl http://localhost:8000/health
-
-# Generate a 3-second video (81 frames at 24fps)
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=A majestic bald eagle soaring over a misty mountain valley at dawn, golden sunlight breaking through clouds" \
-  -F "negative_prompt=blurry, low quality, distorted, watermark" \
-  -F "model=diffusers/LTX-2.3-Diffusers" \
-  -F "num_frames=81" \
-  -F "fps=24" \
-  -F "size=768x512" \
-  -F "num_inference_steps=30" \
-  -F "guidance_scale=4.0" \
-  -F "seed=42"
-
-# Generate a 10-second video (241 frames)
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=A cozy Japanese ramen shop at night in the rain, steam rising from bowls, neon signs reflecting on wet cobblestone streets" \
-  -F "model=diffusers/LTX-2.3-Diffusers" \
-  -F "num_frames=241" \
-  -F "fps=24" \
-  -F "size=768x512" \
-  -F "num_inference_steps=30" \
-  -F "guidance_scale=4.0"
-
-# Generate a 20-second video (481 frames)
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=An underwater coral reef teeming with tropical fish, sea turtles gliding gracefully, National Geographic documentary style" \
-  -F "model=diffusers/LTX-2.3-Diffusers" \
-  -F "num_frames=481" \
-  -F "fps=24" \
-  -F "size=768x512" \
-  -F "num_inference_steps=30" \
-  -F "guidance_scale=4.0"
-```
-
-### I2V Verification
-
-Start the server with `--model-class-name LTX23ImageToVideoPipeline`, then pass
-an image reference. The I2V pipeline requires `input_reference` or
-`image_reference` because the first frame is the conditioning input.
-
-```bash
-curl http://localhost:8000/health
-
-curl -X POST http://localhost:8000/v1/videos/sync \
-  -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
-  -F "input_reference=@/absolute/path/to/reference.png" \
-  -F "model=diffusers/LTX-2.3-Diffusers" \
-  -F "num_frames=81" \
-  -F "fps=24" \
-  -F "size=768x512" \
-  -F "num_inference_steps=30" \
-  -F "guidance_scale=4.0" \
-  -F "seed=42" \
-  -o ltx23_i2v.mp4
-```
-
-### Notes
-
-- Memory usage: run a latest-head serving smoke or DFX benchmark on the target
-  GPU before publishing load/peak VRAM numbers.
-- Checkpoint format: use a Diffusers-format checkpoint such as
-  `diffusers/LTX-2.3-Diffusers`; the upstream `Lightricks/LTX-2.3` repository ships
-  raw safetensors and does not contain the subfolder configs required by this
-  pipeline loader.
-- Key flags:
-  - `--stage-init-timeout 600`: Gives the initial `torch.compile` warmup enough time on large checkpoints
-  - `--model-class-name LTX23Pipeline`: Selects the LTX-2.3 text-to-video pipeline (not LTX-2)
-  - `--model-class-name LTX23ImageToVideoPipeline`: Selects the LTX-2.3 image-to-video pipeline
-  - `input_reference=@...`: Uploads the I2V reference image directly
-- Audio: 48kHz AAC via BWE vocoder, automatically synced with video
-- CPU offloading: Text encoder (Gemma-3-12B), connectors, VAE, audio VAE, and vocoder stay on CPU and are moved to GPU only when needed
-- Supported resolutions: 768x512, 512x384 (must be divisible by 32)
-- Frame rate: 24 fps
-- Duration: Controlled by `num_frames` (frames = duration_seconds * 24 + 1)
-- Known limitations:
-  - LTX-2.3 I2V currently supports first-frame image conditioning, matching the LTX image-conditioning contract.
-  - Requires `diffusers >= 0.38.0` (not yet on PyPI, install from git)
-
-## Hardware Support
-
-## GPU
-
-### 1x NVIDIA RTX PRO 6000 Blackwell (96GB)
-
-#### Environment
-
-- OS: Ubuntu 22.04
-- Python: 3.10+
-- Driver / runtime: CUDA 13.0, Driver 580.126.09
-- vLLM version: Match the PR or release checkout being validated
-- vLLM-Omni version or commit: Use the commit you are deploying from
-
-### Benchmarking
-
-Do not copy latency or VRAM numbers between commits or machines. Record them
-only after running a latest-head sweep on the target hardware.
-
-For formal PR benchmarking, reuse
-`tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
-`tests/dfx/perf/scripts/run_diffusion_benchmark.py`. Each case in that file includes a JSON **`mark`** (`hardware_marks` + `full_model` + `diffusion`) for local `-m` filtering. That config captures the
-single-device eager baseline and CFG-parallel=2 case for a small 384x512,
-25-frame, 20-step workload.

recipes/LTX/LTX-2.md

@@ -1,121 +1,303 @@
-# LTX-2
+# LTX-2 Family
 
-> Text-to-video and image-to-video serving
+> LTX-2 and LTX-2.3 text-to-video and image-to-video generation with synchronized audio
 
-## Summary
+## Pipelines
 
-- Vendor: Lightricks
-- Model: `Lightricks/LTX-2`
-- Task: Text-to-video generation and image-to-video generation
-- Mode: Online serving with the OpenAI-compatible video API
-- Maintainer: Community
+| Model | Checkpoint | Task | `--model-class-name` | Batching |
+|---|---|---|---|---|
+| LTX-2 | `Lightricks/LTX-2` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
+| LTX-2 distilled | `rootonchair/LTX-2-19b-distilled` | Two-stage T2V/I2V | `LTX2DistilledPipeline` | No |
+| LTX-2.3 | `diffusers/LTX-2.3-Diffusers` | One-stage T2V/I2V | `LTX2Pipeline` | Yes |
 
-## When to use this recipe
+`LTX2Pipeline` is the unified one-stage entry. Checkpoint metadata selects the
+LTX-2 or LTX-2.3 profile; omitting an image selects T2V, while one initial
+image selects I2V. Both one-stage repositories declare this class, so
+`--model-class-name` is optional. LTX-2.3 requires the Diffusers checkpoint;
+the raw `Lightricks/LTX-2.3` safetensors repository is not directly loadable.
 
-Use this recipe when you want to deploy the base `Lightricks/LTX-2` model
-for both text-to-video and image-to-video generation. The same model ID is used for
-both tasks; switch the pipeline class with `--model-class-name` depending on
-whether you want T2V or I2V.
+`LTX2DistilledPipeline` is the unified distilled entry: omit `image` for T2V
+or provide one initial image for I2V. As in the official LTX pipelines,
+distilled mode is selected explicitly rather than inferred from the checkpoint
+name or metadata; always pass `--model-class-name LTX2DistilledPipeline`.
+
+## API Migration
+
+Only `req` may be passed positionally to `LTX2Pipeline`; every optional
+`forward` argument is keyword-only:
+
+```python
+# No longer supported
+pipe(req, prompt)
+pipe(req, image, prompt)
+
+# Supported
+pipe(req, prompt=prompt)
+pipe(req, image=image, prompt=prompt)
+```
+
+The consolidation also removes these registry names without aliases:
+
+| Removed name | Replacement |
+|---|---|
+| `LTX23Pipeline` | `LTX2Pipeline`; checkpoint metadata selects LTX-2.3 |
+| `LTX2ImageToVideoPipeline` | `LTX2Pipeline` with `image=` |
+| `LTX23ImageToVideoPipeline` | `LTX2Pipeline` with `image=`; checkpoint metadata selects LTX-2.3 |
+| `LTX2TwoStagesPipeline` | `LTX2DistilledPipeline` |
+| `LTX2ImageToVideoTwoStagesPipeline` | `LTX2DistilledPipeline` with `image=` |
+
+Passing any second positional argument now raises `TypeError`. These changes
+affect direct Python callers and explicit `--model-class-name` overrides;
+offline and serving entrypoints already use named fields and are unaffected.
+
+## One-Stage Defaults
+
+| Parameter | LTX-2 | LTX-2.3 |
+|---|---:|---:|
+| Width × height | 768 × 512 | 768 × 512 |
+| Frames / frame rate | 121 / 24 | 121 / 24 |
+| Denoise steps | 40 | 30 |
+| Video/audio CFG | 3.0 / 7.0 | 3.0 / 7.0 |
+| Video/audio STG | 1.0 / 1.0 | 1.0 / 1.0 |
+| Video/audio modality guidance | 3.0 / 3.0 | 3.0 / 3.0 |
+| Video/audio rescale | 0.7 / 0.7 | 0.7 / 0.7 |
+| Video/audio STG blocks | `[29]` / `[29]` | `[28]` / `[28]` |
+
+The model recipe also supplies the default negative prompt. Top-level
+`guidance_scale` overrides both video and audio CFG values. The online video
+API defaults `num_frames` to `1`, so set it explicitly; the offline LTX scripts
+default to `121`.
+
+## Distilled Two-Stage Defaults
+
+| Parameter | Value |
+|---|---:|
+| Final width × height | 1536 × 1024 |
+| Stage 1 width × height | 768 × 512 |
+| Frames / frame rate | 121 / 24 |
+| Stage 1 / Stage 2 denoise steps | 8 / 3 |
+| Guidance | Fixed positive-only |
+
+The dimensions passed to the API are the final output dimensions and must be
+divisible by 64. Both one- and two-stage requests require `num_frames = 8k+1`.
+The distilled recipe fixes both sigma schedules; `num_inference_steps` must be
+`8` when supplied.
+
+## Serving
+
+Start either checkpoint:
+
+```bash
+vllm serve Lightricks/LTX-2 --omni --stage-init-timeout 600
+```
+
+```bash
+vllm serve diffusers/LTX-2.3-Diffusers --omni --stage-init-timeout 600
+```
+
+Start the distilled checkpoint with its two-stage entry:
+
+```bash
+vllm serve rootonchair/LTX-2-19b-distilled --omni \
+  --model-class-name LTX2DistilledPipeline --stage-init-timeout 600
+```
+
+The same server handles T2V and I2V. A T2V request using the selected recipe's
+default guidance is:
+
+```bash
+curl -X POST http://localhost:8000/v1/videos/sync \
+  -F "prompt=A cinematic close-up of ocean waves at golden hour." \
+  -F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
+  -F "size=768x512" \
+  -F "num_frames=121" \
+  -F "fps=24" \
+  -F "seed=42" \
+  -o ltx_t2v.mp4
+```
+
+For I2V, provide exactly one initial image:
+
+```bash
+curl -X POST http://localhost:8000/v1/videos/sync \
+  -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
+  -F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
+  -F "input_reference=@/absolute/path/to/reference.png" \
+  -F "size=768x512" \
+  -F "num_frames=121" \
+  -F "fps=24" \
+  -F "seed=42" \
+  -o ltx_i2v.mp4
+```
+
+Use `image_reference` for a URL or JSON-safe image reference. Do not provide it
+together with `input_reference`.
+
+## Guidance
+
+LTX supports independent video/audio CFG, spatio-temporal guidance (STG),
+cross-modality guidance, and rescaling:
+
+| Parameter | Default | Effect | Alias |
+|---|---:|---|---|
+| `video_cfg_scale` | 3.0 | Video text CFG; `1.0` disables it | `video_cfg_guidance_scale` |
+| `audio_cfg_scale` | 7.0 | Audio text CFG; `1.0` disables it | `audio_cfg_guidance_scale` |
+| `video_stg_scale` | 1.0 | Video STG; `0.0` disables it | `video_stg_guidance_scale` |
+| `audio_stg_scale` | 1.0 | Audio STG; `0.0` disables it | `audio_stg_guidance_scale` |
+| `video_modality_scale` | 3.0 | Audio-to-video guidance; `1.0` disables it | `a2v_guidance_scale` |
+| `audio_modality_scale` | 3.0 | Video-to-audio guidance; `1.0` disables it | `v2a_guidance_scale` |
+| `video_rescale_scale` | 0.7 | Video guidance rescale; `0.0` disables it | — |
+| `audio_rescale_scale` | 0.7 | Audio guidance rescale; `0.0` disables it | — |
+| `video_stg_blocks` | `[29]` / `[28]` | Perturbed video transformer blocks | — |
+| `audio_stg_blocks` | `[29]` / `[28]` | Perturbed audio transformer blocks | — |
+
+The STG block defaults are shown as LTX-2 / LTX-2.3. A canonical name wins
+over its alias. Explicit top-level `guidance_scale` wins over both CFG fields;
+omit it for independent video/audio CFG.
+
+Pass these fields in online `extra_params` or offline `--extra-body`:
+
+```bash
+# Add to the curl request above
+-F 'extra_params={"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
+
+# Add to text_to_video.py or image_to_video.py
+--extra-body '{"video_cfg_scale":3.0,"audio_cfg_scale":7.0}'
+```
+
+### Python API
+
+Normal per-request values belong in `OmniDiffusionSamplingParams`. Pass the
+resulting `DiffusionRequestBatch` as the only positional argument:
+
+```python
+from vllm_omni.diffusion.request import OmniDiffusionRequest
+from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+params = OmniDiffusionSamplingParams(
+    width=768,
+    height=512,
+    num_frames=121,
+    frame_rate=24.0,
+    num_inference_steps=30,
+    seed=42,
+    extra_args={"video_cfg_scale": 3.0, "audio_cfg_scale": 7.0},
+)
+req = DiffusionRequestBatch([
+    OmniDiffusionRequest(
+        prompt={"prompt": "Cherry blossoms moving in a light breeze"},
+        sampling_params=params,
+        request_id="ltx-example",
+    )
+])
+
+t2v_output = pipe(req)
+i2v_output = pipe(req, image=image)
+```
+
+Direct keyword arguments are low-level fallbacks, for example
+`pipe(req, prompt=prompt, width=768)`. Request fields take precedence unless
+noted below.
+
+### Complete `forward` Surface
+
+| Argument | Type/default | Meaning and constraints |
+|---|---|---|
+| `req` | `DiffusionRequestBatch`, required | Only positional argument; contains prompts and per-request sampling parameters. |
+| `image` | image or batch, `None` | Direct value wins over request images; no image selects T2V. I2V accepts one image per prompt, and a batch cannot mix T2V/I2V. |
+| `prompt` | string or list, `None` | Positive-text fallback; request prompts win. Mutually exclusive with `prompt_embeds`. |
+| `negative_prompt` | string or list, `None` | Fallback after request values and before the recipe default; mutually exclusive with negative embeddings. |
+| `height` | `int`, `None` | Request → direct value → recipe default; divisible by 32 one-stage or 64 distilled. |
+| `width` | `int`, `None` | Same precedence and alignment as `height`. |
+| `num_frames` | `int`, `None` | Request → direct value → recipe default; must be `8k+1` and also determines audio duration with `frame_rate`. |
+| `frame_rate` | `float`, `None` | Request `frame_rate` → request `fps` → direct value → recipe default. |
+| `num_inference_steps` | `int`, `None` | Request → direct value → recipe default; minimum 2 one-stage, fixed at 8 for distilled Stage 1. Custom one-stage `sigmas` determine actual steps. |
+| `sigmas` | list of float, `None` | One-stage only. Request values win; every request in a fused batch must use the same schedule. |
+| `timesteps` | list of int, `None` | Compatibility slot; LTX accepts only `None`. Use `sigmas`. |
+| `guidance_scale` | `float`, `None` | One-stage common video/audio CFG fallback; an explicit request value wins. Distilled two-stage uses fixed positive-only guidance. |
+| `guidance_rescale` | `float`, `None` | Accepts only `None` or `0.0`; use the modality rescale fields. |
+| `noise_scale` | `float`, `0.0` | Compatibility slot; LTX accepts only `0.0`. |
+| `num_videos_per_prompt` | `int`, `1` | Output-count fallback; positive request `num_outputs_per_prompt` wins. |
+| `generator` | generator or list, `None` | Explicit RNG; otherwise request generators/seeds are collated. Lists must match the effective output batch. |
+| `latents` | tensor, `None` | Request tensors win. One-stage accepts packed `[B, S, C]` and validated unpacked video layouts. |
+| `audio_latents` | tensor, `None` | One-stage initial audio latents; request tensors win and are collated. |
+| `prompt_embeds` | tensor, `None` | Precomputed positive conditioning; requires `prompt_attention_mask` and cannot accompany `prompt`. |
+| `negative_prompt_embeds` | tensor, `None` | One-stage precomputed negative conditioning; requires its mask and cannot accompany `negative_prompt`. |
+| `prompt_attention_mask` | tensor, `None` | Mask for positive embeddings; request `prompt_attention_mask` or `attention_mask` wins. |
+| `negative_prompt_attention_mask` | tensor, `None` | Mask for negative embeddings; request negative mask fields win. |
+| `decode_timestep` | float or list, `0.0` | Video-VAE decode timestep; request value wins. Lists may match 1, prompt batch, or output batch. |
+| `decode_noise_scale` | float or list, `None` | Same list rules; request value wins. Defaults to `decode_timestep`. |
+| `output_type` | string, `"np"` | Request value wins. `"np"` decodes output; `"latent"` skips VAE/vocoder decode. |
+| `return_dict` | `bool`, `True` | Compatibility slot; only `True` is accepted. |
+| `attention_kwargs` | dict, `None` | Per-call values are unsupported; configure attention at engine startup. |
+| `max_sequence_length` | `int`, `None` | Request → direct value → tokenizer limit. |
+
+Request-object naming differs only in a few places: `num_videos_per_prompt`
+maps to `num_outputs_per_prompt`; images and prompt text/embeddings live in the
+request prompt payload; LTX guidance fields live in sampling `extra_args`.
+
+### Recipe-Specific Request Capabilities
+
+| Override | One-stage | Distilled two-stage |
+|---|---|---|
+| Guidance | Supported | Fixed positive-only |
+| Negative prompt/embeddings | Supported | Rejected |
+| `num_inference_steps` | Supported | Fixed at 8 for Stage 1; Stage 2 uses 3 |
+| Custom `sigmas` | Supported | Rejected; both phases use fixed schedules |
+| Video/audio latents | Supported | Rejected |
+
+These capability checks apply equally to direct `forward` keywords and
+values in `OmniDiffusionSamplingParams`; unsupported values fail instead of
+being silently ignored.
+
+### Custom Sigma Schedules
+
+One-stage Python requests may set final scheduler boundaries directly:
+
+```python
+params = OmniDiffusionSamplingParams(sigmas=[1.0, 0.75, 0.5, 0.25])
+```
+
+Each nonterminal value produces one denoise step, and a terminal `0.0` is
+appended when omitted. This schedule overrides `num_inference_steps`. All
+requests in a fused batch must use the same list. The video form API and
+bundled offline CLI do not currently expose `sigmas`.
+
+### Constraints
+
+- LTX rejects non-default `timesteps`, `flow_shift`, `guidance_rescale`,
+  `noise_scale`, `attention_kwargs`, and `return_dict=False`. Use final
+  `sigmas`, modality rescale fields, startup attention configuration, and the
+  standard `DiffusionOutput` respectively.
+- Fused one-stage requests must resolve to identical LTX guidance and sigma
+  schedules. Keep concurrent requests guidance-homogeneous or use
+  `--max-num-seqs 1`.
+- `--cfg-parallel-size 2` supports only the CFG-only plan: disable STG and
+  modality guidance and set both rescale fields to `0.0`.
+
+## Operational Notes
+
+- LTX-2 one-stage previously loaded and peaked at about 73.5 GiB on one H200
+  141GB; remeasure on the deployed commit and hardware.
+- LTX-2.3 includes a 22B transformer, Gemma encoder, two VAEs, and a vocoder.
+  Start on a 96GB-class GPU or use CPU/layerwise offload on smaller devices.
+- The distilled upsampler uses the configured pipeline dtype and participates
+  in component discovery/device placement; it remains resident during denoise
+  offload because it is used only at the phase boundary.
+- The output audio sample rate comes from the loaded components and is not a
+  request parameter.
+- For benchmarks, use `tests/dfx/perf/tests/test_ltx2_vllm_omni.json` with
+  `tests/dfx/perf/scripts/run_diffusion_benchmark.py`.
+- Current two-stage support is limited to the distilled checkpoint; official
+  LTX two-stage/HQ execution remains out of scope.
 
 ## References
 
-- Model: <https://huggingface.co/Lightricks/LTX-2>
-- User guide:
-  [`docs/user_guide/examples/online_serving/text_to_video.md`](../../docs/user_guide/examples/online_serving/text_to_video.md)
-- Example guide:
-  [`examples/online_serving/text_to_video/README.md`](../../examples/online_serving/text_to_video/README.md)
-  [`examples/online_serving/image_to_video/README.md`](../../examples/online_serving/image_to_video/README.md)
-
-## Hardware Support
-
-## GPU
-
-### 1x H200 141GB (Text-to-Video)
-
-#### Environment
-
-- OS: Ubuntu 22.04.5 LTS
-- Python: 3.12+
-- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU
-- vLLM version: Match the repository requirements from your current checkout
-- vLLM-Omni version or commit: Use the commit you are deploying from
-
-#### Command
-
-```bash
-vllm serve Lightricks/LTX-2 \
-  --omni \
-  --model-class-name LTX2Pipeline
-```
-
-#### Verification
-
-```bash
-curl http://localhost:8000/health
-
-curl -X POST http://localhost:8000/v1/videos \
-  -F "prompt=A cinematic close-up of ocean waves at golden hour." \
-  -F "negative_prompt=worst quality, inconsistent motion, blurry, jittery, distorted" \
-  -F "width=512" \
-  -F "height=768" \
-  -F "num_frames=121" \
-  -F "fps=24" \
-  -F "num_inference_steps=40" \
-  -F "guidance_scale=4.0" \
-  -o ltx2_t2v_h200.mp4
-```
-
-#### Notes
-
-- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference.
-- Key flags:
-  - `--model-class-name LTX2Pipeline` selects the text-to-video pipeline explicitly.
-- Known limitations:
-  - This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine.
-
-### 1x H200 141GB (Image-to-Video)
-
-#### Environment
-
-- OS: Ubuntu 22.04.5 LTS
-- Python: 3.12+
-- Driver / runtime: NVIDIA CUDA environment with H200 141GB GPU
-- vLLM version: Match the repository requirements from your current checkout
-- vLLM-Omni version or commit: Use the commit you are deploying from
-
-#### Command
-
-```bash
-vllm serve Lightricks/LTX-2 \
-  --omni \
-  --model-class-name LTX2ImageToVideoPipeline
-```
-
-#### Verification
-
-```bash
-curl http://localhost:8000/health
-
-curl -X POST http://localhost:8000/v1/videos/sync \
-  -F "prompt=A plush toy astronaut gently waving while the camera slowly pushes in." \
-  -F "input_reference=@/absolute/path/to/reference.png" \
-  -F "width=512" \
-  -F "height=768" \
-  -F "num_frames=121" \
-  -F "fps=24" \
-  -F "num_inference_steps=40" \
-  -F "guidance_scale=4.0" \
-  -o ltx2_i2v_h200.mp4
-```
-
-#### Notes
-
-- Memory usage: Model loads at ~73.5GB, peaks at ~73.5 GiB during inference.
-- Key flags:
-  - `--model-class-name LTX2ImageToVideoPipeline` switches the same `Lightricks/LTX-2` weights into image-to-video mode.
-  - `input_reference=@...` uploads the reference image directly
-- Known limitations:
-  - Use `image_reference` when you want to pass a URL or JSON-safe image reference
-instead of uploading a file.
-  - Do not send `input_reference` and `image_reference` in the same request.
-  - This recipe is tested on 1x H200 validation; add measured latency and VRAM after running it on your target machine.
+- <https://huggingface.co/Lightricks/LTX-2>
+- <https://huggingface.co/Lightricks/LTX-2.3>
+- <https://huggingface.co/diffusers/LTX-2.3-Diffusers>
+- [Online video generation](../../docs/user_guide/examples/online_serving/text_to_video.md)
+- [Request batching](../../docs/user_guide/diffusion/request_batching.md)
+- [T2V offline example](../../examples/offline_inference/text_to_video/text_to_video.md)
+- [I2V offline example](../../examples/offline_inference/image_to_video/README.md)

recipes/NVIDIA/Nemotron-Labs-Audex.md

@@ -0,0 +1,142 @@
+# Nemotron-Labs Audex (2B / 30B-A3B)
+
+> Speech suite on 1×H100: TTS, text-to-audio, audio understanding, and cascaded speech-to-speech
+
+## Summary
+
+- Vendor: NVIDIA
+- Model: `nvidia/Nemotron-Labs-Audex-2B`, `nvidia/Nemotron-Labs-Audex-30B-A3B`
+- Task: Text-to-speech (English), text-to-audio (sound effects), audio understanding (ASR / audio QA), cascaded speech-to-speech
+- Mode: Online serving (`/v1/audio/speech`, `/v1/chat/completions`) and offline batch inference
+- Maintainer: Community
+
+## When to use this recipe
+
+Use this recipe to serve either Audex checkpoint through one of its four
+deployment pipelines. Both sizes share the same repo layout, codec token
+space, and streaming causal speech decoder (16 kHz output); they differ in
+the thinker: the 2B is a dense LM, the 30B-A3B is a hybrid Mamba + MoE
+NemotronH (~3B active parameters). Pick the pipeline by task:
+
+| pipeline (`vllm_omni/deploy/<name>[_30b].yaml`) | audio in | text out | speech out | general audio out | endpoint |
+|---|---|---|---|---|---|
+| `audex_tts` | ❌ | ❌ | ✅ | ❌ | `/v1/audio/speech` |
+| `audex_tta` | ❌ | ❌ | ❌ | ✅ | `/v1/audio/speech` |
+| `audex_thinker_only` | ✅ | ✅ | ❌ | ❌ | `/v1/chat/completions` |
+| `audex_s2s` | ✅ | ✅ | ✅ | ❌ | both |
+
+## References
+
+- Model cards: [Nemotron-Labs-Audex-2B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-2B), [Nemotron-Labs-Audex-30B-A3B](https://huggingface.co/nvidia/Nemotron-Labs-Audex-30B-A3B)
+- Offline examples: [`examples/offline_inference/audex/`](../../examples/offline_inference/audex/)
+- Online examples: [`examples/online_serving/audex/`](../../examples/online_serving/audex/)
+
+## Hardware Support
+
+## GPU
+
+### 1×H100 80GB — Audex-2B
+
+#### Environment
+
+- OS: Linux
+- Python: 3.12+
+- CUDA: 12.x
+- vLLM version: 0.24.0
+- vLLM-Omni version or commit: PR #4976 (branch `audex`)
+
+#### Command
+
+**Online serving** (TTS is the default pipeline; the repo root's
+`model_type` auto-resolves `audex_tts.yaml`):
+
+```bash
+vllm-omni serve nvidia/Nemotron-Labs-Audex-2B \
+    --host 0.0.0.0 --port 8097 \
+    --trust-remote-code --omni
+# other modes: add --stage-configs-path vllm_omni/deploy/audex_{tta,thinker_only,s2s}.yaml
+# or use the launcher: MODE=s2s PORT=8098 examples/online_serving/audex/run_server.sh
+```
+
+**Offline batch inference** (one script per task, each defaults to its
+correct deploy yaml):
+
+```bash
+python examples/offline_inference/audex/text_to_speech.py \
+    --texts "Hello world." --cfg-scale 1.5 --output-dir results/audex_wavs
+python examples/offline_inference/audex/text_to_audio.py     # sound effects
+python examples/offline_inference/audex/audio_qa.py          # ASR / audio QA
+python examples/offline_inference/audex/speech_to_speech.py  # 3-pass cascade
+```
+
+#### Verification
+
+```bash
+curl -X POST http://localhost:8097/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "model": "nvidia/Nemotron-Labs-Audex-2B",
+        "input": "Hello, how are you?",
+        "response_format": "wav",
+        "extra_params": {"cfg_scale": 1.5}
+    }' --output hello.wav
+# or the universal client (all four modes):
+python examples/online_serving/audex/client.py --mode tts --output hello.wav
+```
+
+#### Notes
+
+- Memory usage: TTS deploy — stage 0 (thinker) 0.4, stage 1 (decoder) 0.25
+  GPU-memory fraction; both stages share GPU 0.
+- Key flags: `--trust-remote-code` and `--omni` are required.
+- Classifier-free guidance: `extra_params.cfg_scale` (1.0 = off; 1.5 is the
+  official quality setting; TTA effectively requires its default 3.0).
+  Measured en-24 self-transcribed CER: 6.87% guided vs 7.24% unguided.
+- Output: 16 kHz mono WAV. English, single built-in voice (no cloning).
+- TTA decodes through the external XCodec1 checkpoint
+  (`hf-audio/xcodec-hubert-general-balanced`, auto-downloaded; override
+  with `XCODEC1_PATH`).
+
+### 1×H100 80GB — Audex-30B-A3B
+
+#### Environment
+
+Same as the 2B row above.
+
+#### Command
+
+The 30B REQUIRES its explicit deploy yaml (the shared HF `model_type`
+auto-resolves to the 2B-tuned config otherwise):
+
+```bash
+vllm-omni serve nvidia/Nemotron-Labs-Audex-30B-A3B \
+    --host 0.0.0.0 --port 8097 \
+    --trust-remote-code --omni \
+    --stage-configs-path vllm_omni/deploy/audex_tts_30b.yaml
+# or: SIZE=30b MODE=tts examples/online_serving/audex/run_server.sh
+```
+
+Offline: same four scripts with
+`--model nvidia/Nemotron-Labs-Audex-30B-A3B --deploy-config vllm_omni/deploy/audex_<mode>_30b.yaml`.
+
+#### Verification
+
+Same curl / `client.py` as the 2B with
+`"model": "nvidia/Nemotron-Labs-Audex-30B-A3B"` (the server validates the
+model id). First launch downloads ~60 GB.
+
+#### Notes
+
+- Memory usage: ~61 GiB weights on one H100 80 GB (stage 0 at 0.85 GPU
+  fraction + decoder at 0.08 on the same card); verified with healthy KV
+  headroom. If long sequences OOM, set `tensor_parallel_size: 2` on
+  stage 0 (documented in the yaml headers).
+- `enable_prefix_caching: false` on stage 0 — hybrid Mamba constraint.
+- CUDA graphs stay ON for the thinker: the NemotronH stage captures
+  cleanly and eager mode costs ~7× single-stream TTS RTF (review-measured
+  0.30 vs 2.08).
+- Quality (en-24 self-transcribed CER): 0.72% unguided baseline; 0.72% at
+  cfg 1.5 with guided temperature 0.1.
+- Known limitation: the S2S chat pass needs the official
+  formatting-instruction prompt (the bundled examples/client send it);
+  a bare transcript turn makes the 30B answer in speech-codec tokens.

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) |
@@ -57,6 +56,7 @@ recipes/
 | [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
 | [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
 | [`OpenMOSS/MOSS-TTS.md`](./OpenMOSS/MOSS-TTS.md) | Online + offline multilingual TTS (MOSS-TTS family, 8B) | 1x H100 80GB |
+| [`NVIDIA/Nemotron-Labs-Audex.md`](./NVIDIA/Nemotron-Labs-Audex.md) | TTS / text-to-audio / audio understanding / cascaded S2S (Audex 2B + 30B-A3B) | 1x H100 80GB |
 | [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-2512.md`](./Qwen/Qwen-Image-2512.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt FP8 / mixed FP8/NVFP4 | 1x A800 80GB / 2x B200 |
 | [`Qwen/Qwen-Image-Edit.md`](./Qwen/Qwen-Image-Edit.md) | Text-guided single-image editing | 1x or 2x H200 141GB |

tests/config/test_omni_config.py

@@ -392,6 +392,11 @@ def test_sub_config_fields_match_rfc_scopes():
         "enable_multithread_weight_load",
         "num_weight_load_threads",
         "disable_autocast",
+        # Per-stage checkpoint resolution for repos whose stages live in
+        # subfolders (e.g. Audex): mirrors StagePipelineConfig on the
+        # legacy engine-args path.
+        "model_subdir",
+        "tokenizer_subdir",
     }
     assert {f.name for f in fields(OmniStageLoadConfig)} == {
         "load_format",

tests/core/sched/test_omni_ar_scheduler_logprobs.py

@@ -0,0 +1,183 @@
+"""Tests for the AR sampled-token logprob contract."""
+
+from __future__ import annotations
+
+from collections import defaultdict
+from types import SimpleNamespace
+
+import numpy as np
+import pytest
+from vllm.v1.engine import FinishReason
+from vllm.v1.outputs import LogprobsLists
+from vllm.v1.request import RequestStatus
+
+from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler, _slice_sampled_logprobs
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _logprobs(
+    token_ids: list[list[int]],
+    values: list[list[float]],
+    *,
+    cumulative_rows: list[int] | None = None,
+) -> LogprobsLists:
+    return LogprobsLists(
+        logprob_token_ids=np.asarray(token_ids, dtype=np.int32),
+        logprobs=np.asarray(values, dtype=np.float32),
+        sampled_token_ranks=np.zeros(len(token_ids), dtype=np.int32),
+        cu_num_generated_tokens=cumulative_rows,
+    )
+
+
+def test_slices_requested_rows_and_preserves_sampled_token_values() -> None:
+    all_logprobs = _logprobs(
+        [[3, 8], [4, 9], [11, 2], [12, 1]],
+        [[-0.3, -1.0], [-0.4, -1.1], [-0.11, -2.0], [-0.12, -2.1]],
+        cumulative_rows=[0, 2],
+    )
+
+    sliced = _slice_sampled_logprobs(all_logprobs, req_index=1, sampled_token_ids=[11, 12])
+
+    np.testing.assert_array_equal(sliced.logprob_token_ids[:, 0], [11, 12])
+    np.testing.assert_allclose(sliced.logprobs[:, 0], [-0.11, -0.12])
+
+
+def test_rejects_missing_logprobs() -> None:
+    with pytest.raises(RuntimeError, match="model runner returned none"):
+        _slice_sampled_logprobs(None, req_index=0, sampled_token_ids=[7])
+
+
+def test_rejects_row_count_mismatch() -> None:
+    logprobs = _logprobs([[7, 1]], [[-0.7, -1.0]])
+
+    with pytest.raises(RuntimeError, match="row count does not match"):
+        _slice_sampled_logprobs(logprobs, req_index=0, sampled_token_ids=[7, 8])
+
+
+def test_rejects_sampled_token_mismatch() -> None:
+    logprobs = _logprobs([[9, 7]], [[-0.9, -1.0]])
+
+    with pytest.raises(RuntimeError, match="generated_token=7 logprob_token=9"):
+        _slice_sampled_logprobs(logprobs, req_index=0, sampled_token_ids=[7])
+
+
+def test_rejects_non_finite_sampled_logprob() -> None:
+    logprobs = _logprobs([[7, 1]], [[float("nan"), -1.0]])
+
+    with pytest.raises(RuntimeError, match="non-finite values at rows \\[0\\]"):
+        _slice_sampled_logprobs(logprobs, req_index=0, sampled_token_ids=[7])
+
+
+class _Request:
+    def __init__(self, request_id: str) -> None:
+        self.request_id = request_id
+        self.client_index = 0
+        self.status = RequestStatus.RUNNING
+        self.sampling_params = SimpleNamespace(num_logprobs=1)
+        self.num_computed_tokens = 0
+        self.num_in_flight_tokens = 0
+        self.num_output_placeholders = 0
+        self.has_encoder_inputs = False
+        self.pooling_params = None
+        self.resumable = True
+        self.stop_reason = None
+        self.trace_headers = None
+        self.num_nans_in_logits = 0
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+    def get_finished_reason(self):
+        return RequestStatus.get_finished_reason(self.status)
+
+    def take_events(self):
+        return None
+
+    def take_prefill_stats(self):
+        return None
+
+
+class _RequestQueue(list):
+    def remove_requests(self, requests) -> None:
+        for request in requests:
+            if request in self:
+                self.remove(request)
+
+
+def test_invalid_logprobs_finish_only_the_affected_scheduler_request() -> None:
+    """A bad runner response must not bring down unrelated batch requests."""
+    bad = _Request("bad")
+    good = _Request("good")
+    scheduler = SimpleNamespace(
+        perf_metrics=None,
+        connector=None,
+        chunk_transfer_adapter=None,
+        requests={bad.request_id: bad, good.request_id: good},
+        running=[bad, good],
+        waiting=_RequestQueue(),
+        skipped_waiting=_RequestQueue(),
+        structured_output_manager=SimpleNamespace(should_advance=lambda _request: False),
+        transfer_triggered_requests=set(),
+        active_kv_transfers=set(),
+        pending_stop_after_extraction=set(),
+        waiting_for_transfer_free=set(),
+        _new_prompt_len_snapshot={},
+        finished_req_ids=set(),
+        finished_req_ids_dict=defaultdict(set),
+        kv_cache_manager=SimpleNamespace(take_events=lambda: None),
+        kv_event_publisher=SimpleNamespace(publish=lambda _events: None),
+    )
+    update_calls: list[str] = []
+
+    def update_request(request, token_ids):
+        update_calls.append(request.request_id)
+        return token_ids, False
+
+    def free_request(request):
+        scheduler.requests.pop(request.request_id, None)
+        scheduler.finished_req_ids.add(request.request_id)
+        scheduler.finished_req_ids_dict[request.client_index].add(request.request_id)
+        # vLLM 0.26 contract: (kv_xfer_params, ec_xfer_params)
+        return None, None
+
+    scheduler._update_request_with_output = update_request
+    scheduler._process_kv_transfer_trigger = lambda _request, _tokens: False
+    scheduler._handle_stopped_request = lambda _request: True
+    scheduler._free_request = free_request
+    scheduler.make_spec_decoding_stats = lambda *args, **kwargs: None
+    scheduler.make_stats = lambda *args, **kwargs: None
+    scheduler._capture_omni_connector_output = lambda _output: None
+
+    scheduler_output = SimpleNamespace(
+        num_scheduled_tokens={"bad": 1, "good": 1},
+        scheduled_spec_decode_tokens={},
+        num_invalid_spec_tokens=0,
+    )
+    model_runner_output = SimpleNamespace(
+        sampled_token_ids=[[7], [8]],
+        logprobs=_logprobs(
+            [[9, 1], [8, 2]],
+            [[-0.9, -1.0], [-0.8, -1.0]],
+            cumulative_rows=[0, 1],
+        ),
+        prompt_logprobs_dict={},
+        pooler_output=None,
+        num_nans_in_logits=None,
+        kv_connector_output=None,
+        cudagraph_stats=None,
+        req_id_to_index={"bad": 0, "good": 1},
+        routed_experts=None,
+    )
+
+    outputs = OmniARScheduler.update_from_output(scheduler, scheduler_output, model_runner_output)
+    output_by_id = {output.request_id: output for output in outputs[0].outputs}
+
+    assert update_calls == ["good"]
+    assert bad.status is RequestStatus.FINISHED_ERROR
+    assert bad.stop_reason is not None and "misaligned" in bad.stop_reason
+    assert "bad" not in scheduler.requests
+    assert output_by_id["bad"].finish_reason is FinishReason.ERROR
+    assert output_by_id["bad"].new_token_ids == []
+    assert output_by_id["good"].new_token_ids == [8]
+    np.testing.assert_array_equal(output_by_id["good"].new_logprobs.logprob_token_ids[:, 0], [8])

tests/core/sched/test_omni_scheduling_coordinator.py

@@ -121,6 +121,10 @@ class TestFullPayloadCoordinatorSelection(unittest.TestCase):
             ("Qwen3TTSCode2Wav", "code2wav"),
             ("MiniCPMO45Code2Wav", "code2wav"),
             ("CosyVoice3Model", "cosyvoice3_code2wav"),
+            # audex: sync-path TTS decoder (async_chunk: false deploys) and
+            # the TTA XCodec1 stage (always sync full-payload).
+            ("AudexCode2Wav", "audex_code2wav"),
+            ("AudexXCodec1", "audex_xcodec"),
             ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),
             ("DyninOmniForConditionalGeneration", "token2image"),
             ("DyninOmniForConditionalGeneration", "token2audio"),

tests/dfx/perf/scripts/run_diffusion_benchmark.py

@@ -355,6 +355,17 @@ def _kill_process_tree(pid: int) -> None:
 # ---------------------------------------------------------------------------
 
 
+def _resolve_offline_model(model: str) -> str:
+    """Under HF_HUB_OFFLINE, resolve a HF repo id to its local snapshot dir."""
+    import huggingface_hub
+
+    if not model or os.path.isdir(model) or not huggingface_hub.constants.HF_HUB_OFFLINE:
+        return model
+    from huggingface_hub import snapshot_download
+
+    return snapshot_download(model, local_files_only=True)
+
+
 class DiffusionServer:
     """Start a vLLM-Omni diffusion model server as a subprocess.
 
@@ -372,7 +383,7 @@ class DiffusionServer:
         port: int | None = None,
     ) -> None:
         self.server_cfg: dict[str, Any] = server_cfg
-        self.model = server_cfg["model"]
+        self.model = _resolve_offline_model(server_cfg["model"])
         self.serve_args = server_cfg["serve_args"]
         self.host = "127.0.0.1"
         self.port = port if port is not None else get_open_port(self.host)

tests/dfx/perf/tests/test_qwen3_omni_async_chunk.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 2
         }

tests/dfx/perf/tests/test_qwen3_omni_no_async_chunk.json

@@ -5,7 +5,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 2
         }

tests/dfx/perf/tests/test_qwen_image_edit_2511_vllm_omni.json

@@ -66,7 +66,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 4
         }

tests/dfx/perf/tests/test_wan22_i2v_vllm_omni.json

@@ -55,7 +55,8 @@
       {
         "hardware_marks": {
           "res": {
-            "cuda": "H100"
+            "cuda": "H100",
+            "npu": "A3"
           },
           "num_cards": 2
         }

tests/dfx/reliability/invalid_param_test/test_invalid_audio_speech.py

@@ -1,6 +1,6 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""HTTP + WebSocket validation for Qwen3-TTS Base and Higgs Audio V3 ``/v1/audio/speech`` (and related routes)."""
+"""HTTP + WebSocket validation for Qwen3-TTS Base, Higgs Audio V3 and Audex TTA ``/v1/audio/speech`` (and related routes)."""
 
 from __future__ import annotations
 
@@ -70,6 +70,19 @@ _HIGGS_AUDIO_V3_SPEECH = [
     ),
 ]
 
+_AUDEX_TTA_CAPTION = "Heavy rain falling on a tin roof."
+_AUDEX_TTA_SPEECH = [
+    pytest.param(
+        OmniServerParams(
+            model="nvidia/Nemotron-Labs-Audex-2B",
+            stage_config_path=get_deploy_config_path("audex_tta.yaml"),
+            server_args=_SPEECH_SERVER_ARGS,
+        ),
+        id="audex_tta",
+        marks=_L4_SPEECH_HW,
+    ),
+]
+
 
 def _apply_batch_overrides(body: dict[str, object], *, loc: str, overrides: dict[str, object]) -> None:
     if loc == "batch":
@@ -325,6 +338,37 @@ def test_speech_higgs_invalid_field_values(
     )
 
 
+# Audex TTA (``audex_tta`` deployment): the serving layer builds the caption
+# prompt and injects the RVQ phase contract + default guidance server-side, so
+# guidance knobs and speech-only fields must be rejected before generation.
+@pytest.mark.parametrize(
+    "overrides, err_message",
+    [
+        pytest.param({"extra_params": {"cfg_scale": 99}}, "cfg_scale", id="cfg_scale_out_of_range"),
+        pytest.param({"extra_params": {"cfg_scale": "big"}}, "cfg_scale", id="cfg_scale_wrong_type"),
+        pytest.param({"extra_params": {"tta_rvq": {}}}, "tta_rvq", id="tta_rvq_not_overridable"),
+        pytest.param({"input": "   "}, "caption", id="caption_whitespace_only"),
+        pytest.param({"voice": "alloy"}, "voice", id="voice_not_supported"),
+        pytest.param(
+            {"ref_audio": "https://example.com/ref.wav"},
+            "reference audio",
+            id="ref_audio_not_supported",
+        ),
+    ],
+)
+@pytest.mark.parametrize("omni_server", _AUDEX_TTA_SPEECH, indirect=True)
+def test_speech_audex_tta_invalid_field_values(
+    omni_server: OmniServer,
+    openai_client: OpenAIClientHandler,
+    overrides: dict[str, object],
+    err_message: str | tuple[str, ...],
+) -> None:
+    body = {"model": omni_server.model, "input": _AUDEX_TTA_CAPTION, **overrides}
+    openai_client.send_audio_speech_http_request(
+        {"json": body, "timeout": 120, "err_code": 400, "err_message": err_message}
+    )
+
+
 # ─────────────────────────────────────────────────────────────────────────────
 # POST /v1/audio/speech/batch (JSON)
 # ─────────────────────────────────────────────────────────────────────────────

tests/diffusion/models/dmd2/test_dmd2_request_sanitization.py

@@ -8,7 +8,6 @@ import torch
 from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     LTX2I2VDMD2Pipeline,
-    LTX2ImageToVideoPipeline,
     LTX2Pipeline,
     LTX2T2VDMD2Pipeline,
 )
@@ -25,7 +24,7 @@ _DMD2_BASE = {
     WanT2VDMD2Pipeline: Wan22Pipeline,
     WanI2VDMD2Pipeline: Wan22I2VPipeline,
     LTX2T2VDMD2Pipeline: LTX2Pipeline,
-    LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
+    LTX2I2VDMD2Pipeline: LTX2Pipeline,
     FluxDMD2Pipeline: FluxPipeline,
     QwenImageDMD2Pipeline: QwenImagePipeline,
 }

tests/diffusion/models/dmd2/test_dmd2_scheduler.py

@@ -8,7 +8,6 @@ import torch
 from vllm_omni.diffusion.models.flux.pipeline_flux import FluxDMD2Pipeline, FluxPipeline
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     LTX2I2VDMD2Pipeline,
-    LTX2ImageToVideoPipeline,
     LTX2Pipeline,
     LTX2T2VDMD2Pipeline,
 )
@@ -27,7 +26,7 @@ _DMD2_BASE = {
     WanT2VDMD2Pipeline: Wan22Pipeline,
     WanI2VDMD2Pipeline: Wan22I2VPipeline,
     LTX2T2VDMD2Pipeline: LTX2Pipeline,
-    LTX2I2VDMD2Pipeline: LTX2ImageToVideoPipeline,
+    LTX2I2VDMD2Pipeline: LTX2Pipeline,
     FluxDMD2Pipeline: FluxPipeline,
     QwenImageDMD2Pipeline: QwenImagePipeline,
 }

tests/diffusion/models/ltx2/test_ltx2_conditioning.py

@@ -7,6 +7,7 @@ from types import SimpleNamespace
 
 import pytest
 import torch
+from PIL import Image
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -16,17 +17,31 @@ def _make_ltx23_request_pipe(cls):
     torch.nn.Module.__init__(pipe)
     pipe.device = torch.device("cpu")
     pipe.tokenizer_max_length = 99
+    pipe.vae_spatial_compression_ratio = 32
+    pipe.vae_temporal_compression_ratio = 8
     return pipe
 
 
 class TestLTXImageToVideoForwardStages:
-    def test_forward_resolves_request_image_and_delegates_to_shared_forward_impl(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+    def test_i2v_pil_preprocessing_matches_official_center_crop(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_conditioning import _preprocess_i2v_pil_images
+
+        pixels = torch.zeros(4, 8, 3, dtype=torch.uint8)
+        pixels[:, :, 0] = torch.arange(8, dtype=torch.uint8)
+        image = Image.fromarray(pixels.numpy())
+
+        actual = _preprocess_i2v_pil_images(image, height=4, width=4)
+        expected = pixels[:, 2:6].permute(2, 0, 1).unsqueeze(0).float() / 127.5 - 1.0
+
+        torch.testing.assert_close(actual, expected)
+
+    def test_forward_resolves_request_image_and_delegates_to_shared_recipe_runtime(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-        pipe = _make_ltx23_request_pipe(LTX23ImageToVideoPipeline)
+        pipe = _make_ltx23_request_pipe(LTX2Pipeline)
         image = torch.zeros(3, 8, 8)
         req = DiffusionRequestBatch(
             [
@@ -48,27 +63,45 @@ class TestLTXImageToVideoForwardStages:
         )
         seen = {}
 
-        def fake_forward_impl(req_arg, request_inputs, **kwargs):
+        def fake_run_recipe(req_arg, request_inputs, **kwargs):
             seen["req"] = req_arg
             seen["request_inputs"] = request_inputs
             seen["kwargs"] = kwargs
             return ["i2v-delegated"]
 
-        object.__setattr__(pipe, "_forward_impl", fake_forward_impl)
+        object.__setattr__(pipe, "_run_recipe", fake_run_recipe)
 
-        output = pipe.forward(req, noise_scale=0.5)
+        output = pipe.forward(req)
 
         assert output == ["i2v-delegated"]
         assert seen["req"] is req
         assert seen["request_inputs"].prompt == ["make the image move"]
         assert seen["request_inputs"].negative_prompt == ["jitter"]
         assert seen["kwargs"]["image"] is image
-        assert seen["kwargs"]["noise_scale"] == 0.5
+        assert seen["kwargs"]["request_sigmas"] is None
+
+    def test_unified_entry_selects_t2v_without_images_and_rejects_mixed_batches(self):
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        request_inputs = SimpleNamespace(latents=None)
+        t2v_req = SimpleNamespace(prompts=["text only", {"prompt": "also text only"}])
+
+        assert pipe._resolve_request_image(t2v_req, None, request_inputs) is None
+
+        mixed_req = SimpleNamespace(
+            prompts=[
+                "text only",
+                {"prompt": "image conditioned", "multi_modal_data": {"image": torch.zeros(3, 8, 8)}},
+            ]
+        )
+        with pytest.raises(ValueError, match="cannot mix text-to-video and image-to-video"):
+            pipe._resolve_request_image(mixed_req, None, request_inputs)
 
     def test_denoise_timestep_kwargs_masks_video_only(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        pipe = object.__new__(LTX2Pipeline)
         ts = torch.tensor([2.0, 4.0])
         denoise_ctx = SimpleNamespace(
             conditioning_mask=torch.tensor([[1.0, 0.0]]),
@@ -79,30 +112,32 @@ class TestLTXImageToVideoForwardStages:
             ts,
             SimpleNamespace(cfg_parallel_ready=False),
             denoise_ctx,
+            video_token_count=2,
+            audio_token_count=1,
         )
 
         torch.testing.assert_close(kwargs["timestep"], torch.tensor([[0.0, 2.0], [4.0, 0.0]]))
-        torch.testing.assert_close(kwargs["audio_timestep"], ts)
+        torch.testing.assert_close(kwargs["audio_timestep"], ts[:, None])
         torch.testing.assert_close(kwargs["sigma"], ts)
 
 
 class TestLTXImageToVideoConditioning:
     def test_ltx23_i2v_supports_image_input(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        assert LTX23ImageToVideoPipeline.support_image_input is True
+        assert LTX2Pipeline.support_image_input is True
 
     def test_ltx23_i2v_rejects_multi_image_prompt_list(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         image = object()
 
-        assert LTX23ImageToVideoPipeline._resolve_single_prompt_image([image]) is image
+        assert LTX2Pipeline._resolve_single_prompt_image([image]) is image
         with pytest.raises(ValueError, match="exactly one image per prompt"):
-            LTX23ImageToVideoPipeline._resolve_single_prompt_image([object(), object()])
+            LTX2Pipeline._resolve_single_prompt_image([object(), object()])
 
     def test_ltx23_i2v_additional_image_resolution_is_tensor_safe(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         image = torch.zeros(1, 3, 4, 4)
         additional = {
@@ -111,19 +146,24 @@ class TestLTXImageToVideoConditioning:
             "image": torch.ones_like(image),
         }
 
-        assert LTX23ImageToVideoPipeline._resolve_additional_image(additional) is image
+        assert LTX2Pipeline._resolve_additional_image(additional) is image
 
     def test_ltx23_i2v_packed_latents_are_not_noised(self, monkeypatch):
         import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
         import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
         pipe.vae_spatial_compression_ratio = 1
         pipe.vae_temporal_compression_ratio = 1
         pipe.transformer_spatial_patch_size = 1
         pipe.transformer_temporal_patch_size = 1
+        object.__setattr__(
+            pipe,
+            "_encode_i2v_image_latents",
+            lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
+        )
 
         def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
             raise AssertionError("packed I2V latents should not be noised")
@@ -134,7 +174,7 @@ class TestLTXImageToVideoConditioning:
         latents = torch.tensor([[[10.0, 11.0], [20.0, 21.0], [30.0, 31.0]]])
 
         out, conditioning_mask = pipe.prepare_latents(
-            image=None,
+            image=torch.zeros(1, 3, 1, 1),
             batch_size=1,
             num_channels_latents=2,
             height=1,
@@ -147,14 +187,14 @@ class TestLTXImageToVideoConditioning:
         )
 
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
-        torch.testing.assert_close(out, latents)
+        torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [20.0, 21.0], [30.0, 31.0]]]))
 
     def test_ltx23_i2v_5d_latents_noise_preserves_conditioning_frame(self, monkeypatch):
         import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
         import vllm_omni.diffusion.models.ltx2.ltx2_latents as ltx2_latents
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
         pipe.vae_spatial_compression_ratio = 1
         pipe.vae_temporal_compression_ratio = 1
@@ -165,8 +205,16 @@ class TestLTXImageToVideoConditioning:
             latents_std=torch.ones(2),
             config=SimpleNamespace(scaling_factor=1.0),
         )
+        object.__setattr__(
+            pipe,
+            "_encode_i2v_image_latents",
+            lambda *_args, **_kwargs: torch.tensor([[[[[40.0]]], [[[41.0]]]]]),
+        )
+
+        sampled_shapes = []
 
         def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            sampled_shapes.append(tuple(shape))
             return torch.ones(shape, device=device, dtype=dtype)
 
         monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
@@ -175,7 +223,7 @@ class TestLTXImageToVideoConditioning:
         latents = torch.tensor([[[[[10.0]], [[20.0]], [[30.0]]], [[[11.0]], [[21.0]], [[31.0]]]]])
 
         out, conditioning_mask = pipe.prepare_latents(
-            image=None,
+            image=torch.zeros(1, 3, 1, 1),
             batch_size=1,
             num_channels_latents=2,
             height=1,
@@ -187,20 +235,64 @@ class TestLTXImageToVideoConditioning:
             latents=latents,
         )
 
+        torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
+        torch.testing.assert_close(out, torch.tensor([[[40.0, 41.0], [1.0, 1.0], [1.0, 1.0]]]))
+        assert sampled_shapes == [(1, 3, 2)]
+
+    def test_ltx23_i2v_image_noise_is_sampled_after_packing(self, monkeypatch):
+        import vllm_omni.diffusion.models.ltx2.ltx2_conditioning as ltx2_conditioning
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        torch.nn.Module.__init__(pipe)
+        pipe.vae_spatial_compression_ratio = 1
+        pipe.vae_temporal_compression_ratio = 1
+        pipe.transformer_spatial_patch_size = 1
+        pipe.transformer_temporal_patch_size = 1
+        pipe.vae = SimpleNamespace(
+            encode=lambda image: image,
+            latents_mean=torch.zeros(2),
+            latents_std=torch.ones(2),
+        )
+        monkeypatch.setattr(
+            ltx2_conditioning,
+            "retrieve_latents",
+            lambda *_args, **_kwargs: torch.tensor([[[[[10.0]]], [[[11.0]]]]]),
+        )
+        sampled_shapes = []
+
+        def fake_randn_tensor(shape, generator=None, device=None, dtype=None):
+            sampled_shapes.append(tuple(shape))
+            return torch.ones(shape, device=device, dtype=dtype)
+
+        monkeypatch.setattr(ltx2_conditioning, "randn_tensor", fake_randn_tensor)
+
+        out, conditioning_mask = pipe.prepare_latents(
+            image=torch.zeros(1, 3, 1, 1),
+            batch_size=1,
+            num_channels_latents=2,
+            height=1,
+            width=1,
+            num_frames=3,
+            noise_scale=1.0,
+            dtype=torch.float32,
+            device=torch.device("cpu"),
+        )
+
+        assert sampled_shapes == [(1, 3, 2)]
         torch.testing.assert_close(conditioning_mask, torch.tensor([[1.0, 0.0, 0.0]]))
         torch.testing.assert_close(out, torch.tensor([[[10.0, 11.0], [1.0, 1.0], [1.0, 1.0]]]))
 
     def test_ltx23_i2v_video_step_preserves_conditioning_frame(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23ImageToVideoPipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        pipe = object.__new__(LTX23ImageToVideoPipeline)
+        pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
         pipe.transformer_spatial_patch_size = 1
         pipe.transformer_temporal_patch_size = 1
 
         class FakeScheduler:
-            def step(self, noise_pred, t, latents, return_dict=False):
-                return (latents + noise_pred + t,)
+            sigmas = torch.tensor([1.0, 0.5])
 
         pipe.scheduler = FakeScheduler()
         latents = torch.tensor([[[1.0], [2.0], [3.0]]])
@@ -209,11 +301,103 @@ class TestLTXImageToVideoConditioning:
         out = pipe._step_video_latents_i2v(
             noise_pred,
             latents,
-            torch.tensor(0.5),
+            0,
             latent_num_frames=3,
             latent_height=1,
             latent_width=1,
         )
 
         torch.testing.assert_close(out[:, :1], latents[:, :1])
-        torch.testing.assert_close(out[:, 1:], latents[:, 1:] + noise_pred[:, 1:] + 0.5)
+        torch.testing.assert_close(out[:, 1:], latents[:, 1:] - 0.5 * noise_pred[:, 1:])
+
+    def test_shared_step_adapter_uses_official_euler_for_t2v(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
+
+        pipeline = SimpleNamespace(scheduler=SimpleNamespace(sigmas=torch.tensor([1.0, 0.5])))
+        adapter = LTXVideoAudioStepAdapter(
+            pipeline,
+            SimpleNamespace(sigmas=torch.tensor([1.0, 0.25])),
+            latent_num_frames=1,
+            latent_height=1,
+            latent_width=1,
+            image_conditioned=False,
+        )
+        video = torch.tensor([[[2.0]]])
+        audio = torch.tensor([[[4.0]]])
+
+        ((video, audio),) = adapter.step(
+            (torch.tensor([[[2.0]]]), torch.tensor([[[4.0]]])),
+            (torch.tensor(1.0), torch.tensor(1.0)),
+            (video, audio),
+        )
+
+        torch.testing.assert_close(video, torch.tensor([[[1.0]]]))
+        torch.testing.assert_close(audio, torch.tensor([[[1.0]]]))
+
+    def test_shared_step_adapter_uses_official_euler_for_i2v_audio(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXVideoAudioStepAdapter
+
+        class FakePipeline:
+            scheduler = object()
+
+            def __init__(self):
+                self.step_indices = []
+
+            def _step_video_latents_i2v(
+                self,
+                noise_pred,
+                latents,
+                step_index,
+                latent_num_frames,
+                latent_height,
+                latent_width,
+            ):
+                del noise_pred, latent_num_frames, latent_height, latent_width
+                self.step_indices.append(step_index)
+                return latents
+
+        pipeline = FakePipeline()
+        adapter = LTXVideoAudioStepAdapter(
+            pipeline,
+            SimpleNamespace(sigmas=torch.tensor([1.0, 0.5, 0.0])),
+            latent_num_frames=1,
+            latent_height=1,
+            latent_width=1,
+            image_conditioned=True,
+        )
+        video = torch.tensor([[[1.0]]])
+        audio = torch.tensor([[[2.0]]])
+        video_velocity = torch.zeros_like(video)
+        audio_velocity = torch.tensor([[[4.0]]])
+
+        ((video, audio),) = adapter.step(
+            (video_velocity, audio_velocity),
+            (torch.tensor(1.0), torch.tensor(1.0)),
+            (video, audio),
+        )
+        torch.testing.assert_close(audio, torch.tensor([[[0.0]]]))
+
+        ((video, audio),) = adapter.step(
+            (video_velocity, audio_velocity),
+            (torch.tensor(0.5), torch.tensor(0.5)),
+            (video, audio),
+        )
+        torch.testing.assert_close(audio, torch.tensor([[[-2.0]]]))
+        assert pipeline.step_indices == [0, 1]
+
+    def test_i2v_guidance_uses_zero_sigma_for_conditioned_tokens(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_denoise import LTXDenoiseContext
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
+
+        pipe = object.__new__(LTX2Pipeline)
+        denoise_ctx = LTXDenoiseContext(
+            latents=torch.empty(1, 3, 1),
+            audio_latents=torch.empty(1, 1, 1),
+            video_coords=torch.empty(1),
+            audio_coords=torch.empty(1),
+            conditioning_mask=torch.tensor([[1.0, 0.0, 0.0]]),
+        )
+
+        actual = pipe._video_guidance_model_sigma(torch.tensor(0.75), denoise_ctx)
+
+        torch.testing.assert_close(actual, torch.tensor([[[0.0], [0.75], [0.75]]]))

tests/diffusion/models/ltx2/test_ltx2_latents.py

@@ -4,7 +4,7 @@ import pytest
 import torch
 
 from vllm_omni.diffusion.models.ltx2 import ltx2_latents as latent_ops
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline, LTX23Pipeline
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -23,8 +23,51 @@ def _make_pipeline(pipeline_cls, sequence_parallel_size: int = 1):
     return pipeline
 
 
+def test_prepare_video_latents_samples_directly_in_packed_token_space():
+    pipeline = _make_pipeline(LTX2Pipeline)
+    pipeline.vae_spatial_compression_ratio = 8
+    pipeline.vae_temporal_compression_ratio = 8
+    pipeline.transformer_spatial_patch_size = 2
+    pipeline.transformer_temporal_patch_size = 1
+
+    expected_generator = torch.Generator().manual_seed(42)
+    expected = torch.randn((1, 32, 16), generator=expected_generator)
+    actual = pipeline.prepare_latents(
+        batch_size=1,
+        num_channels_latents=4,
+        height=64,
+        width=64,
+        num_frames=9,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
+def test_prepare_audio_latents_samples_directly_in_packed_token_space():
+    pipeline = _make_pipeline(LTX2Pipeline)
+
+    expected_generator = torch.Generator().manual_seed(42)
+    expected = torch.randn((1, 3, 32), generator=expected_generator)
+    actual, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
+        batch_size=1,
+        num_channels_latents=2,
+        num_mel_bins=64,
+        audio_latent_length=3,
+        dtype=torch.float32,
+        device=torch.device("cpu"),
+        generator=torch.Generator().manual_seed(42),
+    )
+
+    assert original_num_frames == 3
+    assert padded_num_frames == 3
+    torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
+
 def test_prepare_audio_latents_pads_generated_dummy_length_for_sp():
-    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=2)
+    pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=2)
 
     latents, original_num_frames, padded_num_frames = pipeline.prepare_audio_latents(
         batch_size=1,
@@ -66,8 +109,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack():
     padded = torch.cat([original, torch.full((1, 2, 4), 999.0)], dim=1)
 
     unpadded = latent_ops.unpad_audio_latents(padded, 10)
-    unpacked = latent_ops.unpack_audio_latents(unpadded, latent_length=10, num_mel_bins=2)
-    expected = latent_ops.unpack_audio_latents(original, latent_length=10, num_mel_bins=2)
+    unpacked = latent_ops.unpack_audio_latents(unpadded, num_mel_bins=2)
+    expected = latent_ops.unpack_audio_latents(original, num_mel_bins=2)
 
     assert unpacked.shape == (1, 2, 10, 2)
     assert not (unpacked == 999.0).any()
@@ -75,7 +118,8 @@ def test_unpad_audio_latents_restores_original_frames_before_unpack():
 
 
 def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
-    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
+    pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
+    pipeline.preserve_sp_padded_audio_duration = True
     latents = torch.arange(96, dtype=torch.float32).view(1, 2, 12, 4)
 
     audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
@@ -97,7 +141,7 @@ def test_prepare_audio_latents_accepts_already_padded_4d_latents_for_sp():
 
 
 def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
-    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
+    pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
     latents = torch.zeros(1, 2, 13, 4)
 
     audio_latent_length = pipeline._resolve_audio_latent_length(10, latents)
@@ -106,7 +150,7 @@ def test_resolve_audio_latent_length_preserves_legacy_4d_shape_inference():
 
 
 def test_prepare_audio_latents_rejects_incompatible_provided_length():
-    pipeline = _make_pipeline(LTX23Pipeline, sequence_parallel_size=4)
+    pipeline = _make_pipeline(LTX2Pipeline, sequence_parallel_size=4)
     latents = torch.zeros(1, 11, 4)
 
     with pytest.raises(ValueError, match="incompatible audio frame count"):

tests/diffusion/models/ltx2/test_ltx2_parallel.py

@@ -15,9 +15,9 @@ class TestCFGParallelHelpers:
     """Test LTX-2.3 CFG helper math without loading model weights."""
 
     def test_combine_cfg_noise_matches_x0_space_formula(self):
-        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX23Pipeline
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-        pipe = object.__new__(LTX23Pipeline)
+        pipe = object.__new__(LTX2Pipeline)
         video_sample = torch.tensor([[[1.0, -2.0]]])
         audio_sample = torch.tensor([[[0.5, 3.0]]])
         video_pos = torch.tensor([[[0.2, -0.3]]])
@@ -50,93 +50,53 @@ class TestCFGParallelHelpers:
         assert torch.allclose(video_combined, expected_video)
         assert torch.allclose(audio_combined, expected_audio)
 
-    def test_two_rank_cfg_parallel_smoke_uses_rank_local_branch_and_x0_formula(self, monkeypatch):
-        from vllm_omni.diffusion.models.ltx2 import ltx2_guidance
-        from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
-
-        pipe = object.__new__(ltx23.LTX23Pipeline)
-        video_sample = torch.tensor([[[1.0, -2.0]]])
-        audio_sample = torch.tensor([[[0.5, 3.0, -1.0]]])
-        video_pos = torch.tensor([[[0.2, -0.3]]])
-        video_neg = torch.tensor([[[-0.4, 0.1]]])
-        audio_pos = torch.tensor([[[0.7, -0.2, 0.3]]])
-        audio_neg = torch.tensor([[[0.1, 0.4, -0.5]]])
-        video_sigma = torch.tensor(0.25)
-        audio_sigma = torch.tensor(0.5)
-        scale = 4.0
-
-        class FakeCfgGroup:
-            def all_gather(self, tensor, separate_tensors=True):
-                assert separate_tensors
-                if tensor.shape == video_pos.shape:
-                    return [video_pos, video_neg]
-                return [audio_pos, audio_neg]
-
-        monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
-        monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
-
-        expected_video = ltx2_guidance.combine_velocity_via_x0(
-            video_sample,
-            video_pos,
-            video_neg,
-            video_sigma,
-            scale,
-        )
-        expected_audio = ltx2_guidance.combine_velocity_via_x0(
-            audio_sample,
-            audio_pos,
-            audio_neg,
-            audio_sigma,
-            scale,
+    def test_cfg_parallel_rejects_non_cfg_guidance(self):
+        from vllm_omni.diffusion.models.ltx2.ltx2_guidance import (
+            LTX_GUIDANCE_EXECUTOR,
+            LTXGuidancePlan,
         )
+        from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
 
-        for rank, expected_branch in ((0, "positive"), (1, "negative")):
-            calls = []
-            monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda rank=rank: rank)
+        plan = LTXGuidancePlan.build(LTX23_ONE_STAGE_RECIPE.request_guidance)
 
-            def fake_predict_noise(**kwargs):
-                calls.append(kwargs["branch"])
-                if kwargs["branch"] == "positive":
-                    return video_pos, audio_pos
-                return video_neg, audio_neg
+        with pytest.raises(ValueError, match="CFG-only guidance"):
+            LTX_GUIDANCE_EXECUTOR.validate_cfg_world_size(plan, 2)
 
-            object.__setattr__(pipe, "predict_noise", fake_predict_noise)
-            video_combined, audio_combined = pipe.predict_noise_with_parallel_cfg(
-                true_cfg_scale=scale,
-                positive_kwargs={"branch": "positive"},
-                negative_kwargs={"branch": "negative"},
-                cfg_normalize=False,
-                video_latents=video_sample,
-                audio_latents=audio_sample,
-                video_sigma=video_sigma,
-                audio_sigma=audio_sigma,
-            )
+    def test_cfg_parallel_dummy_warms_supported_cfg_only_plan(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
+        from vllm_omni.diffusion.models.ltx2.ltx2_recipes import LTX23_ONE_STAGE_RECIPE
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
-            assert calls == [expected_branch]
-            torch.testing.assert_close(video_combined, expected_video)
-            torch.testing.assert_close(audio_combined, expected_audio)
+        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
+        pipe = object.__new__(LTX2Pipeline)
+        req = SimpleNamespace(is_dummy_run=lambda: True)
+        request_inputs = SimpleNamespace(guidance=LTX23_ONE_STAGE_RECIPE.request_guidance)
 
-        assert "_cfg_video_latents" not in pipe.__dict__
-        assert "_cfg_audio_latents" not in pipe.__dict__
+        cfg_parallel_ready = pipe._setup_forward_runtime(req, request_inputs, attention_kwargs=None)
+
+        assert cfg_parallel_ready is True
+        assert pipe._guidance_plan.names == ("cond", "uncond")
+        assert pipe._guidance_plan.spec.video.cfg_scale == 3.0
+        assert pipe._guidance_plan.spec.audio.cfg_scale == 7.0
 
 
 class TestCFGParallelForwardPath:
     """Test the LTX-2.3 CFG-parallel denoising path without loading model weights."""
 
-    @pytest.mark.parametrize("pipeline_cls_name", ["LTX2Pipeline", "LTX23Pipeline"])
-    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch, pipeline_cls_name):
-        from vllm_omni.diffusion.models.ltx2 import ltx2_pipeline_runtime
+    def test_forward_collates_request_prompt_embeds_and_mask_aliases(self, monkeypatch):
+        from vllm_omni.diffusion.models.ltx2 import ltx2_runtime
         from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-        pipeline_cls = getattr(ltx23, pipeline_cls_name)
-        pipe = object.__new__(pipeline_cls)
+        pipe = object.__new__(ltx23.LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
         pipe.device = torch.device("cpu")
         pipe.tokenizer_max_length = 4
-        monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
+        pipe.vae_spatial_compression_ratio = 32
+        pipe.vae_temporal_compression_ratio = 8
+        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 1)
 
         class StopAtEncodePromptError(Exception):
             pass
@@ -224,7 +184,7 @@ class TestCFGParallelForwardPath:
         ("frame_rate_input", "audio_sampling_rate", "expected_frame_rate"),
         [(1.0, 1, 1.0), (None, 24, 24.0)],
     )
-    def test_forward_cfg_parallel_steps_video_and_audio_scheduler(
+    def test_forward_cfg_parallel_uses_shared_video_audio_euler_step(
         self,
         monkeypatch,
         cfg_rank,
@@ -233,13 +193,13 @@ class TestCFGParallelForwardPath:
         audio_sampling_rate,
         expected_frame_rate,
     ):
-        from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_pipeline_runtime
+        from vllm_omni.diffusion.models.ltx2 import ltx2_denoise, ltx2_guidance, ltx2_runtime
         from vllm_omni.diffusion.models.ltx2 import pipeline_ltx2 as ltx23
         from vllm_omni.diffusion.request import OmniDiffusionRequest
         from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
         from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
-        pipe = object.__new__(ltx23.LTX23Pipeline)
+        pipe = object.__new__(ltx23.LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
         pipe.device = torch.device("cpu")
         pipe.tokenizer_max_length = 1
@@ -278,7 +238,7 @@ class TestCFGParallelForwardPath:
                     return [audio_pos, audio_neg]
                 raise AssertionError(f"Unexpected gathered tensor: {tensor}")
 
-        monkeypatch.setattr(ltx2_pipeline_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
+        monkeypatch.setattr(ltx2_runtime, "get_classifier_free_guidance_world_size", lambda: 2)
         monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_world_size", lambda: 2)
         monkeypatch.setattr(ltx2_guidance, "get_classifier_free_guidance_rank", lambda: cfg_rank)
         monkeypatch.setattr(ltx2_guidance, "get_cfg_group", lambda: FakeCfgGroup())
@@ -314,7 +274,7 @@ class TestCFGParallelForwardPath:
 
             def __call__(self, prompt_embeds, prompt_attention_mask, padding_side):
                 assert padding_side == "left"
-                assert prompt_embeds.shape[0] == 2
+                assert prompt_embeds.shape[0] == 1
                 return prompt_embeds, prompt_embeds, prompt_attention_mask
 
         rope_video_fps: list[float] = []
@@ -380,7 +340,16 @@ class TestCFGParallelForwardPath:
                 num_frames=1,
                 frame_rate=frame_rate_input,
                 num_inference_steps=2,
-                guidance_scale=4.0,
+                extra_args={
+                    "video_cfg_scale": 3.0,
+                    "audio_cfg_scale": 7.0,
+                    "video_stg_scale": 0.0,
+                    "audio_stg_scale": 0.0,
+                    "video_modality_scale": 1.0,
+                    "audio_modality_scale": 1.0,
+                    "video_rescale_scale": 0.0,
+                    "audio_rescale_scale": 0.0,
+                },
                 latents=video_latents,
                 audio_latents=audio_latents,
                 output_type="latent",
@@ -388,35 +357,31 @@ class TestCFGParallelForwardPath:
             request_id="ltx23-cfg-parallel-forward-test",
         )
 
-        output = pipe.forward(DiffusionRequestBatch(requests=[req]))[0]
+        output = pipe.forward(DiffusionRequestBatch(requests=[req]), sigmas=[0.25, 0.25])[0]
 
         expected_video_noise = ltx2_guidance.combine_velocity_via_x0(
             video_latents,
             video_pos,
             video_neg,
             pipe.scheduler.sigmas[0],
-            4.0,
+            3.0,
         )
         expected_audio_noise = ltx2_guidance.combine_velocity_via_x0(
             audio_latents,
             audio_pos,
             audio_neg,
             pipe.scheduler.sigmas[0],
-            4.0,
+            7.0,
         )
-        scheduler_call_names = [call[0] for call in pipe.scheduler.calls]
-        assert scheduler_call_names == ["video", "audio", "video", "audio"]
+        assert pipe.scheduler.calls == []
         assert len(pipe.transformer.calls) == 2
-        torch.testing.assert_close(pipe.scheduler.calls[0][1], expected_video_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[1][1], expected_audio_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[2][1], expected_video_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[3][1], expected_audio_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[2][3], video_latents - expected_video_noise)
-        torch.testing.assert_close(pipe.scheduler.calls[3][3], audio_latents - expected_audio_noise)
 
         video_out, audio_out = output.output
-        torch.testing.assert_close(video_out, (video_latents - 2 * expected_video_noise).reshape(1, 2, 1, 1, 1))
-        torch.testing.assert_close(audio_out, (audio_latents - 2 * expected_audio_noise).reshape(1, 1, 1, 2))
+        sigma_delta = pipe.scheduler.sigmas[-1] - pipe.scheduler.sigmas[0]
+        expected_video = video_latents + expected_video_noise * sigma_delta
+        expected_audio = audio_latents + expected_audio_noise * sigma_delta
+        torch.testing.assert_close(video_out, expected_video.reshape(1, 2, 1, 1, 1))
+        torch.testing.assert_close(audio_out, expected_audio.reshape(1, 1, 1, 2))
 
         # fps regression guard: an omitted request fps (frame_rate_input=None) must resolve
         # to the model's own 24.0 default, not crash on None; a provided rate is passed through.

tests/diffusion/models/ltx2/test_ltx2_transformer.py

@@ -10,11 +10,12 @@ from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransforme
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
-def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
+def test_ltx_rms_norm_no_affine_matches_official_torch_module():
     norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
 
     assert "weight" not in dict(norm.named_parameters())
-    assert "weight" in dict(norm.named_buffers())
+    assert norm.weight is None
+    assert "weight" not in dict(norm.named_buffers())
     assert "weight" not in norm.state_dict()
 
 
@@ -50,3 +51,16 @@ def test_ltx_transformer_exposes_hsdp_shard_conditions_for_blocks():
             matched.append(name)
 
     assert matched == ["transformer_blocks.0", "transformer_blocks.1"]
+
+
+@pytest.mark.parametrize("rope_type", ["interleaved", "split"])
+def test_ltx_sp_plan_shards_video_and_audio_timesteps_together(rope_type):
+    root_plan = LTX2VideoTransformer3DModel._build_sp_plan(rope_type)[""]
+
+    video_timestep = root_plan["timestep"]
+    audio_timestep = root_plan["audio_timestep"]
+
+    assert video_timestep.split_dim == audio_timestep.split_dim == 1
+    assert video_timestep.expected_dims == audio_timestep.expected_dims == 2
+    assert not video_timestep.split_output
+    assert not audio_timestep.split_output

tests/diffusion/models/ltx2/test_ltx2_vae.py

@@ -14,7 +14,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 class TestLTXOutputRank:
     @pytest.mark.parametrize(
         ("distributed_vae_state", "expected_decode_calls"),
-        [(False, 0), (True, 1), (RuntimeError("unavailable"), 0)],
+        [(False, 0), (True, 1), (RuntimeError("unavailable"), None)],
     )
     def test_non_output_rank_only_enters_collective_vae_decode(
         self,
@@ -22,7 +22,8 @@ class TestLTXOutputRank:
         distributed_vae_state,
         expected_decode_calls,
     ):
-        import vllm_omni.diffusion.models.ltx2.pipeline_ltx2 as ltx
+        import vllm_omni.diffusion.models.ltx2.ltx2_runtime as ltx_runtime
+        from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import LTX2Pipeline
 
         class FakeVae:
             dtype = torch.float32
@@ -40,23 +41,30 @@ class TestLTXOutputRank:
                 self.decode_calls += 1
                 return (torch.ones(1, 1),)
 
-        monkeypatch.setattr(ltx.torch.distributed, "is_initialized", lambda: True)
-        monkeypatch.setattr(ltx.torch.distributed, "get_rank", lambda: 1)
-        pipe = object.__new__(ltx.LTX2Pipeline)
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "is_initialized", lambda: True)
+        monkeypatch.setattr(ltx_runtime.torch.distributed, "get_rank", lambda: 1)
+        pipe = object.__new__(LTX2Pipeline)
         torch.nn.Module.__init__(pipe)
         pipe.vae = FakeVae()
 
-        output = pipe._decode_output(
-            latents=torch.ones(1, 1),
-            audio_latents=torch.ones(1, 1),
-            output_type="np",
-            connector_prompt_embeds=torch.ones(1, 1),
-            generator=None,
-            device=torch.device("cpu"),
-            decode_timestep=0.0,
-            decode_noise_scale=None,
-            prompt_batch_size=1,
-        )
+        decode_kwargs = {
+            "latents": torch.ones(1, 1),
+            "audio_latents": torch.ones(1, 1),
+            "output_type": "np",
+            "connector_prompt_embeds": torch.ones(1, 1),
+            "generator": None,
+            "device": torch.device("cpu"),
+            "decode_timestep": 0.0,
+            "decode_noise_scale": None,
+            "prompt_batch_size": 1,
+        }
+        if isinstance(distributed_vae_state, Exception):
+            with pytest.raises(RuntimeError, match="unavailable"):
+                pipe._decode_output(**decode_kwargs)
+            assert pipe.vae.decode_calls == 0
+            return
+
+        output = pipe._decode_output(**decode_kwargs)
 
         assert pipe.vae.decode_calls == expected_decode_calls
         assert output.output[0].numel() == 0

tests/distributed/omni_connectors/test_bagel_shared_memory_connector.py

@@ -228,6 +228,7 @@ def _resolve_deploy_config(config_path: str, run_level: str) -> str:
 
 
 @pytest.mark.advanced_model
+@pytest.mark.skip(reason="#issue 5475")
 @pytest.mark.diffusion
 @hardware_test(res={"cuda": "H100", "rocm": "MI325"})
 def test_bagel_img2img_shared_memory_connector(run_level):

tests/dreamzero/test_deploy_configs.py

@@ -0,0 +1,57 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""Every shipped DreamZero deploy config must select the AR-Diffusion engine.
+
+``DreamZeroPipeline.__init__`` rejects any other engine, so a config missing
+``engine_backend`` fails at worker startup rather than at request time. That
+only surfaces in slow multi-GPU e2e serving tests (see issue #5046), so assert
+it here on CPU instead.
+"""
+
+from pathlib import Path
+from types import SimpleNamespace
+
+import pytest
+import yaml
+
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
+from vllm_omni.experimental.ar_diffusion.engine import ARDiffusionEngine
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+# Discover siblings of a known config so a future dreamzero*.yaml is covered
+# without editing this test.
+DEPLOY_DIR = Path(get_deploy_config_path("dreamzero.yaml")).parent
+
+
+def _dreamzero_stages() -> list[tuple[str, int, dict]]:
+    stages = []
+    for name in sorted(path.name for path in DEPLOY_DIR.glob("dreamzero*.yaml")):
+        config = yaml.safe_load(Path(get_deploy_config_path(name)).read_text())
+        for index, stage in enumerate(config.get("stages") or []):
+            stages.append((name, index, stage))
+    return stages
+
+
+def test_dreamzero_deploy_configs_exist() -> None:
+    assert _dreamzero_stages(), f"no DreamZero deploy configs found under {DEPLOY_DIR}"
+
+
+@pytest.mark.parametrize(
+    ("config_name", "stage_index", "stage"),
+    _dreamzero_stages(),
+    ids=lambda value: value if isinstance(value, str) else None,
+)
+def test_dreamzero_stage_selects_ar_diffusion_engine(config_name: str, stage_index: int, stage: dict) -> None:
+    backend = stage.get("engine_backend")
+    assert backend, (
+        f"{config_name} stage {stage_index} does not set engine_backend; "
+        "DreamZeroPipeline requires the AR-Diffusion engine"
+    )
+    # Resolving also proves the import path is real, not just a lookalike string.
+    engine_class = DiffusionEngine.resolve_engine_class(SimpleNamespace(engine_backend=backend))
+    assert issubclass(engine_class, ARDiffusionEngine), (
+        f"{config_name} stage {stage_index} resolves to {engine_class.__name__}, not an AR-Diffusion engine"
+    )

tests/e2e/accuracy/ltx/__init__.py

@@ -0,0 +1 @@
+"""LTX family official-reference accuracy tests."""

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()

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

tests/e2e/accuracy/test_ltx2_3_video_similarity.py

@@ -1,413 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-SSIM/PSNR accuracy tests for LTX-2.3.
-
-1. **Transformer parity** (``test_ltx2_3_transformer_matches_diffusers``):
-   Swaps our custom transformer into diffusers' ``LTX2Pipeline`` to measure
-   numerical parity in isolation.  Thresholds: SSIM >= 0.95, PSNR >= 28 dB.
-   Result: SSIM 0.999987 (bit-identical).
-
-2. **Full pipeline** (``test_ltx2_3_pipeline_matches_diffusers``):
-   Runs the full vLLM-Omni serving stack (``OmniServer`` -> HTTP API) and
-   compares per-frame against stock diffusers.  Currently skipped because
-   the OmniServer subprocess creates a different RNG state than in-process
-   diffusers, producing different initial latents from the same seed.
-   This is a test infrastructure limitation, not a model accuracy issue.
-"""
-
-from __future__ import annotations
-
-import gc
-import os
-import tempfile
-from pathlib import Path
-
-import diffusers
-import numpy as np
-import pytest
-import requests
-import torch
-from PIL import Image
-
-from tests.e2e.accuracy.helpers import compute_image_ssim_psnr, model_output_dir
-from tests.helpers.env import run_post_test_cleanup, run_pre_test_cleanup
-from tests.helpers.mark import hardware_test
-from tests.helpers.runtime import OmniServer
-
-# Parse diffusers version for compatibility check
-_DIFFUSERS_VERSION = tuple(int(x) for x in diffusers.__version__.split(".")[:2] if x.isdigit())
-_DIFFUSERS_038 = _DIFFUSERS_VERSION >= (0, 38)
-
-MODEL_ID = "diffusers/LTX-2.3-Diffusers"
-MODEL_ENV_VAR = "VLLM_TEST_LTX23_MODEL"
-PROMPT = "A lighthouse on a rocky cliff at sunset, waves crashing below, golden hour lighting"
-NEGATIVE_PROMPT = "blurry, low quality, distorted, watermark"
-WIDTH = 512
-HEIGHT = 384
-NUM_FRAMES = 25  # ~1 second at 24fps
-NUM_INFERENCE_STEPS = 20
-GUIDANCE_SCALE = 4.0
-SEED = 42
-
-# Transformer-swap test: near-identical output expected
-TRANSFORMER_SSIM_THRESHOLD = 0.95
-TRANSFORMER_PSNR_THRESHOLD = 28.0
-
-# Full-pipeline test: allows minor divergence from RNG / pipeline differences
-PIPELINE_SSIM_THRESHOLD = 0.94
-PIPELINE_PSNR_THRESHOLD = 28.0
-
-
-def _model_name() -> str:
-    return os.environ.get(MODEL_ENV_VAR, MODEL_ID)
-
-
-def _local_files_only(model: str) -> bool:
-    return Path(model).exists()
-
-
-# ---------------------------------------------------------------------------
-# Frame extraction helpers
-# ---------------------------------------------------------------------------
-
-
-def _video_to_frames(video_np: np.ndarray) -> list[Image.Image]:
-    """Convert numpy video to list of PIL Images."""
-    while video_np.ndim > 4:
-        video_np = video_np[0]
-    if video_np.dtype in (np.float32, np.float64, np.float16):
-        video_np = np.clip(video_np * 255, 0, 255).astype(np.uint8)
-    return [Image.fromarray(video_np[t]) for t in range(video_np.shape[0])]
-
-
-def _extract_diffusers_frames(result) -> list[Image.Image]:
-    """Extract frames from diffusers pipeline output."""
-    video = result.frames
-    if isinstance(video, np.ndarray):
-        return _video_to_frames(video)
-    if isinstance(video, list):
-        if isinstance(video[0], list):
-            return [img.convert("RGB") for img in video[0]]
-        if isinstance(video[0], Image.Image):
-            return [img.convert("RGB") for img in video]
-    raise ValueError(f"Unexpected output type: {type(video)}")
-
-
-def _extract_mp4_frames(mp4_bytes: bytes) -> list[Image.Image]:
-    """Extract frames from an MP4 video using ffmpeg."""
-    import subprocess
-
-    with tempfile.TemporaryDirectory() as tmpdir:
-        mp4_path = os.path.join(tmpdir, "video.mp4")
-        with open(mp4_path, "wb") as f:
-            f.write(mp4_bytes)
-
-        # Extract video frames as PNG files using ffmpeg
-        frame_pattern = os.path.join(tmpdir, "frame_%04d.png")
-        subprocess.run(
-            ["ffmpeg", "-i", mp4_path, "-vsync", "0", frame_pattern],
-            capture_output=True,
-            check=True,
-        )
-
-        # Load frames in order
-        frames = []
-        i = 1
-        while True:
-            fpath = os.path.join(tmpdir, f"frame_{i:04d}.png")
-            if not os.path.exists(fpath):
-                break
-            frames.append(Image.open(fpath).convert("RGB").copy())
-            i += 1
-        return frames
-
-
-# ---------------------------------------------------------------------------
-# Comparison helper
-# ---------------------------------------------------------------------------
-
-
-def _assert_video_similarity(
-    *,
-    model_name: str,
-    vllm_frames: list[Image.Image],
-    diffusers_frames: list[Image.Image],
-    ssim_threshold: float,
-    psnr_threshold: float,
-) -> tuple[float, float]:
-    """Compare video frames and assert SSIM/PSNR meet thresholds."""
-    min_frames = min(len(vllm_frames), len(diffusers_frames))
-    assert min_frames > 0, "No frames to compare"
-
-    ssim_scores = []
-    psnr_scores = []
-    for i in range(min_frames):
-        ssim_val, psnr_val = compute_image_ssim_psnr(
-            prediction=vllm_frames[i],
-            reference=diffusers_frames[i],
-        )
-        ssim_scores.append(ssim_val)
-        psnr_scores.append(psnr_val)
-
-    avg_ssim = sum(ssim_scores) / len(ssim_scores)
-    avg_psnr = sum(psnr_scores) / len(psnr_scores)
-
-    print(f"\n{model_name} video similarity ({min_frames} frames):")
-    print(f"  SSIM: avg={avg_ssim:.6f}, min={min(ssim_scores):.6f}, threshold>={ssim_threshold:.6f}")
-    print(f"  PSNR: avg={avg_psnr:.6f} dB, min={min(psnr_scores):.6f} dB, threshold>={psnr_threshold:.6f} dB")
-
-    assert avg_ssim >= ssim_threshold, f"SSIM below threshold: got {avg_ssim:.6f}, expected >= {ssim_threshold:.6f}."
-    assert avg_psnr >= psnr_threshold, f"PSNR below threshold: got {avg_psnr:.6f}, expected >= {psnr_threshold:.6f}."
-    return avg_ssim, avg_psnr
-
-
-# ---------------------------------------------------------------------------
-# Diffusers baseline (shared by both tests)
-# ---------------------------------------------------------------------------
-
-
-def _run_diffusers_baseline(model: str, output_dir: Path) -> list[Image.Image]:
-    """Generate video using stock diffusers LTX2Pipeline."""
-    from diffusers import LTX2Pipeline
-
-    run_pre_test_cleanup()
-    pipe = None
-    try:
-        pipe = LTX2Pipeline.from_pretrained(
-            model, torch_dtype=torch.bfloat16, local_files_only=_local_files_only(model)
-        ).to("cuda")
-
-        generator = torch.Generator(device="cuda").manual_seed(SEED)
-        result = pipe(
-            prompt=PROMPT,
-            negative_prompt=NEGATIVE_PROMPT,
-            width=WIDTH,
-            height=HEIGHT,
-            num_frames=NUM_FRAMES,
-            num_inference_steps=NUM_INFERENCE_STEPS,
-            guidance_scale=GUIDANCE_SCALE,
-            generator=generator,
-            output_type="np",
-        )
-        frames = _extract_diffusers_frames(result)
-        for i, f in enumerate(frames):
-            f.save(output_dir / f"diffusers_frame_{i:04d}.png")
-        return frames
-    finally:
-        del pipe
-        gc.collect()
-        if torch.cuda.is_available():
-            torch.accelerator.empty_cache()
-        run_post_test_cleanup()
-
-
-# ---------------------------------------------------------------------------
-# Test 1: Transformer-swap parity
-# ---------------------------------------------------------------------------
-
-
-def _run_with_custom_transformer(model: str, output_dir: Path) -> list[Image.Image]:
-    """Run diffusers pipeline with our custom transformer swapped in."""
-    from contextlib import nullcontext
-
-    from diffusers import LTX2Pipeline
-    from vllm.config import VllmConfig, set_current_vllm_config
-    from vllm.distributed.parallel_state import init_distributed_environment, initialize_model_parallel
-
-    from vllm_omni.diffusion.models.ltx2 import (
-        create_transformer_from_config,
-        load_transformer_config,
-    )
-
-    vllm_config = VllmConfig()
-    ctx = set_current_vllm_config(vllm_config)
-    ctx.__enter__()
-
-    if not torch.distributed.is_initialized():
-        os.environ.setdefault("MASTER_ADDR", "localhost")
-        os.environ.setdefault("MASTER_PORT", "29503")
-        os.environ.setdefault("RANK", "0")
-        os.environ.setdefault("WORLD_SIZE", "1")
-        init_distributed_environment(world_size=1, rank=0, local_rank=0)
-        initialize_model_parallel(tensor_model_parallel_size=1)
-
-    local = _local_files_only(model)
-    pipe = LTX2Pipeline.from_pretrained(model, torch_dtype=torch.bfloat16, local_files_only=local)
-
-    transformer_config = load_transformer_config(model, "transformer", local)
-    our_transformer = create_transformer_from_config(transformer_config)
-
-    diffusers_state = dict(pipe.transformer.named_parameters())
-
-    def _weight_iter():
-        for name, param in diffusers_state.items():
-            yield name, param.data
-
-    our_transformer.load_weights(_weight_iter())
-    our_transformer = our_transformer.to(dtype=torch.bfloat16, device="cuda").eval()
-
-    # Compatibility shims for diffusers pipeline
-    our_transformer.dtype = torch.bfloat16
-    if not hasattr(our_transformer, "cache_context"):
-        our_transformer.cache_context = lambda name: nullcontext()
-
-    del pipe.transformer
-    pipe.transformer = our_transformer
-    for name, component in pipe.components.items():
-        if name != "transformer" and hasattr(component, "to"):
-            try:
-                component.to("cuda")
-            except Exception:
-                pass
-
-    generator = torch.Generator(device="cuda").manual_seed(SEED)
-    result = pipe(
-        prompt=PROMPT,
-        negative_prompt=NEGATIVE_PROMPT,
-        width=WIDTH,
-        height=HEIGHT,
-        num_frames=NUM_FRAMES,
-        num_inference_steps=NUM_INFERENCE_STEPS,
-        guidance_scale=GUIDANCE_SCALE,
-        generator=generator,
-        output_type="np",
-    )
-    frames = _extract_diffusers_frames(result)
-    for i, f in enumerate(frames):
-        f.save(output_dir / f"vllm_transformer_frame_{i:04d}.png")
-
-    del pipe, result, our_transformer
-    gc.collect()
-    if torch.cuda.is_available():
-        torch.accelerator.empty_cache()
-    return frames
-
-
-@pytest.mark.full_model
-@pytest.mark.benchmark
-@pytest.mark.diffusion
-@pytest.mark.skipif(
-    not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
-)
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_ltx2_3_transformer_matches_diffusers(accuracy_artifact_root: Path) -> None:
-    """Transformer-level parity: swap our transformer into diffusers pipeline.
-
-    Isolates transformer numerical accuracy from pipeline-level differences.
-    Both runs use diffusers' denoising loop, CFG, scheduler, and RNG.
-    """
-    model = _model_name()
-    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
-
-    diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
-    vllm_frames = _run_with_custom_transformer(model=model, output_dir=output_dir)
-
-    _assert_video_similarity(
-        model_name=f"{MODEL_ID} (transformer-swap)",
-        vllm_frames=vllm_frames,
-        diffusers_frames=diffusers_frames,
-        ssim_threshold=TRANSFORMER_SSIM_THRESHOLD,
-        psnr_threshold=TRANSFORMER_PSNR_THRESHOLD,
-    )
-
-
-# ---------------------------------------------------------------------------
-# Test 2: Full pipeline (OmniServer → HTTP API vs diffusers)
-# ---------------------------------------------------------------------------
-
-
-def _run_vllm_omni_serving(model: str, output_dir: Path) -> list[Image.Image]:
-    """Generate video via the full vLLM-Omni serving stack."""
-    server_args = [
-        "--model-class-name",
-        "LTX23Pipeline",
-        "--stage-init-timeout",
-        "600",
-    ]
-    with OmniServer(model, server_args, use_omni=True) as server:
-        # Submit generation request
-        response = requests.post(
-            f"http://{server.host}:{server.port}/v1/videos",
-            files={
-                "prompt": (None, PROMPT),
-                "negative_prompt": (None, NEGATIVE_PROMPT),
-                "model": (None, server.model),
-                "num_frames": (None, str(NUM_FRAMES)),
-                "fps": (None, "24"),
-                "size": (None, f"{WIDTH}x{HEIGHT}"),
-                "num_inference_steps": (None, str(NUM_INFERENCE_STEPS)),
-                "guidance_scale": (None, str(GUIDANCE_SCALE)),
-                "seed": (None, str(SEED)),
-            },
-            timeout=120,
-        )
-        response.raise_for_status()
-        video_id = response.json()["id"]
-
-        # Poll for completion
-        import time
-
-        for _ in range(120):
-            status_resp = requests.get(
-                f"http://{server.host}:{server.port}/v1/videos/{video_id}",
-                timeout=30,
-            )
-            status_resp.raise_for_status()
-            status = status_resp.json()["status"]
-            if status == "completed":
-                break
-            if status in ("error", "failed"):
-                raise RuntimeError(f"Video generation failed: {status_resp.json()}")
-            time.sleep(5)
-        else:
-            raise TimeoutError(f"Video generation timed out after 600s (id={video_id})")
-
-        # Download video content
-        content_resp = requests.get(
-            f"http://{server.host}:{server.port}/v1/videos/{video_id}/content",
-            timeout=120,
-        )
-        content_resp.raise_for_status()
-        mp4_bytes = content_resp.content
-
-    # Save MP4
-    mp4_path = output_dir / "vllm_omni_pipeline.mp4"
-    with open(mp4_path, "wb") as f:
-        f.write(mp4_bytes)
-
-    # Extract frames
-    frames = _extract_mp4_frames(mp4_bytes)
-    for i, frame in enumerate(frames):
-        frame.save(output_dir / f"vllm_pipeline_frame_{i:04d}.png")
-    return frames
-
-
-@pytest.mark.full_model
-@pytest.mark.benchmark
-@pytest.mark.diffusion
-@pytest.mark.skipif(
-    not _DIFFUSERS_038, reason="LTX-2.3 requires diffusers >= 0.38.0 for cross_attn_mod and BWE vocoder"
-)
-@hardware_test(res={"cuda": "H100"}, num_cards=1)
-def test_ltx2_3_pipeline_matches_diffusers(accuracy_artifact_root: Path) -> None:
-    """Full-pipeline parity: vLLM-Omni serving stack vs diffusers.
-
-    Runs the complete vLLM-Omni OmniServer (subprocess, HTTP API, video
-    encoding) and compares per-frame against stock diffusers output.
-    Follows the Wan2.2 / Qwen Image pattern with seed-based determinism.
-    """
-    model = _model_name()
-    output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID)
-
-    diffusers_frames = _run_diffusers_baseline(model=model, output_dir=output_dir)
-    vllm_frames = _run_vllm_omni_serving(model=model, output_dir=output_dir)
-
-    _assert_video_similarity(
-        model_name=f"{MODEL_ID} (full-pipeline)",
-        vllm_frames=vllm_frames,
-        diffusers_frames=diffusers_frames,
-        ssim_threshold=PIPELINE_SSIM_THRESHOLD,
-        psnr_threshold=PIPELINE_PSNR_THRESHOLD,
-    )

tests/e2e/offline_inference/test_audex_s2s.py

@@ -0,0 +1,262 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E test for the Audex cascaded speech-to-speech pipeline.
+
+One ``audex_s2s`` deployment serves all three official
+passes. The routing contract is the hard gate here: text-modality passes
+(ASR, chat) finish at stage 0 and must never produce stage-1 audio; the
+audio-modality TTS pass streams through the causal speech decoder.
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+import re
+import time
+
+import numpy as np
+import pytest
+from vllm.assets.audio import AudioAsset
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import concat_audio
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+SAMPLE_RATE = 16_000
+
+ASR_PROMPT = (
+    "<|im_start|>user\n<so_embedding>\nTranscribe the input speech.<|im_end|>\n<|im_start|>assistant\n<think></think>"
+)
+
+# Official chat-pass recipe (examples/offline_inference/audex/speech_to_speech.py,
+# mirroring the official cascaded web demo) — 30B-A3B ONLY. The
+# formatting-instruction prefix is load-bearing on the 30B-A3B: without it a bare
+# transcript turn is treated as a TTS request and the model emits literal
+# <speechcodec_N> tokens instead of text, and under the yaml's near-greedy
+# defaults its reasoning block never closes. The 2B must NOT use this recipe:
+# it is long-verified stable on a bare prompt + engine-default sampling, while
+# the official temp-1.0 seed-0 draw hallucinates off-language text (yet closes
+# </think> cleanly, so the seed retry never fires) and fails the round-trip
+# WER gate — measured, reproduced byte-identically.
+SYSTEM_PROMPT = "You are a helpful and harmless assistant.\n\nYou are not allowed to use any tools."
+TEXT_INSTRUCTION = (
+    "SYSTEM & FORMATTING INSTRUCTION: You are Audex, created by NVIDIA based on the Nemotron-Cascade-2 architecture. "
+    "You may output your reasoning, followed by your final response. "
+    "You may format your reasoning block however you like. "
+    "However, your final response must strictly follow these rules: "
+    "* Write in plain, unformatted prose like a book or newspaper article. "
+    "* Do not use markdown, bullet points, lists, or headers. "
+    "* Standard numbers, abbreviations, and symbols are acceptable. "
+    "* [CRITICAL] You must press enter after every single sentence, placing each sentence on its own separate line."
+)
+CHAT_MAX_TOKENS = 2048  # official --text-max-tokens default
+# The official demo samples this pass unseeded; some trajectories deliberate
+# past the token budget without ever closing the think block, so retry a few
+# deterministic seeds until one closes.
+CHAT_SEEDS = (0, 1, 2)
+
+_audex_deployment = get_deploy_config_path("audex_s2s.yaml")
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": True}),
+        id="s2s_full",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_s2s_30b.yaml"), {"async_chunk": True}),
+            id="s2s_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+def _normalize(text: str) -> list[str]:
+    return "".join(c.lower() if c.isalnum() or c.isspace() else " " for c in text).split()
+
+
+def _wer(hyp: str, ref: str) -> float:
+    hyp_words, ref_words = _normalize(hyp), _normalize(ref)
+    if not ref_words:
+        return 0.0
+    rows = list(range(len(hyp_words) + 1))
+    for i, ref_word in enumerate(ref_words, 1):
+        prev, rows[0] = rows[0], i
+        for j, hyp_word in enumerate(hyp_words, 1):
+            cur = min(rows[j] + 1, rows[j - 1] + 1, prev + (ref_word != hyp_word))
+            prev, rows[j] = rows[j], cur
+    return rows[-1] / len(ref_words)
+
+
+def _chat_prompt(user_text: str) -> str:
+    # Official chat-pass recipe: system prompt + formatting instruction
+    # prefixed to the user text, thinking ENABLED via an open <think> block
+    # (chat_template.jinja with enable_thinking=True renders
+    # "<|im_start|>assistant\n<think>\n"). Priming the block open keeps the
+    # reasoning inside <think>...</think> so _extract_final_response can drop
+    # it; with a closed block the 30B leaks the reasoning as plain text.
+    return (
+        f"<|im_start|>system\n{SYSTEM_PROMPT}<|im_end|>\n"
+        f"<|im_start|>user\n{TEXT_INSTRUCTION}\n\n{user_text}<|im_end|>\n"
+        "<|im_start|>assistant\n<think>\n"
+    )
+
+
+def _extract_final_response(text: str) -> str:
+    # Mirror the official demo's extract_spoken_answer: keep only what
+    # follows the LAST closing think tag, falling back to the full text when
+    # the model skipped reasoning or never closed the block. Benign no-op for
+    # the 2B, whose answers carry no think block.
+    if "</think>" in text:
+        answer = text.rsplit("</think>", 1)[-1].strip()
+        if answer:
+            return answer
+    return text.strip()
+
+
+def _chat_params(runner: OmniRunner, seed: int):
+    # Official demo text-pass sampling (--text-temperature 1.0, --text-top-p
+    # 0.95, --text-max-tokens 2048). The yaml stage-0 defaults (temperature
+    # 0.1, top_k 80) are TTS-thinker settings; under near-greedy decoding the
+    # 30B loops inside its reasoning block until the token budget runs out
+    # and never emits </think>.
+    params = copy.deepcopy(runner.omni.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    stage0.temperature = 1.0
+    stage0.top_p = 0.95
+    stage0.top_k = -1
+    stage0.max_tokens = CHAT_MAX_TOKENS
+    stage0.seed = seed
+    return params
+
+
+def _chat_once(runner: OmniRunner, prompt: str, seed: int) -> tuple[str, bool]:
+    (req_output,) = runner.omni.generate([{"prompt": prompt, "modalities": ["text"]}], _chat_params(runner, seed))
+    out = req_output.outputs[0]
+    text = (out.text or "").strip()
+    truncated = getattr(out, "finish_reason", None) == "length" or len(out.token_ids or []) >= CHAT_MAX_TOKENS
+    return text, truncated
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+def test_audex_s2s_cascade_round_trip(omni_runner: OmniRunner, run_level: str) -> None:
+    """ASR → chat → TTS over one deployment; text passes never emit audio."""
+    real_weights = run_level in {"advanced_model", "full_model"}
+    # Spoken question: vLLM's public audio asset (downloaded on first use).
+    audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    # Pass 1 — ASR (text-final).
+    asr_outputs = omni_runner.omni.generate(
+        [
+            {
+                "prompt": ASR_PROMPT,
+                "multi_modal_data": {"audio": (audio, sr)},
+                "modalities": ["text"],
+            }
+        ]
+    )
+    assert len(asr_outputs) == 1
+    asr_out = asr_outputs[0]
+    transcript = (asr_out.outputs[0].text or "").strip()
+    mm = getattr(asr_out, "multimodal_output", None) or {}
+    stage1_audio = concat_audio(mm.get("audio")) if "audio" in mm else np.zeros((0,), dtype=np.float32)
+    assert stage1_audio.size == 0, "Text-modality ASR pass produced stage-1 audio (routing regression)"
+
+    # Pass 2 — chat (text-final). The recipe is size-conditional (see the
+    # SYSTEM_PROMPT comment above): the 30B requires the official recipe,
+    # the 2B regresses under it.
+    if omni_runner.model_name == _AUDEX_30B_MODEL:
+        # Official recipe. The raw completion is
+        # <reasoning...></think><answer>; speak only the final response.
+        # Retry with the next seed only when the budget ran out before the
+        # reasoning block closed (a natural stop without </think> is a
+        # complete answer); if every seed spins, keep the full text
+        # (official demo behavior).
+        chat_prompt = _chat_prompt(transcript or "Hello.")
+        raw = ""
+        for seed in CHAT_SEEDS:
+            raw, truncated = _chat_once(omni_runner, chat_prompt, seed)
+            if "</think>" in raw or not truncated:
+                break
+        answer = _extract_final_response(raw)
+    else:
+        # 2B: bare prompt with a closed think block, engine-default sampling
+        # (yaml near-greedy), no retry; the answer is used as-is.
+        chat_prompt = f"<|im_start|>user\n{transcript or 'Hello.'}<|im_end|>\n<|im_start|>assistant\n<think></think>"
+        (chat_output,) = omni_runner.omni.generate([{"prompt": chat_prompt, "modalities": ["text"]}])
+        answer = (chat_output.outputs[0].text or "").strip()
+
+    if real_weights:
+        assert transcript, "ASR pass produced an empty transcript"
+        assert answer, "Chat pass produced an empty answer"
+
+    # Pass 3 — TTS (audio-final; streams through code2wav). Both final
+    # stages may emit an output record; collect the audio-bearing one.
+    # The formatting instruction puts every sentence on its own line;
+    # flatten to plain prose for the TTS thinker.
+    tts_text = re.sub(r"\s+", " ", answer).strip() if answer else "The weather is nice today."
+    tts_text = tts_text[:200]
+    tts_prompt = {"prompt": build_cond_prompt(tts_text), "modalities": ["audio"]}
+    tts_outputs = omni_runner.omni.generate([tts_prompt])
+    assert tts_outputs, "TTS pass returned no outputs"
+    speech = np.zeros((0,), dtype=np.float32)
+    for req_output in tts_outputs:
+        mm = getattr(req_output, "multimodal_output", None) or {}
+        if "audio" in mm:
+            speech = concat_audio(mm["audio"])
+            break
+    assert speech.size > 0, "TTS pass produced empty audio"
+
+    if real_weights:
+        rms = float(np.sqrt(np.mean(np.square(speech))))
+        assert rms > 1e-3, f"TTS-pass audio near-silent (rms={rms})"
+
+        # Pass 4 — output verification: transcribe the synthesized answer
+        # with the SAME deployment's ASR pass and compare to the answer
+        # text (self-contained; no external ASR dependency).
+        verify_outputs = omni_runner.omni.generate(
+            [
+                {
+                    "prompt": ASR_PROMPT,
+                    "multi_modal_data": {"audio": (speech, SAMPLE_RATE)},
+                    "modalities": ["text"],
+                }
+            ]
+        )
+        verify_text = (verify_outputs[0].outputs[0].text or "").strip()
+        assert verify_text, "Output-verification ASR pass returned empty text"
+        # The full checkpoint prefixes transcripts with a language-tag
+        # sentence and quotes the content; extract the quoted payload.
+        if "'" in verify_text and len(verify_text.split("'")) >= 3:
+            verify_text = verify_text.split("'")[1].strip()
+        wer = _wer(verify_text, tts_text)
+        # Hard gate: the synthesized WAV must transcribe back to the answer
+        # (lenient threshold absorbs ASR quoting/prefix differences).
+        assert wer <= 0.5, f"Output WAV does not match the answer (WER={wer:.2f}): {verify_text!r} vs {tts_text!r}"
+        print(f"[trend] Audex S2S output-ASR WER vs answer: {wer:.3f} ({verify_text!r})")
+
+    # Final pass — time-to-first-audio for the streaming TTS path. This
+    # MUST be the last engine call: the py_generator API closes the engine
+    # when the generator is exhausted.
+    t_start = time.perf_counter()
+    first_audio_s: float | None = None
+    for req_output in omni_runner.omni.generate([tts_prompt], py_generator=True):
+        mm = getattr(req_output, "multimodal_output", None) or {}
+        if "audio" in mm and concat_audio(mm["audio"]).size > 0:
+            if first_audio_s is None:
+                first_audio_s = time.perf_counter() - t_start
+    assert first_audio_s is not None, "streaming TTS never produced an audio-bearing chunk"
+    print(f"[trend] Audex S2S TTS-pass first-audio latency: {first_audio_s * 1000.0:.0f} ms")

tests/e2e/offline_inference/test_audex_s2t.py

@@ -0,0 +1,106 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E smoke test for Audex audio understanding (thinker-only).
+
+Single-stage ``audex_thinker_only`` pipeline on
+``checkpoint_folder_full``: WAV (+ instruction) in, text out. The input clip
+is vLLM's public ``mary_had_lamb`` audio asset (downloaded on first use), so
+its transcript is known.
+
+WER is a trend metric (recorded in the assertion message, no numeric hard
+gate); the hard gates are structural: non-empty, non-degenerate text.
+"""
+
+from __future__ import annotations
+
+import os
+
+import pytest
+from vllm.assets.audio import AudioAsset
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+
+# Full transcript of the mary_had_lamb asset (Edison's phonograph recording);
+# a shortened reference would count the extra spoken words as insertions and
+# inflate WER past the trend threshold.
+_REFERENCE = (
+    "the first words i spoke in the original phonograph a little piece of practical poetry "
+    "mary had a little lamb its fleece was white as snow and everywhere that mary went the "
+    "lamb was sure to go"
+)
+
+ASR_PROMPT = (
+    "<|im_start|>user\n<so_embedding>\nTranscribe the input speech.<|im_end|>\n<|im_start|>assistant\n<think></think>"
+)
+
+_audex_deployment = get_deploy_config_path("audex_thinker_only.yaml")
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": False}),
+        id="thinker_only",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_thinker_only_30b.yaml"), {"async_chunk": False}),
+            id="thinker_only_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+def _normalize(text: str) -> list[str]:
+    return "".join(c.lower() if c.isalnum() or c.isspace() else " " for c in text).split()
+
+
+def _wer(hyp: str, ref: str) -> float:
+    hyp_words, ref_words = _normalize(hyp), _normalize(ref)
+    if not ref_words:
+        return 0.0
+    rows = list(range(len(hyp_words) + 1))
+    for i, ref_word in enumerate(ref_words, 1):
+        prev, rows[0] = rows[0], i
+        for j, hyp_word in enumerate(hyp_words, 1):
+            cur = min(
+                rows[j] + 1,
+                rows[j - 1] + 1,
+                prev + (ref_word != hyp_word),
+            )
+            prev, rows[j] = rows[j], cur
+    return rows[-1] / len(ref_words)
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_asr_smoke(omni_runner: OmniRunner, run_level: str) -> None:
+    """Transcribing the public audio asset yields coherent, non-degenerate text."""
+    audio, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
+    outputs = omni_runner.omni.generate([{"prompt": ASR_PROMPT, "multi_modal_data": {"audio": (audio, sr)}}])
+
+    assert len(outputs) == 1
+    text = outputs[0].outputs[0].text if outputs[0].outputs else ""
+    assert isinstance(text, str)
+
+    if run_level in {"advanced_model", "full_model"}:
+        words = _normalize(text)
+        assert words, f"Empty transcription: {text!r}"
+        # Degenerate repetition guard: no word may dominate the output.
+        most_common = max(words.count(w) for w in set(words))
+        assert most_common <= max(3, len(words) // 2), f"Degenerate transcript: {text!r}"
+        wer = _wer(text, _REFERENCE)
+        # Trend metric (recorded, not gated): substantially-correct check only.
+        assert wer <= 0.9, f"Transcript unrelated to reference (WER={wer:.2f}): {text!r}"
+        print(f"[trend] Audex S2T WER vs reference: {wer:.3f} ({text!r})")

tests/e2e/offline_inference/test_audex_tta.py

@@ -0,0 +1,145 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E smoke test for Audex text-to-audio (TTA).
+
+Caption → RVQ-phase-masked <audiocodec_*> tokens → XCodec1 waveform through
+the ``audex_tta`` pipeline. RVQ phase validity is enforced by
+the mask logits processor (unit-tested against the official validator); the
+hard gates here are decode-side: non-silent, finite audio at the documented
+sample rate. CFG runs at the official TTA default (3.0).
+"""
+
+from __future__ import annotations
+
+import copy
+import os
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import concat_audio
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+SAMPLE_RATE = 16_000
+
+# The acceptance gate requires eight captions decoding to non-silent,
+# finite audio through XCodec1.
+CAPTIONS = (
+    "Heavy rain falling on a tin roof.",
+    "A dog barking in the distance while birds chirp.",
+    "Ocean waves crashing on a rocky shore.",
+    "A crackling campfire at night with crickets.",
+    "Thunder rumbling during a storm.",
+    "A steam train whistle blowing as it departs.",
+    "Wind chimes ringing in a gentle breeze.",
+    "Footsteps crunching through fresh snow.",
+)
+
+_audex_deployment = get_deploy_config_path("audex_tta.yaml")
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": False}),
+        id="tta_full_payload",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_tta_30b.yaml"), {"async_chunk": False}),
+            id="tta_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+@pytest.fixture(scope="module")
+def tta_tokenizer_bits():
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+    from vllm_omni.model_executor.models.audex.tta import build_tta_phase_token_ids
+
+    root = ensure_audex_snapshot(_audex_model, profile="tta")
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, "checkpoint_folder_audiogen"), trust_remote_code=True)
+    phase_token_ids, start_tid, end_tid = build_tta_phase_token_ids(tokenizer)
+    return tokenizer, phase_token_ids, start_tid, end_tid
+
+
+def _tta_sampling_params(runner: OmniRunner, cond_prompt: str, tokenizer_bits, cfg_scale: float = 3.0):
+    from vllm_omni.model_executor.models.audex.prompt import build_tta_null_prompt
+
+    tokenizer, phase_token_ids, start_tid, end_tid = tokenizer_bits
+    params = copy.deepcopy(runner.omni.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    if stage0.extra_args is None:
+        stage0.extra_args = {}
+    stage0.extra_args["tta_rvq"] = {
+        "phase_token_ids": phase_token_ids,
+        "start_tid": start_tid,
+        "end_tid": end_tid,
+        "codec_cap": 4000,
+        "start_in_prompt": True,
+    }
+    if cfg_scale > 1.0:
+        stage0.extra_args.update(
+            {
+                "cfg_scale": float(cfg_scale),
+                "cfg_role": "cond",
+                "cfg_pair_id": f"e2e-tta-{abs(hash(cond_prompt)) % 10_000_000}",
+                "cfg_null_prompt": build_tta_null_prompt(cond_prompt, tokenizer),
+            }
+        )
+    return params
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_tta_eight_captions(omni_runner: OmniRunner, run_level: str, tta_tokenizer_bits) -> None:
+    """All eight captions decode through XCodec1 to non-silent, finite audio.
+
+    RVQ phase validity is enforced by the mask logits processor by
+    construction; the final user output carries only decoded audio (the
+    audio-final pipeline emits no stage-0 token record), and disabling the
+    mask in-engine as an e2e control would require a second engine boot
+    with a different logits-processor config. The mask's sensitivity is
+    therefore proven at unit level (see test_audex_tta.py's mask tests and
+    the unmasked-sampling control) against the official validator.
+    """
+    real_weights = run_level in {"advanced_model", "full_model"}
+    for caption in CAPTIONS:
+        prompt = build_tta_cond_prompt(caption)
+        params = _tta_sampling_params(omni_runner, prompt, tta_tokenizer_bits)
+        outputs = omni_runner.omni.generate([prompt], params)
+
+        assert len(outputs) == 1, f"{caption!r}: expected one output, got {len(outputs)}"
+        audio_mm = outputs[0].multimodal_output
+        assert "audio" in audio_mm, f"{caption!r}: no audio output: {list(audio_mm.keys())}"
+        audio = concat_audio(audio_mm["audio"])
+        assert audio.size > 0, f"{caption!r}: generated audio is empty"
+        assert np.isfinite(audio).all(), f"{caption!r}: audio contains NaN/Inf"
+
+        sr_val = audio_mm.get("sr", SAMPLE_RATE)
+        if isinstance(sr_val, list) and sr_val:
+            sr_val = sr_val[-1]
+        if hasattr(sr_val, "item"):
+            sr_val = sr_val.item()
+        assert int(sr_val) == SAMPLE_RATE, f"{caption!r}: unexpected sample_rate={sr_val}"
+
+        if real_weights:
+            duration_s = audio.size / SAMPLE_RATE
+            # The decode cap is 500 frames = 10 s at 50 fps.
+            assert 0.1 <= duration_s <= 10.5, f"{caption!r}: unexpected duration={duration_s:.3f}s"
+            rms = float(np.sqrt(np.mean(np.square(audio))))
+            assert rms > 1e-4, f"{caption!r}: audio is near-silent (rms={rms})"

tests/e2e/offline_inference/test_audex_tts.py

@@ -0,0 +1,161 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Offline E2E smoke test for Audex (Nemotron-Labs-Audex-2B) TTS.
+
+Passes the HF repo ROOT and verifies the 2-stage pipeline (thinker →
+streaming causal speech decoder): root-path/subfolder resolution, codec-token
+extraction, per-request streaming sessions, and the lookahead flush at EOS.
+"""
+
+from __future__ import annotations
+
+import os
+
+import numpy as np
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.media import concat_audio
+from tests.helpers.runtime import OmniRunner
+from tests.helpers.stage_config import get_deploy_config_path
+from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+MODEL_DIR_ENV = "VLLM_OMNI_AUDEX_MODEL_DIR"
+SAMPLE_RATE = 16_000
+
+SYNTH_TEXTS = (
+    "The weather is so good, and I want to enjoy the beautiful morning in the park.",
+    "The quick brown fox jumps over the lazy dog.",
+)
+
+
+_audex_deployment = get_deploy_config_path("audex_tts.yaml")
+# Collection must not download anything: pass the repo id (or a local
+# override) and let the engine's stage-init path resolve/download the
+# required snapshot subset at execution time (ensure_audex_snapshot).
+_audex_model = os.environ.get(MODEL_DIR_ENV) or MODEL
+_OMNI_RUNNER_PARAMS = [
+    pytest.param(
+        (_audex_model, _audex_deployment, {"async_chunk": True}),
+        id="async_chunk",
+    ),
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    _OMNI_RUNNER_PARAMS.append(
+        pytest.param(
+            (_AUDEX_30B_MODEL, get_deploy_config_path("audex_tts_30b.yaml"), {"async_chunk": True}),
+            id="tts_30b",
+        )
+    )
+pytestmark = [
+    pytest.mark.slow,
+    pytest.mark.tts,
+    pytest.mark.parametrize("omni_runner", _OMNI_RUNNER_PARAMS, indirect=True),
+]
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_tts_smoke(omni_runner: OmniRunner, run_level: str) -> None:
+    """Audex TTS from the repo root should produce sane 16 kHz audio per prompt.
+
+    At ``core_model`` level the runner loads DUMMY weights (structural smoke
+    only); real-speech assertions apply from ``advanced_model`` up.
+    """
+    prompts = [build_cond_prompt(text) for text in SYNTH_TEXTS]
+    outputs = omni_runner.omni.generate(prompts)
+
+    real_weights = run_level in {"advanced_model", "full_model"}
+    assert len(outputs) == len(SYNTH_TEXTS), f"expected {len(SYNTH_TEXTS)} outputs, got {len(outputs)}"
+    for output in outputs:
+        audio_mm = output.multimodal_output
+        assert "audio" in audio_mm, f"No audio output found: {list(audio_mm.keys())}"
+        audio = concat_audio(audio_mm["audio"])
+        assert audio.size > 0, "Generated audio is empty"
+
+        sr_val = audio_mm.get("sr", SAMPLE_RATE)
+        if isinstance(sr_val, list) and sr_val:
+            sr_val = sr_val[-1]
+        if hasattr(sr_val, "item"):
+            sr_val = sr_val.item()
+        assert int(sr_val) == SAMPLE_RATE, f"Unexpected sample_rate={sr_val}"
+
+        if real_weights:
+            duration_s = audio.size / SAMPLE_RATE
+            assert 0.5 <= duration_s <= 20.0, f"Unexpected duration={duration_s:.3f}s"
+            # Real speech, not silence.
+            rms = float(np.sqrt(np.mean(np.square(audio))))
+            assert rms > 1e-3, f"Audio is near-silent (rms={rms})"
+
+
+def _cfg_sampling_params(runner: OmniRunner, cfg_scale: float, pair_id: str, cond_prompt: str):
+    """Stage sampling params carrying the CFG pair contract for one request."""
+    import copy
+
+    from transformers import AutoTokenizer
+
+    from vllm_omni.model_executor.models.audex.checkpoint import ensure_audex_snapshot
+    from vllm_omni.model_executor.models.audex.prompt import build_null_prompt
+
+    # Resolve the ACTIVE param's model (2B or 30B), not the module default:
+    # the null prompt must be built with the same tokenizer the engine uses.
+    root = ensure_audex_snapshot(runner.model_name)
+    tokenizer = AutoTokenizer.from_pretrained(os.path.join(root, "checkpoint_folder_audiogen"), trust_remote_code=True)
+    params = copy.deepcopy(runner.omni.resolve_sampling_params_list(None))
+    stage0 = params[0]
+    if stage0.extra_args is None:
+        stage0.extra_args = {}
+    stage0.extra_args.update(
+        {
+            "cfg_scale": float(cfg_scale),
+            "cfg_role": "cond",
+            "cfg_pair_id": pair_id,
+            "cfg_null_prompt": build_null_prompt(cond_prompt, tokenizer),
+        }
+    )
+    return params
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_cfg_guided_single_stream(omni_runner: OmniRunner, run_level: str) -> None:
+    """A guided request must yield exactly ONE audio stream (companion suppressed)."""
+    prompt = build_cond_prompt(SYNTH_TEXTS[0])
+    params = _cfg_sampling_params(omni_runner, 1.5, "e2e-cfg-pair", prompt)
+    outputs = omni_runner.omni.generate([prompt], params)
+
+    assert len(outputs) == 1, f"expected exactly one output (no companion leak), got {len(outputs)}"
+    audio = concat_audio(outputs[0].multimodal_output["audio"])
+    assert audio.size > 0, "Guided generation produced empty audio"
+    if run_level in {"advanced_model", "full_model"}:
+        rms = float(np.sqrt(np.mean(np.square(audio))))
+        assert rms > 1e-3, f"Guided audio is near-silent (rms={rms})"
+        duration_s = audio.size / SAMPLE_RATE
+        assert 0.5 <= duration_s <= 20.0, f"Unexpected duration={duration_s:.3f}s"
+
+
+@pytest.mark.advanced_model
+@hardware_test(res={"cuda": "L4"}, num_cards=1)
+def test_audex_offline_cfg_disabled_is_byte_identical(omni_runner: OmniRunner, run_level: str) -> None:
+    """cfg_scale=1.0 must be indistinguishable from not passing cfg at all."""
+    prompt = build_cond_prompt(SYNTH_TEXTS[1])
+
+    baseline = omni_runner.omni.generate([prompt])
+    base_audio = concat_audio(baseline[0].multimodal_output["audio"])
+
+    import copy
+
+    params = copy.deepcopy(omni_runner.omni.resolve_sampling_params_list(None))
+    if params[0].extra_args is None:
+        params[0].extra_args = {}
+    params[0].extra_args["cfg_scale"] = 1.0
+    disabled = omni_runner.omni.generate([prompt], params)
+    disabled_audio = concat_audio(disabled[0].multimodal_output["audio"])
+
+    assert base_audio.shape == disabled_audio.shape, (
+        f"cfg_scale=1.0 changed output length: {base_audio.shape} vs {disabled_audio.shape}"
+    )
+    np.testing.assert_array_equal(base_audio, disabled_audio)

tests/e2e/online_serving/test_audex_tta.py

@@ -0,0 +1,84 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online tests for Audex (Nemotron-Labs-Audex-2B) text-to-audio.
+
+Covers /v1/audio/speech through the ``audex_tta`` deployment:
+the ``audex_tta`` adapter builds the caption prompt and the serving layer
+injects the RVQ phase contract plus the official default guidance
+(cfg_scale 3.0) server-side.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+CAPTION = "Heavy rain falling on a tin roof."
+SAMPLE_RATE = 16_000
+
+audex_tta_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("audex_tta.yaml"),
+            server_args=["--trust-remote-code"],
+        ),
+        id="audex_tta",
+    )
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    audex_tta_server_params.append(
+        pytest.param(
+            OmniServerParams(
+                model=_AUDEX_30B_MODEL,
+                stage_config_path=get_deploy_config_path("audex_tta_30b.yaml"),
+                server_args=["--trust-remote-code"],
+            ),
+            id="online_tta_30b",
+        )
+    )
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_tta_server_params, indirect=True)
+def test_audex_tta_online_smoke(omni_server, openai_client, run_level) -> None:
+    """A caption returns a valid finite non-silent WAV at the default cfg 3.0."""
+    import io
+    import wave
+
+    import numpy as np
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    resp = requests.post(
+        url,
+        json={"model": omni_server.model, "input": CAPTION, "response_format": "wav"},
+        timeout=600,
+    )
+    assert resp.status_code == 200, resp.text[:200]
+    assert resp.content[:4] == b"RIFF", "not a WAV response"
+
+    with wave.open(io.BytesIO(resp.content)) as wav:
+        assert wav.getframerate() == SAMPLE_RATE, f"unexpected sample rate {wav.getframerate()}"
+        frames = wav.readframes(wav.getnframes())
+        duration_s = wav.getnframes() / wav.getframerate()
+
+    audio = np.frombuffer(frames, dtype=np.int16).astype(np.float32) / 32768.0
+    assert audio.size > 0, "empty audio body"
+    assert np.isfinite(audio).all(), "audio contains NaN/Inf"
+
+    if run_level in {"advanced_model", "full_model"}:
+        # The decode cap is 500 frames = 10 s at 50 fps.
+        assert 0.1 <= duration_s <= 10.5, f"unexpected duration={duration_s:.3f}s"
+        rms = float(np.sqrt(np.mean(np.square(audio))))
+        assert rms > 1e-4, f"audio is near-silent (rms={rms})"

tests/e2e/online_serving/test_audex_tts.py

@@ -0,0 +1,230 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""E2E online tests for Audex (Nemotron-Labs-Audex-2B) TTS.
+
+Covers /v1/audio/speech non-streaming and audio-byte streaming, plus the
+Audex request policies: single built-in voice (missing/empty/"default" only),
+no reference audio, and cfg_scale reserved at 1.0.
+"""
+
+import os
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+import pytest
+
+from tests.helpers.mark import hardware_test
+from tests.helpers.runtime import OmniServerParams
+from tests.helpers.stage_config import get_deploy_config_path
+
+pytestmark = [pytest.mark.slow, pytest.mark.tts]
+
+MODEL = "nvidia/Nemotron-Labs-Audex-2B"
+PROMPT = "The weather is so good, and I want to enjoy the beautiful morning in the park."
+
+audex_server_params = [
+    pytest.param(
+        OmniServerParams(
+            model=MODEL,
+            stage_config_path=get_deploy_config_path("audex_tts.yaml"),
+            server_args=["--trust-remote-code"],
+        ),
+        id="audex",
+    )
+]
+# 30B-A3B variant: opt-in (pulls ~60 GB); enable with AUDEX_E2E_30B=1.
+_AUDEX_30B_MODEL = os.environ.get("VLLM_OMNI_AUDEX_30B_MODEL_DIR") or "nvidia/Nemotron-Labs-Audex-30B-A3B"
+if os.environ.get("AUDEX_E2E_30B") == "1":
+    audex_server_params.append(
+        pytest.param(
+            OmniServerParams(
+                model=_AUDEX_30B_MODEL,
+                stage_config_path=get_deploy_config_path("audex_tts_30b.yaml"),
+                server_args=["--trust-remote-code"],
+            ),
+            id="online_tts_30b",
+        )
+    )
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_nonstream(omni_server, openai_client) -> None:
+    """Plain English TTS, non-streaming WAV response."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": False,
+        "response_format": "wav",
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_stream_audio(omni_server, openai_client) -> None:
+    """Plain English TTS, raw audio-byte streaming."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": True,
+        "stream_format": "audio",
+        "response_format": "wav",
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_max_new_tokens_caps_duration(omni_server, openai_client, run_level) -> None:
+    """max_new_tokens must cap stage-0 generation (review P2 regression)."""
+    import io
+    import wave
+
+    import requests
+
+    if run_level not in {"advanced_model", "full_model"}:
+        pytest.skip("duration cap needs real weights")
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    resp = requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": PROMPT,
+            "response_format": "wav",
+            "max_new_tokens": 25,
+        },
+    )
+    assert resp.status_code == 200, resp.text
+    with wave.open(io.BytesIO(resp.content)) as wav:
+        duration_s = wav.getnframes() / wav.getframerate()
+    # 25 codec frames at 50 fps = 0.5 s (+ small lookahead tail); the full
+    # sentence would be ~3 s, so a capped result proves the override applied.
+    assert duration_s <= 1.0, f"max_new_tokens ignored: duration={duration_s:.2f}s"
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_request_policies(omni_server, openai_client) -> None:
+    """Unsupported voice / CFG / empty input must be rejected with 400."""
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+
+    ok = requests.post(url, json={"model": omni_server.model, "input": PROMPT, "voice": "default"})
+    assert ok.status_code == 200, ok.text
+
+    for bad_payload, needle in (
+        ({"input": PROMPT, "voice": "alloy"}, "voice"),
+        ({"input": "   "}, "non-empty"),
+        ({"input": PROMPT, "extra_params": {"cfg_scale": 0.5}}, "cfg_scale"),
+        ({"input": PROMPT, "extra_params": {"cfg_scale": 10.5}}, "cfg_scale"),
+        ({"input": PROMPT, "extra_params": {"cfg_scale": "big"}}, "cfg_scale"),
+        ({"input": PROMPT, "extra_params": {"cfg_role": "cond"}}, "cfg_role"),
+        ({"input": PROMPT, "extra_params": {"cfg_pair_id": "x"}}, "cfg_pair_id"),
+        ({"input": PROMPT, "ref_audio": "https://example.com/ref.wav"}, "reference audio"),
+    ):
+        resp = requests.post(url, json={"model": omni_server.model, **bad_payload})
+        assert resp.status_code == 400, f"{bad_payload} -> {resp.status_code}: {resp.text[:200]}"
+        assert needle in resp.text, f"{bad_payload} error missing {needle!r}: {resp.text[:200]}"
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_cfg_nonstream(omni_server, openai_client) -> None:
+    """Guided (cfg_scale=1.5) non-streaming request returns one WAV."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": False,
+        "response_format": "wav",
+        "extra_params": {"cfg_scale": 1.5},
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_cfg_stream_audio(omni_server, openai_client) -> None:
+    """Guided (cfg_scale=1.5) streaming request yields a single audio stream."""
+    request_config = {
+        "model": omni_server.model,
+        "input": PROMPT,
+        "stream": True,
+        "stream_format": "audio",
+        "response_format": "wav",
+        "extra_params": {"cfg_scale": 1.5},
+    }
+    openai_client.send_audio_speech_request(request_config)
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_mixed_cfg_concurrency(omni_server, openai_client) -> None:
+    """Mixed guided/unguided requests at concurrency 4 all complete with audio."""
+    import concurrent.futures
+
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    payloads = [
+        {"model": omni_server.model, "input": f"{PROMPT} Take {i}.", "response_format": "wav"}
+        | ({"extra_params": {"cfg_scale": 1.5}} if i % 2 == 0 else {})
+        for i in range(4)
+    ]
+
+    def _post(payload):
+        return requests.post(url, json=payload, timeout=300)
+
+    with concurrent.futures.ThreadPoolExecutor(max_workers=4) as pool:
+        responses = list(pool.map(_post, payloads))
+
+    for payload, resp in zip(payloads, responses):
+        assert resp.status_code == 200, f"{payload} -> {resp.status_code}: {resp.text[:200]}"
+        assert resp.content[:4] == b"RIFF", f"{payload}: not a WAV response"
+        assert len(resp.content) > 44, f"{payload}: empty audio body"
+
+
+@hardware_test(res={"cuda": "H100"}, num_cards=1)
+@pytest.mark.parametrize("omni_server", audex_server_params, indirect=True)
+def test_audex_tts_stream_abort_leaves_server_healthy(omni_server, openai_client) -> None:
+    """Aborting a guided streaming request mid-flight must not wedge the server.
+
+    The abort cascades to the CFG companion (cfg_companion_tracker); a
+    follow-up guided request on the same server must still succeed.
+    """
+    import requests
+
+    url = f"{openai_client.base_url.rstrip('/')}/v1/audio/speech"
+    long_input = " ".join([PROMPT] * 4)
+
+    with requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": long_input,
+            "stream": True,
+            "stream_format": "audio",
+            "response_format": "wav",
+            "extra_params": {"cfg_scale": 1.5},
+        },
+        stream=True,
+        timeout=300,
+    ) as resp:
+        assert resp.status_code == 200, resp.text
+        # Read one chunk, then abort the connection mid-stream.
+        next(resp.iter_content(chunk_size=4096))
+
+    follow_up = requests.post(
+        url,
+        json={
+            "model": omni_server.model,
+            "input": PROMPT,
+            "response_format": "wav",
+            "extra_params": {"cfg_scale": 1.5},
+        },
+        timeout=300,
+    )
+    assert follow_up.status_code == 200, follow_up.text[:200]
+    assert follow_up.content[:4] == b"RIFF" and len(follow_up.content) > 44

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 @@ HSDP_ARGS = ["--use-hsdp", "--hsdp-shard-size", "2"]
 def _cases():
     cases = []
 
-    # T2V: LTX2TwoStagesPipeline
     cases.append(
         pytest.param(
             OmniServerParams(
@@ -35,15 +33,15 @@ def _cases():
                 server_args=[
                     *HSDP_ARGS,
                     "--model-class-name",
-                    "LTX2TwoStagesPipeline",
+                    "LTX2DistilledPipeline",
                 ],
             ),
+            False,
             id="t2v_hsdp",
             marks=PARALLEL_MARKS,
         )
     )
 
-    # I2V: LTX2ImageToVideoTwoStagesPipeline
     cases.append(
         pytest.param(
             OmniServerParams(
@@ -51,9 +49,10 @@ def _cases():
                 server_args=[
                     *HSDP_ARGS,
                     "--model-class-name",
-                    "LTX2ImageToVideoTwoStagesPipeline",
+                    "LTX2DistilledPipeline",
                 ],
             ),
+            True,
             id="i2v_hsdp",
             marks=PARALLEL_MARKS,
         )
@@ -62,25 +61,22 @@ def _cases():
     return cases
 
 
-@pytest.mark.parametrize("omni_server", _cases(), indirect=True)
+@pytest.mark.parametrize(("omni_server", "is_i2v"), _cases(), indirect=["omni_server"])
 def test_ltx2_two_stage_hsdp(
     omni_server: OmniServer,
+    is_i2v: bool,
     openai_client: OpenAIClientHandler,
 ):
-    is_i2v = any("ImageToVideo" in arg for arg in omni_server.serve_args)
-
-    # The two-stage pipeline generates at the requested resolution then 2x
-    # upsamples via LTX2LatentUpsamplerModel, so output dimensions differ
-    # from the request.  Omit height/width from form_data so the assertion
-    # helper skips the dimension check (it only asserts when the key is
-    # present).  The pipeline falls back to its own defaults.
+    # Keep CI small while exercising the fixed 8+3-step distilled recipe.
+    # Height and width describe the final output; Stage 1 runs at half size.
     form_data = {
         "prompt": PROMPT,
         "model": omni_server.model,
+        "height": 128,
+        "width": 128,
         "num_frames": 9,
         "fps": 8,
-        "num_inference_steps": 2,
-        "guidance_scale": 1.0,
+        "num_inference_steps": 8,
         "seed": 42,
     }
 

tests/e2e/online_serving/test_qwen3_tts_customvoice.py

@@ -12,7 +12,6 @@ import os
 os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
 
 import pytest
-from vllm.platforms import current_platform
 
 from tests.helpers.mark import hardware_test
 from tests.helpers.runtime import OmniServerParams
@@ -21,6 +20,7 @@ from tests.helpers.stage_config import (
     get_deploy_config_stage,
     modify_stage_config,
 )
+from vllm_omni.platforms import current_omni_platform
 
 MODEL = "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"
 
@@ -71,7 +71,6 @@ default_tts_server_params = [
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.tts
-@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA Graph startup test requires CUDA")
 @hardware_test(res={"cuda": "L4"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", default_tts_server_params, indirect=True)
 def test_default_cuda_graph_startup(omni_server) -> None:
@@ -93,7 +92,11 @@ def test_default_cuda_graph_startup(omni_server) -> None:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(
+    current_omni_platform.is_npu(),
+    reason="#issue 5479",
+)
+@hardware_test(res={"cuda": "L4", "npu": "A3"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_001(omni_server, openai_client) -> None:
     """
@@ -119,7 +122,11 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
 @pytest.mark.core_model
 @pytest.mark.advanced_model
 @pytest.mark.tts
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@pytest.mark.skipif(
+    current_omni_platform.is_npu(),
+    reason="#issue 5479",
+)
+@hardware_test(res={"cuda": "L4", "npu": "A3"}, num_cards=1)
 @pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
 def test_text_to_audio_002(omni_server, openai_client) -> None:
     """

tests/entrypoints/openai_api/test_audex_serving_guards.py

@@ -0,0 +1,376 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Serving-layer regressions for the Audex zero-codec-token contract.
+
+A zero-codec-token (or phase-invalid TTA) request reaches the serving layer
+as a stream whose only audio payloads are zero-size tensors. The streaming
+generator must not turn that into a header-only/empty successful stream: the
+raw-bytes generator raises (before any WAV header), and the SSE wrapper
+surfaces ``speech.audio.error`` instead of ``speech.audio.done``. The
+contract covers both audex pipelines (``audex`` TTS and ``audex_tta``).
+"""
+
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from vllm_omni.entrypoints.openai.serving_speech import (
+    _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES,
+    OmniOpenAIServingSpeech,
+)
+
+
+def _serving(model_type: str = "audex") -> OmniOpenAIServingSpeech:
+    serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+    serving._tts_model_type = model_type
+    serving._mark_ref_audio_artifact_ready_for_request = lambda request_id: None
+    serving._discard_ref_audio_artifact_warmup = lambda request_id: None
+    return serving
+
+
+def _result(audio: torch.Tensor) -> SimpleNamespace:
+    return SimpleNamespace(multimodal_output={"model_outputs": audio, "sr": torch.tensor(16000, dtype=torch.int32)})
+
+
+async def _gen(*results):
+    for res in results:
+        yield res
+
+
+async def _collect(chunks):
+    out = []
+    async for chunk in chunks:
+        out.append(chunk)
+    return out
+
+
+def test_no_audio_guard_covers_both_audex_pipelines():
+    # The non-streaming CreateAudio guard keys off the same constant; pin
+    # its membership so neither pipeline regresses to empty-200 responses.
+    assert _AUDEX_NO_AUDIO_GUARD_MODEL_TYPES == {"audex", "audex_tta"}
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_type", ["audex", "audex_tta"])
+async def test_raw_stream_raises_on_empty_audio_before_any_header(model_type):
+    serving = _serving(model_type)
+    chunks = serving._generate_audio_chunks(
+        _gen(_result(torch.empty(0)), _result(torch.empty(0))),
+        request_id="req-empty",
+        response_format="wav",
+    )
+    with pytest.raises(ValueError, match="no audio"):
+        await _collect(chunks)
+
+
+@pytest.mark.asyncio
+async def test_raw_stream_yields_audio_normally_when_non_empty():
+    serving = _serving()
+    chunks = serving._generate_audio_chunks(
+        _gen(_result(torch.zeros(320)), _result(torch.empty(0))),
+        request_id="req-ok",
+        response_format="wav",
+    )
+    collected = await _collect(chunks)
+    # WAV header + one PCM chunk; the trailing empty chunk is skipped silently.
+    assert len(collected) == 2
+    assert collected[0][:4] == b"RIFF"
+    assert len(collected[1]) == 320 * 2  # int16 PCM
+
+
+@pytest.mark.asyncio
+@pytest.mark.parametrize("model_type", ["audex", "audex_tta"])
+async def test_sse_stream_emits_error_not_done_on_empty_audio(model_type):
+    serving = _serving(model_type)
+    events = serving._generate_audio_sse_events(
+        _gen(_result(torch.empty(0))),
+        request_id="req-empty",
+        response_format="pcm",
+    )
+    collected = await _collect(events)
+    assert any("speech.audio.error" in event for event in collected)
+    assert not any("speech.audio.done" in event for event in collected)
+
+
+@pytest.mark.asyncio
+async def test_non_audex_empty_chunks_keep_legacy_behavior():
+    """Other models keep their existing semantics (no new rejection)."""
+    serving = _serving(model_type="qwen3_tts")
+    chunks = serving._generate_audio_chunks(
+        _gen(_result(torch.empty(0))),
+        request_id="req-other",
+        response_format="pcm",
+    )
+    await _collect(chunks)  # must not raise
+
+
+# ---------------------------------------------------------------- CFG serving surface
+
+import re  # noqa: E402
+
+from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest  # noqa: E402
+from vllm_omni.entrypoints.openai.tts_adapters.audex import AudexAdapter  # noqa: E402
+
+
+def _adapter() -> AudexAdapter:
+    return AudexAdapter.__new__(AudexAdapter)
+
+
+def _speech_request(**extra_params) -> OpenAICreateSpeechRequest:
+    payload = {"model": "audex", "input": "Hello world."}
+    if extra_params:
+        payload["extra_params"] = extra_params
+    return OpenAICreateSpeechRequest(**payload)
+
+
+class _MarkerTokenizer:
+    """Special markers count one token; words and punctuation split."""
+
+    _TOKEN_RE = re.compile(r"<unk>|<\|[^|]*\|>|<[a-z_]+>|\w+|[^\w\s]")
+
+    def encode(self, text: str) -> list[int]:
+        return [0] * len(self._TOKEN_RE.findall(text))
+
+
+class TestAudexCfgValidation:
+    def test_no_cfg_scale_accepted(self):
+        assert _adapter().validate(_speech_request()) is None
+
+    @pytest.mark.parametrize("scale", [1.0, 1.5, 3.0, 10.0])
+    def test_in_range_scale_accepted(self, scale):
+        assert _adapter().validate(_speech_request(cfg_scale=scale)) is None
+
+    @pytest.mark.parametrize("scale", [0.5, 0.0, -1.0, 10.5, "big"])
+    def test_out_of_range_or_malformed_scale_rejected(self, scale):
+        error = _adapter().validate(_speech_request(cfg_scale=scale))
+        assert error is not None and "cfg_scale" in error
+
+    def test_explicit_null_scale_treated_as_absent(self):
+        assert _adapter().validate(_speech_request(cfg_scale=None)) is None
+        serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": None})
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": "p"}, params)
+        assert out.extra_args == {}
+
+    @pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt"])
+    def test_internal_cfg_keys_rejected(self, key):
+        error = _adapter().validate(_speech_request(**{key: "x"}))
+        assert error is not None and key in error
+
+
+class TestAudexCfgInjection:
+    def _serving_with_fake_tokenizer(self):
+        serving = _serving()
+        serving._get_audex_tokenizer = lambda: _MarkerTokenizer()
+        return serving
+
+    def _cond_prompt(self, text: str = "Hello world.") -> str:
+        from vllm_omni.model_executor.models.audex.prompt import build_cond_prompt
+
+        return build_cond_prompt(text)
+
+    def test_no_extra_args_is_untouched(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args=None)
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        assert out.extra_args is None
+        assert params.extra_args is None
+
+    def test_scale_one_leaves_params_byte_identical(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
+        out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": self._cond_prompt()}, params)
+        assert out.extra_args == {}
+
+    def test_guided_request_gains_pair_contract(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
+        cond = self._cond_prompt("Some transcription text here.")
+        out = serving._inject_audex_cfg_pair_args("req-42", {"prompt": cond}, params)
+
+        extra = out.extra_args
+        assert extra["cfg_scale"] == 1.5
+        assert extra["cfg_role"] == "cond"
+        assert extra["cfg_pair_id"] == "req-42"
+        null_prompt = extra["cfg_null_prompt"]
+        assert "<unk>" in null_prompt
+        assert "Some transcription" not in null_prompt
+        tok = _MarkerTokenizer()
+        assert len(tok.encode(null_prompt)) == len(tok.encode(cond))
+
+    def test_missing_text_prompt_raises(self):
+        serving = self._serving_with_fake_tokenizer()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.5})
+        with pytest.raises(ValueError, match="adapter-built text prompt"):
+            serving._inject_audex_cfg_pair_args("req-1", {"prompt_token_ids": [1]}, params)
+
+    def test_injection_never_mutates_input_params(self):
+        """The input may be the engine's SHARED default params (review P2):
+        pair state must land only on the returned clone."""
+        serving = self._serving_with_fake_tokenizer()
+        shared = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
+        out = serving._inject_audex_cfg_pair_args("req-7", {"prompt": self._cond_prompt()}, shared)
+        assert out is not shared
+        assert shared.extra_args == {"cfg_scale": 1.5}
+        assert shared.temperature == 0.1
+        assert out.extra_args["cfg_pair_id"] == "req-7"
+        assert out.temperature == 0.05
+
+
+class TestAudexStageDetection:
+    """audex_omni is speech-capable only alongside the speech decoder."""
+
+    def _serving_with_stages(self, stages: list[str]) -> OmniOpenAIServingSpeech:
+        serving = OmniOpenAIServingSpeech.__new__(OmniOpenAIServingSpeech)
+        serving.engine_client = SimpleNamespace(
+            stage_configs=[SimpleNamespace(engine_args=SimpleNamespace(model_stage=s)) for s in stages]
+        )
+        return serving
+
+    def test_full_pipeline_detected_as_tts(self):
+        serving = self._serving_with_stages(["audex_omni", "audex_code2wav"])
+        stage = serving._find_tts_stage()
+        assert stage is not None and stage.engine_args.model_stage == "audex_omni"
+
+    def test_thinker_only_not_detected_as_tts(self):
+        serving = self._serving_with_stages(["audex_omni"])
+        assert serving._find_tts_stage() is None
+
+    def test_tts_pipeline_still_detected(self):
+        serving = self._serving_with_stages(["audex_thinker", "audex_code2wav"])
+        stage = serving._find_tts_stage()
+        assert stage is not None and stage.engine_args.model_stage == "audex_thinker"
+
+
+def test_guided_injection_sets_guided_temperature():
+    """Guided requests decode at the measured CFG temperature (0.05)."""
+    serving = TestAudexCfgInjection()._serving_with_fake_tokenizer()
+    cond = TestAudexCfgInjection()._cond_prompt()
+    params = SimpleNamespace(extra_args={"cfg_scale": 1.5}, temperature=0.1)
+    out = serving._inject_audex_cfg_pair_args("req-1", {"prompt": cond}, params)
+    assert out.temperature == 0.05
+
+    unguided = SimpleNamespace(extra_args={"cfg_scale": 1.0}, temperature=0.1)
+    out = serving._inject_audex_cfg_pair_args("req-2", {"prompt": cond}, unguided)
+    assert out.temperature == 0.1
+
+
+# ---------------------------------------------------------------- TTA serving surface
+
+from vllm_omni.entrypoints.openai.tts_adapters import resolve_adapter  # noqa: E402
+from vllm_omni.entrypoints.openai.tts_adapters.audex_tta import AudexTTAAdapter  # noqa: E402
+
+
+def _tta_adapter() -> AudexTTAAdapter:
+    return AudexTTAAdapter.__new__(AudexTTAAdapter)
+
+
+class _TTAVocabTokenizer(_MarkerTokenizer):
+    """Marker tokenizer + a full audiocodec vocab for phase-id building."""
+
+    def get_vocab(self):
+        vocab = {"<audiogen_start>": 7, "<audiogen_end>": 8}
+        for n in range(8192):
+            vocab[f"<audiocodec_{n}>"] = 1000 + n
+        return vocab
+
+
+class TestAudexTTAAdapterValidation:
+    def test_default_caption_accepted(self):
+        assert _tta_adapter().validate(_speech_request()) is None
+
+    def test_empty_caption_rejected(self):
+        req = _speech_request()
+        req.input = "  "
+        assert _tta_adapter().validate(req) is not None
+
+    @pytest.mark.parametrize("scale", [1.0, 3.0, 10.0])
+    def test_in_range_scale_accepted(self, scale):
+        assert _tta_adapter().validate(_speech_request(cfg_scale=scale)) is None
+
+    @pytest.mark.parametrize("scale", [0.5, 10.5, "big"])
+    def test_out_of_range_scale_rejected(self, scale):
+        error = _tta_adapter().validate(_speech_request(cfg_scale=scale))
+        assert error is not None and "cfg_scale" in error
+
+    @pytest.mark.parametrize("key", ["cfg_role", "cfg_pair_id", "cfg_null_prompt", "tta_rvq"])
+    def test_internal_keys_rejected(self, key):
+        error = _tta_adapter().validate(_speech_request(**{key: "x"}))
+        assert error is not None and key in error
+
+
+class TestAdapterRouting:
+    def test_tta_adapter_resolves_by_model_type(self):
+        adapter_cls = resolve_adapter("audex_tta")
+        assert adapter_cls is AudexTTAAdapter
+        assert adapter_cls.stage_keys == frozenset({"audex_tta_thinker"})
+
+    def test_tts_adapter_never_selects_tta_stage(self):
+        assert "audex_tta_thinker" not in AudexAdapter.stage_keys
+        assert not (AudexAdapter.stage_keys & AudexTTAAdapter.stage_keys)
+
+    def test_detection_separates_tts_and_tta(self):
+        tta_serving = TestAudexStageDetection()._serving_with_stages(["audex_tta_thinker", "audex_xcodec"])
+        tta_serving._tts_stage = tta_serving._find_tts_stage()
+        assert tta_serving._tts_stage.engine_args.model_stage == "audex_tta_thinker"
+        assert tta_serving._detect_tts_model_type() == "audex_tta"
+
+        tts_serving = TestAudexStageDetection()._serving_with_stages(["audex_thinker", "audex_code2wav"])
+        tts_serving._tts_stage = tts_serving._find_tts_stage()
+        assert tts_serving._detect_tts_model_type() == "audex"
+
+
+class TestAudexTTAInjection:
+    def _serving(self):
+        serving = _serving(model_type="audex_tta")
+        serving._get_audex_tokenizer = lambda: _TTAVocabTokenizer()
+        serving._audex_tta_rvq = None
+        return serving
+
+    def _prompt(self):
+        from vllm_omni.model_executor.models.audex.prompt import build_tta_cond_prompt
+
+        return {"prompt": build_tta_cond_prompt("Rain on a tin roof.")}
+
+    def test_default_injection_applies_rvq_and_cfg3(self):
+        serving = self._serving()
+        params = SimpleNamespace(extra_args=None)
+        out = serving._inject_audex_tta_args("req-9", self._prompt(), params)
+
+        extra = out.extra_args
+        rvq = extra["tta_rvq"]
+        assert len(rvq["phase_token_ids"]) == 4
+        assert all(len(ids) == 1024 for ids in rvq["phase_token_ids"])
+        assert (rvq["start_tid"], rvq["end_tid"]) == (7, 8)
+        assert rvq["codec_cap"] == 4000 and rvq["start_in_prompt"] is True
+        assert extra["cfg_scale"] == 3.0
+        assert extra["cfg_role"] == "cond" and extra["cfg_pair_id"] == "req-9"
+        assert "<unk>" in extra["cfg_null_prompt"]
+
+    def test_explicit_scale_one_disables_cfg_but_keeps_rvq(self):
+        serving = self._serving()
+        params = SimpleNamespace(extra_args={"cfg_scale": 1.0})
+        out = serving._inject_audex_tta_args("req-1", self._prompt(), params)
+        assert "tta_rvq" in out.extra_args
+        assert "cfg_scale" not in out.extra_args
+        assert "cfg_role" not in out.extra_args
+
+    def test_rvq_contract_cached_per_process(self):
+        serving = self._serving()
+        out_a = serving._inject_audex_tta_args("a", self._prompt(), SimpleNamespace(extra_args=None))
+        out_b = serving._inject_audex_tta_args("b", self._prompt(), SimpleNamespace(extra_args=None))
+        assert out_a.extra_args["tta_rvq"] is out_b.extra_args["tta_rvq"]
+
+    def test_injection_never_mutates_input_params(self):
+        """The default online path hands the injector an element of the
+        engine's SHARED default_sampling_params_list (review P2): the
+        per-request pair state must land only on the returned clone."""
+        serving = self._serving()
+        shared = SimpleNamespace(extra_args=None)
+        out_1 = serving._inject_audex_tta_args("req-1", self._prompt(), shared)
+        out_2 = serving._inject_audex_tta_args("req-2", self._prompt(), shared)
+        assert shared.extra_args is None
+        assert out_1 is not shared and out_2 is not shared
+        assert out_1.extra_args["cfg_pair_id"] == "req-1"
+        assert out_2.extra_args["cfg_pair_id"] == "req-2"

tests/examples/offline_inference/test_ltx2_cfg_parallel_parity.py

@@ -1,243 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-import hashlib
-import os
-import subprocess
-import sys
-from pathlib import Path
-
-import numpy as np
-import pytest
-from PIL import Image
-
-from tests.helpers.mark import hardware_test
-
-REPO_ROOT = Path(__file__).resolve().parents[3]
-T2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "text_to_video" / "text_to_video.py"
-I2V_EXAMPLE = REPO_ROOT / "examples" / "offline_inference" / "image_to_video" / "image_to_video.py"
-
-T2V_PROMPT = (
-    "At sunrise, a glowing paper lantern boat drifts through a narrow canal between mossy stone walls, "
-    "soft fog above the water, the camera slowly gliding forward as golden reflections shimmer across "
-    "the ripples, cinematic, realistic, highly detailed."
-)
-T2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion, distorted, oversaturated, artifacts"
-I2V_PROMPT = "A cinematic dolly shot of a boat drifting on calm water at sunset"
-I2V_NEGATIVE_PROMPT = "worst quality, blurry, jittery motion"
-
-os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
-
-
-def _get_ltx2_model() -> str:
-    return "Lightricks/LTX-2"
-
-
-def _md5(path: Path) -> str:
-    digest = hashlib.md5(usedforsecurity=False)
-    with path.open("rb") as f:
-        for chunk in iter(lambda: f.read(1024 * 1024), b""):
-            digest.update(chunk)
-    return digest.hexdigest()
-
-
-def _make_deterministic_test_image(path: Path) -> None:
-    """Create a deterministic 256x256 test image for I2V tests."""
-    rng = np.random.RandomState(42)
-    img = Image.fromarray(rng.randint(0, 255, (256, 256, 3), dtype=np.uint8))
-    img.save(path)
-
-
-def _run_and_check(cmd: list[str], env: dict, output_path: Path, expected_md5: str) -> None:
-    result = subprocess.run(cmd, cwd=REPO_ROOT, env=env, capture_output=True, text=True, check=False)
-    assert result.returncode == 0, (
-        f"Command failed (exit {result.returncode}).\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
-    )
-    generated_md5 = _md5(output_path)
-    assert generated_md5 == expected_md5, (
-        f"Unexpected output md5: {generated_md5} != {expected_md5}.\nstdout:\n{result.stdout}\nstderr:\n{result.stderr}"
-    )
-
-
-# ── T2V tests ──
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@pytest.mark.slow
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
-def test_ltx2_t2v_cfg_parallel(tmp_path: Path):
-    """T2V with CFG=4.0, cfg-parallel-size=2."""
-    output = tmp_path / "t2v_cfg4.mp4"
-    env = os.environ.copy()
-    env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
-    cmd = [
-        sys.executable,
-        str(T2V_EXAMPLE),
-        "--model",
-        _get_ltx2_model(),
-        "--prompt",
-        T2V_PROMPT,
-        "--negative-prompt",
-        T2V_NEGATIVE_PROMPT,
-        "--height",
-        "256",
-        "--width",
-        "256",
-        "--num-frames",
-        "145",
-        "--num-inference-steps",
-        "6",
-        "--guidance-scale",
-        "4.0",
-        "--frame-rate",
-        "24",
-        "--fps",
-        "24",
-        "--seed",
-        "42",
-        "--cfg-parallel-size",
-        "2",
-        "--enforce-eager",
-        "--output",
-        str(output),
-    ]
-    _run_and_check(cmd, env, output, expected_md5="08e606b9c522fee4b6f30cee8b77db40")
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.slow
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_ltx2_t2v_no_cfg(tmp_path: Path):
-    """T2V with CFG=1.0 (no classifier-free guidance)."""
-    output = tmp_path / "t2v_nocfg.mp4"
-    env = os.environ.copy()
-    env.setdefault("CUDA_VISIBLE_DEVICES", "0")
-    cmd = [
-        sys.executable,
-        str(T2V_EXAMPLE),
-        "--model",
-        _get_ltx2_model(),
-        "--prompt",
-        T2V_PROMPT,
-        "--height",
-        "256",
-        "--width",
-        "256",
-        "--num-frames",
-        "145",
-        "--num-inference-steps",
-        "6",
-        "--guidance-scale",
-        "1.0",
-        "--frame-rate",
-        "24",
-        "--fps",
-        "24",
-        "--seed",
-        "42",
-        "--enforce-eager",
-        "--output",
-        str(output),
-    ]
-    _run_and_check(cmd, env, output, expected_md5="a83994b94b6e67c54a524e0383c45ce8")
-
-
-# ── I2V tests ──
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.parallel
-@pytest.mark.slow
-@hardware_test(res={"cuda": "L4"}, num_cards=2)
-def test_ltx2_i2v_cfg_parallel(tmp_path: Path):
-    """I2V with CFG=4.0, cfg-parallel-size=2."""
-    test_image = tmp_path / "test_input.png"
-    _make_deterministic_test_image(test_image)
-    output = tmp_path / "i2v_cfg4.mp4"
-    env = os.environ.copy()
-    env.setdefault("CUDA_VISIBLE_DEVICES", "0,1")
-    cmd = [
-        sys.executable,
-        str(I2V_EXAMPLE),
-        "--model",
-        _get_ltx2_model(),
-        "--model-class-name",
-        "LTX2ImageToVideoPipeline",
-        "--image",
-        str(test_image),
-        "--prompt",
-        I2V_PROMPT,
-        "--negative-prompt",
-        I2V_NEGATIVE_PROMPT,
-        "--height",
-        "256",
-        "--width",
-        "256",
-        "--num-frames",
-        "73",
-        "--num-inference-steps",
-        "6",
-        "--guidance-scale",
-        "4.0",
-        "--frame-rate",
-        "24",
-        "--fps",
-        "24",
-        "--seed",
-        "42",
-        "--cfg-parallel-size",
-        "2",
-        "--enforce-eager",
-        "--output",
-        str(output),
-    ]
-    _run_and_check(cmd, env, output, expected_md5="aed7e56084b36373244d8f839b16d115")
-
-
-@pytest.mark.full_model
-@pytest.mark.diffusion
-@pytest.mark.slow
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
-def test_ltx2_i2v_no_cfg(tmp_path: Path):
-    """I2V with CFG=1.0 (no classifier-free guidance)."""
-    test_image = tmp_path / "test_input.png"
-    _make_deterministic_test_image(test_image)
-    output = tmp_path / "i2v_nocfg.mp4"
-    env = os.environ.copy()
-    env.setdefault("CUDA_VISIBLE_DEVICES", "0")
-    cmd = [
-        sys.executable,
-        str(I2V_EXAMPLE),
-        "--model",
-        _get_ltx2_model(),
-        "--model-class-name",
-        "LTX2ImageToVideoPipeline",
-        "--image",
-        str(test_image),
-        "--prompt",
-        I2V_PROMPT,
-        "--height",
-        "256",
-        "--width",
-        "256",
-        "--num-frames",
-        "73",
-        "--num-inference-steps",
-        "6",
-        "--guidance-scale",
-        "1.0",
-        "--frame-rate",
-        "24",
-        "--fps",
-        "24",
-        "--seed",
-        "42",
-        "--enforce-eager",
-        "--output",
-        str(output),
-    ]
-    _run_and_check(cmd, env, output, expected_md5="81b21ede12753e9e14a357a6c548b666")

tests/helpers/media.py

@@ -529,6 +529,26 @@ def decode_b64_image(b64: str):
     return img
 
 
+def concat_audio(audio_val) -> np.ndarray:
+    """Flatten a multimodal audio payload to mono float32 samples.
+
+    Engines return ``multimodal_output["audio"]`` as a tensor, a list of
+    per-chunk tensors (streaming decoders), or an array-like; concatenate
+    in order and return a 1-D ``np.float32`` array (empty when a list has
+    no tensors).
+    """
+    import torch
+
+    if isinstance(audio_val, list):
+        tensors = [t.detach().cpu().float().reshape(-1) for t in audio_val if isinstance(t, torch.Tensor)]
+        if not tensors:
+            return np.zeros((0,), dtype=np.float32)
+        return torch.cat(tensors, dim=-1).numpy().astype(np.float32, copy=False)
+    if isinstance(audio_val, torch.Tensor):
+        return audio_val.detach().cpu().float().reshape(-1).numpy()
+    return np.asarray(audio_val, dtype=np.float32).reshape(-1)
+
+
 def preprocess_text(text):
     import opencc
 
@@ -695,6 +715,7 @@ def convert_audio_bytes_to_text(raw_bytes: bytes, model_size: str = "small") ->
 
 __all__ = [
     "_merge_base64_audio_to_segment",
+    "concat_audio",
     "convert_audio_bytes_to_text",
     "convert_audio_file_to_text",
     "cosine_similarity_text",

tests/model_executor/models/glm_image/test_glm_image_ar.py

@@ -91,6 +91,10 @@ GlmImageRotaryEmbedding = _ar_mod.GlmImageRotaryEmbedding
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
+def test_glm_image_ar_omits_hidden_pooler_payload():
+    assert GlmImageForConditionalGeneration.omni_pooler_payload_include_hidden is False
+
+
 # =============================================================================
 # Helper: Minimal config for testing
 # =============================================================================

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:

tests/model_tests/diffusion/test_alignment.py

@@ -30,13 +30,9 @@ EXCLUDED_MODELS = [
     "OvisImagePipeline",
     "WanPipeline",
     "WanVACEPipeline",
-    "LTX2ImageToVideoPipeline",
-    "LTX2TwoStagesPipeline",
-    "LTX2ImageToVideoTwoStagesPipeline",
+    "LTX2DistilledPipeline",
     "LTX2T2VDMD2Pipeline",
     "LTX2I2VDMD2Pipeline",
-    "LTX23Pipeline",
-    "LTX23ImageToVideoPipeline",
     "StableAudioPipeline",
     "WanImageToVideoPipeline",
     "WanS2VPipeline",

vllm_omni/config/config_factory.py

@@ -45,6 +45,26 @@ logger = init_logger(__name__)
 _DEFAULT_PARALLEL_DEGREE = 1
 
 
+def with_trust_remote_code_override(
+    overrides: Mapping[str, Any],
+    trust_remote_code: bool | None,
+) -> dict[str, Any]:
+    """Merge the tri-state ``trust_remote_code`` into an override mapping.
+
+    Single home for the precedence rule (explicit caller value > deploy
+    yaml per-stage value > vLLM default False): a non-None value becomes an
+    explicit override; ``None`` means "not specified" and leaves the deploy
+    yaml's per-stage setting in effect. The serve ``--trust-remote-code``
+    flag is store_true — its absent-False must be mapped to ``None`` at the
+    CLI boundary before reaching here, since it cannot express an explicit
+    False.
+    """
+    merged = dict(overrides)
+    if trust_remote_code is not None:
+        merged["trust_remote_code"] = trust_remote_code
+    return merged
+
+
 class StageConfigFactory:
     """Factory that loads pipeline YAML and merges CLI overrides.
 
@@ -293,7 +313,7 @@ class StageConfigFactory:
         cls,
         model: str,
         *,
-        trust_remote_code: bool,
+        trust_remote_code: bool | None,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None,
     ) -> VllmOmniConfig | None:
@@ -301,18 +321,19 @@ class StageConfigFactory:
         user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
         pipeline_cfg = cls.get_pipeline_config(
             model=model,
-            trust_remote_code=trust_remote_code,
+            # HF config resolution needs a real bool: transformers treats
+            # None as "prompt for consent", which blocks non-interactively.
+            trust_remote_code=bool(trust_remote_code),
             deploy_config_path=deploy_config_path,
             user_deploy_config=user_deploy_config,
         )
         if pipeline_cfg is None:
             return None
 
-        registry_cli_overrides = {
-            **cli_overrides,
-            "trust_remote_code": trust_remote_code,
-            "model": model,
-        }
+        registry_cli_overrides = with_trust_remote_code_override(
+            {**cli_overrides, "model": model},
+            trust_remote_code,
+        )
         return VllmOmniConfig.from_pipeline_config(
             pipeline_cfg,
             user_deploy_config=user_deploy_config,
@@ -325,7 +346,7 @@ class StageConfigFactory:
         cls,
         model: str,
         *,
-        trust_remote_code: bool,
+        trust_remote_code: bool | None,
         cli_overrides: dict[str, Any],
         deploy_config_path: str | None,
         strategy_specs: Mapping[Any, Any] | None = None,
@@ -339,17 +360,15 @@ class StageConfigFactory:
         user_deploy_config = cls._load_user_deploy_config(deploy_config_path)
         pipeline_cfg = cls.get_pipeline_config(
             model=model,
-            trust_remote_code=trust_remote_code,
+            # See create_from_model: HF resolution needs a real bool.
+            trust_remote_code=bool(trust_remote_code),
             deploy_config_path=deploy_config_path,
             user_deploy_config=user_deploy_config,
         )
         if pipeline_cfg is None:
             return None, None
 
-        legacy_cli_overrides = {
-            **cli_overrides,
-            "trust_remote_code": trust_remote_code,
-        }
+        legacy_cli_overrides = with_trust_remote_code_override(cli_overrides, trust_remote_code)
         return cls._create_legacy_from_registry(
             pipeline_cfg,
             legacy_cli_overrides,

vllm_omni/config/model.py

@@ -119,6 +119,7 @@ class OmniModelConfig(ModelConfig):
 
     stage_id: int = 0
     async_chunk: bool = False
+    retains_state_across_chunks: bool = False
     # Stage-1 active stream slots; 0 keeps legacy chunk-level round-robin.
     active_stream_window: int = 0
     model_stage: str = "thinker"
@@ -146,7 +147,10 @@ class OmniModelConfig(ModelConfig):
 
     @property
     def architectures(self) -> list[str]:
-        if self.model_arch is not None:
+        # Falsy (None or "") means "no stage override": fall back to the
+        # checkpoint config's own architectures. The stage-config builder
+        # emits None; "" is tolerated for legacy callers.
+        if self.model_arch:
             return [self.model_arch]
         return super().architectures
 

vllm_omni/config/omni_config.py

@@ -270,6 +270,11 @@ class OmniStageModelConfig:
     enable_multithread_weight_load: bool = True
     num_weight_load_threads: int = Field(default=4, ge=1)
     disable_autocast: bool = False
+    # Per-stage checkpoint/tokenizer subdirectories under the model root
+    # (e.g. Audex stage 0 → checkpoint_folder_audiogen). Mirrors
+    # StagePipelineConfig.model_subdir/tokenizer_subdir on the legacy path.
+    model_subdir: str | None = None
+    tokenizer_subdir: str | None = None
 
 
 @config
@@ -1117,6 +1122,10 @@ def _build_model_config(
     kwargs = _config_kwargs(engine)
     if "has_sampling_extra_args" not in kwargs:
         kwargs["has_sampling_extra_args"] = bool((default_sampling_params or {}).get("extra_args"))
+    if "model_subdir" not in kwargs and topology.model_subdir is not None:
+        kwargs["model_subdir"] = topology.model_subdir
+    if "tokenizer_subdir" not in kwargs and topology.tokenizer_subdir is not None:
+        kwargs["tokenizer_subdir"] = topology.tokenizer_subdir
     return OmniStageModelConfig(
         default_sampling_params=default_sampling_params,
         **kwargs,

vllm_omni/config/pipeline_registry.py

@@ -33,6 +33,12 @@ from vllm.logger import init_logger
 from vllm_omni.config.stage_config import (
     PipelineConfig,
 )
+from vllm_omni.model_executor.models.audex.pipeline import (
+    AUDEX_S2S_PIPELINE,
+    AUDEX_THINKER_ONLY_PIPELINE,
+    AUDEX_TTA_PIPELINE,
+    AUDEX_TTS_PIPELINE,
+)
 from vllm_omni.model_executor.models.aura_omni.pipeline import AURA_OMNI_PIPELINE
 from vllm_omni.model_executor.models.bagel.pipeline import (
     BAGEL_PIPELINE,
@@ -126,6 +132,14 @@ OMNI_PIPELINES: dict[str, PipelineConfig | PipelineResolverFunc] = {
     "wan2_2_ti2v": WAN2_2_TI2V_PIPELINE,
     "voxcpm2": VOXCPM2_PIPELINE,
     "cosyvoice3": COSYVOICE3_PIPELINE,
+    "audex_tts": AUDEX_TTS_PIPELINE,
+    "audex_tta": AUDEX_TTA_PIPELINE,
+    "audex_thinker_only": AUDEX_THINKER_ONLY_PIPELINE,
+    "audex_s2s": AUDEX_S2S_PIPELINE,
+    # Alias: the Nemotron-Labs-Audex-2B repo-root config.json reports
+    # ``model_type: nemotron_labs_audex``; bare ``vllm-omni serve <repo>``
+    # auto-detects through it and must land on the default (TTS) pipeline.
+    "nemotron_labs_audex": AUDEX_TTS_PIPELINE,
     "mimo_audio": MIMO_AUDIO_PIPELINE,
     "ming_tts": MING_TTS_PIPELINE,
     "ming_tts_moe": MING_TTS_MOE_PIPELINE,

vllm_omni/config/stage_config.py

@@ -223,6 +223,9 @@ class StagePipelineConfig:
     hf_config_name: str | None = None
     engine_output_type: str | None = None
     model_arch: str | None = None
+    # The model keeps per-request execution state while awaiting the next
+    # async chunk, so the parked request continues to consume model capacity.
+    retains_state_across_chunks: bool = False
     sampling_constraints: dict[str, Any] = field(default_factory=dict)
     custom_process_input_func: str | None = None
     custom_process_next_stage_input_func: str | None = None
@@ -821,7 +824,8 @@ def _build_engine_args(
     per-stage StageDeployConfig overrides take precedence when present (e.g.
     ``engine_extras`` can still carry a stage-specific ``dtype``).
     """
-    engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch}
+    engine_args: dict[str, Any] = {"model_arch": ps.model_arch or pipeline.model_arch or None}
+    engine_args["retains_state_across_chunks"] = ps.retains_state_across_chunks
     if ps.execution_type == StageExecutionType.DIFFUSION and ps.model_arch:
         engine_args.setdefault("model_class_name", ps.model_arch)
     if ps.engine_output_type:

vllm_omni/core/sched/omni_ar_scheduler.py

@@ -5,6 +5,7 @@ from collections.abc import Iterable
 from time import time
 from typing import Any
 
+import numpy as np
 from vllm.compilation.cuda_graph import CUDAGraphStat
 from vllm.distributed.kv_events import KVEventBatch
 from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats
@@ -36,6 +37,51 @@ from vllm_omni.outputs import OmniConnectorOutput
 logger = init_logger(__name__)
 
 
+class SampledLogprobContractError(RuntimeError):
+    """The model runner returned unusable sampled-token logprobs."""
+
+
+def _slice_sampled_logprobs(logprobs: Any, req_index: int, sampled_token_ids: list[int]) -> Any:
+    """Slice and validate the sampled-token logprobs for one AR request."""
+    if logprobs is None:
+        raise SampledLogprobContractError("AR logprobs were requested, but the model runner returned none")
+
+    sliced = logprobs.slice_request(req_index, len(sampled_token_ids))
+    token_rows = np.asarray(sliced.logprob_token_ids)
+    value_rows = np.asarray(sliced.logprobs)
+    expected_rows = len(sampled_token_ids)
+
+    if token_rows.ndim != 2 or value_rows.ndim != 2:
+        raise SampledLogprobContractError(
+            "AR sampled-token logprobs must be rank-2 arrays, "
+            f"got token_ids={token_rows.shape} logprobs={value_rows.shape}"
+        )
+    if token_rows.shape[0] != expected_rows or value_rows.shape[0] != expected_rows:
+        raise SampledLogprobContractError(
+            "AR sampled-token logprob row count does not match generated tokens: "
+            f"tokens={expected_rows} token_id_rows={token_rows.shape[0]} "
+            f"logprob_rows={value_rows.shape[0]}"
+        )
+    if expected_rows == 0:
+        return sliced
+    if token_rows.shape[1] == 0 or value_rows.shape[1] == 0:
+        raise SampledLogprobContractError("AR sampled-token logprob rows are empty")
+
+    sampled = np.asarray(sampled_token_ids)
+    if not np.array_equal(token_rows[:, 0], sampled):
+        mismatch = np.flatnonzero(token_rows[:, 0] != sampled)
+        first = int(mismatch[0])
+        raise SampledLogprobContractError(
+            "AR sampled-token logprobs are misaligned: "
+            f"row={first} generated_token={int(sampled[first])} "
+            f"logprob_token={int(token_rows[first, 0])}"
+        )
+    if not np.isfinite(value_rows[:, 0]).all():
+        bad_rows = np.flatnonzero(~np.isfinite(value_rows[:, 0])).tolist()
+        raise SampledLogprobContractError(f"AR sampled-token logprobs contain non-finite values at rows {bad_rows}")
+    return sliced
+
+
 class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
     """Synchronous AutoRegressive scheduler for vLLM-Omni. This class is also
     used as a base class for the OmniARAsyncScheduler and holds most of the
@@ -362,6 +408,27 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
 
             req_index = model_runner_output.req_id_to_index[req_id]
             generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
+            status_before_stop = request.status
+            new_logprobs = None
+            logprob_validation_failed = False
+
+            # Validate before mutating request token state. A bad runner output
+            # is request-local: terminate only this request and keep processing
+            # the rest of the batch.
+            if (
+                generated_token_ids
+                and request.sampling_params is not None
+                and request.sampling_params.num_logprobs is not None
+            ):
+                try:
+                    new_logprobs = _slice_sampled_logprobs(logprobs, req_index, generated_token_ids)
+                except SampledLogprobContractError as exc:
+                    logger.error("Invalid AR sampled-token logprobs for request %s: %s", req_id, exc)
+                    request.status = RequestStatus.FINISHED_ERROR
+                    request.stop_reason = str(exc)
+                    request.resumable = False
+                    generated_token_ids = []
+                    logprob_validation_failed = True
 
             scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
             if scheduled_spec_token_ids and generated_token_ids:
@@ -391,16 +458,14 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
             if request.has_encoder_inputs:
                 self._free_encoder_inputs(request)
 
-            stopped = False
+            stopped = logprob_validation_failed
             is_segment_finished = False
             finished = False
-            new_logprobs = None
             new_token_ids = generated_token_ids
             pooler_output = pooler_outputs[req_index] if pooler_outputs else None
             mm_output = mm_outputs[req_index] if mm_outputs else None
             inter_stage_output = inter_stage_outputs[req_index] if inter_stage_outputs else None
             kv_transfer_params = None
-            status_before_stop = request.status
             finish_reason = None
             routed_experts = None
 
@@ -471,10 +536,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
                 else:
                     stopped_preempted_reqs.add(request)
 
-            # Extract sample logprobs if needed.
-            if request.sampling_params is not None and request.sampling_params.num_logprobs is not None and logprobs:
-                new_logprobs = logprobs.slice_request(req_index, len(new_token_ids))
-
             if num_nans_in_logits is not None and req_id in num_nans_in_logits:
                 request.num_nans_in_logits = num_nans_in_logits[req_id]
 

vllm_omni/core/sched/omni_generation_scheduler.py

@@ -47,6 +47,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
         self.chunk_transfer_adapter = None
         if getattr(model_config, "async_chunk", False):
             self.chunk_transfer_adapter = OmniChunkTransferAdapter(self.vllm_config)
+        self._retains_state_across_chunks = bool(getattr(model_config, "retains_state_across_chunks", False))
         self._pending_finish_reqs: list[Request] = []
         self.input_coordinator: OmniSchedulingCoordinator | None = None
         if uses_full_payload_input_coordinator(model_config):
@@ -145,12 +146,17 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
 
         # Fast path selection and scheduling for one-shot generation requests,
         # independent of pooling_params.
-        while (
-            self.waiting
-            and token_budget > 0
-            and len(self.running) < self.max_num_running_reqs
-            and self._pause_state == PauseState.UNPAUSED
-        ):
+        while self.waiting and token_budget > 0 and self._pause_state == PauseState.UNPAUSED:
+            # Requests waiting for their next chunk are temporarily absent
+            # from `running`, but stateful models still retain their model
+            # runner slot. Mirror vLLM's treatment of
+            # `num_waiting_for_streaming_input` when enforcing max_num_seqs.
+            num_running = len(self.running)
+            if self._retains_state_across_chunks and self.chunk_transfer_adapter is not None:
+                num_running += self.chunk_transfer_adapter.num_running_waiting_for_chunk
+            if num_running >= self.max_num_running_reqs:
+                break
+
             request = self.waiting.peek_request()
             # OMNI: Skip requests that are not in self.requests
             if request.request_id not in self.requests or (

vllm_omni/core/sched/omni_scheduling_coordinator.py

@@ -51,6 +51,12 @@ _FULL_PAYLOAD_INPUT_STAGES: frozenset[tuple[str, str]] = frozenset(
         ("MiniCPMO45Code2Wav", "code2wav"),
         # cosyvoice3: cosyvoice3_talker (Stage 0) -> cosyvoice3_code2wav (Stage 1).
         ("CosyVoice3Model", "cosyvoice3_code2wav"),
+        # audex TTS sync path: thinker (Stage 0) -> streaming decoder (Stage 1).
+        # The default deploy is async_chunk; this covers async_chunk: false.
+        ("AudexCode2Wav", "audex_code2wav"),
+        # audex TTA: tta thinker (Stage 0) -> XCodec1 (Stage 1, always sync
+        # full-payload — CNN codec decoded over the full sequence).
+        ("AudexXCodec1", "audex_xcodec"),
         # indextts2: indextts2_talker (Stage 0) -> indextts2_s2mel_decoder
         # (Stage 1). Stage 1 consumes the complete mel/latent payload.
         ("IndexTTS2S2MelDecoder", "indextts2_s2mel_decoder"),

vllm_omni/deploy/audex_s2s.yaml

@@ -0,0 +1,65 @@
+# Audex (Nemotron-Labs-Audex-2B) full deploy: cascaded speech-to-speech.
+#
+# One audio-capable thinker (checkpoint_folder_full) serves the official
+# three-pass cascade — (1) ASR: audio -> transcript, (2) chat: transcript ->
+# answer, (3) TTS: answer -> speech — plus the streaming causal speech
+# decoder as stage 1. Per-request output modalities route the passes: text
+# passes finish at stage 0; only audio-modality requests stream codec frames
+# into code2wav. CFG applies to the TTS pass via extra_params.cfg_scale
+# (recommended 1.5; 1.0 = off).
+pipeline: audex_s2s
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.5
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_num_batched_tokens: 8192
+    # Sound embeddings (up to 30 clips x 750) + speech-token generation.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_s2s_30b.yaml

@@ -0,0 +1,60 @@
+# Audex 30B-A3B full deploy: cascaded speech-to-speech.
+# See audex_s2s.yaml for the cascade/routing notes and audex_tts_30b.yaml
+# for the 30B sizing notes (hybrid Mamba: no prefix caching; single-H100
+# layout, TP2 fallback; pass this yaml explicitly).
+pipeline: audex_s2s
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.85
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 8192
+    # Sound embeddings (up to 30 clips x 750) + speech-token generation.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.08
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_thinker_only.yaml

@@ -0,0 +1,25 @@
+# Audex (Nemotron-Labs-Audex-2B) thinker-only deploy: audio understanding
+# (ASR / audio QA) -> text. Single-stage configuration of the audio-capable
+# full checkpoint (checkpoint_folder_full): NV-Whisper encoder + relu2
+# projector + 2B dense LM.
+#
+# The official audioqa script decodes ASR at temperature 0; max_tokens 256
+# covers transcription and short QA answers (raise for long-form answers).
+pipeline: audex_thinker_only
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.7
+    trust_remote_code: true
+    enable_prefix_caching: true
+    max_num_batched_tokens: 8192
+    # 30 clips x 750 sound embeddings (900 s cap) + text headroom.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 256
+      seed: 0

vllm_omni/deploy/audex_thinker_only_30b.yaml

@@ -0,0 +1,22 @@
+# Audex 30B-A3B thinker-only deploy: audio understanding (ASR / audio QA).
+# Single-stage NemotronHAudexForConditionalGeneration on
+# checkpoint_folder_full. Hybrid Mamba: no prefix caching. Pass this yaml
+# explicitly (see audex_tts_30b.yaml).
+pipeline: audex_thinker_only
+async_chunk: false
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.9
+    trust_remote_code: true
+    enable_prefix_caching: false
+    max_num_batched_tokens: 8192
+    # 30 clips x 750 sound embeddings (900 s cap) + text headroom.
+    max_model_len: 32768
+    dtype: bfloat16
+    devices: "0"
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 256
+      seed: 0

vllm_omni/deploy/audex_tta.yaml

@@ -0,0 +1,69 @@
+# Audex (Nemotron-Labs-Audex-2B) text-to-audio deploy: TTA thinker → XCodec1.
+#
+# Notable choices:
+#   * Sync full-payload (async_chunk: false): XCodec1 is a CNN codec the
+#     official flow decodes over the full sequence; there is no streaming
+#     reference and chunked decode risks boundary artifacts.
+#   * Stage 0 sampling defaults follow the official TTA script (temperature
+#     1.0, top_k 80; stop on <audiogen_end>; codec cap 4000 tokens = 1000
+#     frames generated, decode truncates at 500 frames ≈ 10 s).
+#   * Classifier-free guidance is effectively mandatory for TTA quality; the
+#     official default is cfg_scale 3.0. The RVQ phase mask is required for
+#     validity — both logits processors are registered here and driven per
+#     request via extra_args (cfg_* and tta_rvq).
+#   * Stage 1's model is the EXTERNAL hf-audio XCodec1 checkpoint (override
+#     with a local path here or via XCODEC1_PATH); its tokenizer is borrowed
+#     from the thinker checkpoint automatically.
+pipeline: audex_tta
+async_chunk: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      audiocodec_token_offset: 196613
+      audiocodec_vocab_size: 8192
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.4
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 8192
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+      - "vllm_omni.model_executor.models.audex.tta:TTARVQPhaseMaskLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 1.0
+      top_k: 80
+      max_tokens: 4200
+      seed: 0
+
+  - stage_id: 1
+    model: hf-audio/xcodec-hubert-general-balanced
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_tta_30b.yaml

@@ -0,0 +1,57 @@
+# Audex 30B-A3B text-to-audio deploy: TTA thinker → XCodec1.
+# See audex_tta.yaml for the task notes and audex_tts_30b.yaml for the
+# 30B sizing notes (hybrid Mamba: no prefix caching; single-H100 layout,
+# TP2 fallback; pass this yaml explicitly).
+pipeline: audex_tta
+async_chunk: false
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      audiocodec_token_offset: 196613
+      audiocodec_vocab_size: 8192
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.85
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 8192
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+      - "vllm_omni.model_executor.models.audex.tta:TTARVQPhaseMaskLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 1.0
+      top_k: 80
+      max_tokens: 4200
+      seed: 0
+
+  - stage_id: 1
+    model: hf-audio/xcodec-hubert-general-balanced
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.08
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_tts.yaml

@@ -0,0 +1,83 @@
+# Audex (Nemotron-Labs-Audex-2B) TTS deploy: thinker → code2wav via
+# shared-memory chunk streaming. Verified on 1x H100.
+#
+# Notable choices:
+#   * The code2wav decoder is a stateful streaming session per request — no
+#     left-context replay; chunks carry only new frames, so
+#     codec_left_context_frames stays 0.
+#   * codec_token_offset / codec_vocab_size mirror the checkpoint tokenizer's
+#     <speechcodec_N> block (pinned by a unit test).
+#   * Stage 0 sampling defaults follow the official TTS script
+#     (temperature 0.1, top_k 80, max_tokens 2048).
+#   * Classifier-free guidance: the stage-0 logits processor + scheduler
+#     pair patches are always installed but idle until a request carries
+#     extra_args.cfg_scale > 1.0 (serving: extra_params.cfg_scale; offline:
+#     --cfg-scale). Guided requests decode at temperature 0.05 (CFG
+#     sharpens the distribution; measured en-24 cfg1.5 CER: 6.87% at
+#     temp 0.05 vs 7.31% at the unguided 0.1 — baseline 7.24%).
+#     Every guided request adds an unconditional companion sequence in
+#     stage 0 — at max_num_seqs 8 that means at most 4 guided requests
+#     decode concurrently; raise max_num_seqs if CFG-heavy traffic needs
+#     more parallelism.
+pipeline: audex_tts
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      # Emit the first chunk early for a fast first audible sample.
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      # Decoder-internal emit granularity (frames per decode call).
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.4
+    trust_remote_code: true
+    enable_prefix_caching: true
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 4096
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.25
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    # Sync scheduling: with async scheduling the runner builds per-request
+    # payloads against the NEXT step's input batch, tripping the one-entry-
+    # per-request check when chunk arrival staggers requests.
+    async_scheduling: false
+    # async_chunk appends codec frames per step; max_model_len must cover the
+    # accumulated stream (max_tokens 2048 codes plus prompt headroom).
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/audex_tts_30b.yaml

@@ -0,0 +1,76 @@
+# Audex 30B-A3B (nvidia/Nemotron-Labs-Audex-30B-A3B) TTS deploy.
+#
+# Same pipeline/topology/token space as the 2B (audex_tts.yaml); the thinker
+# is the hybrid Mamba+MoE NemotronHForCausalLM (~60 GB bf16), so:
+#   * enable_prefix_caching: false — hybrid Mamba stage.
+#   * Single H100 80GB layout: thinker 0.85 + decoder 0.08 share GPU 0.
+#     If weights + state do not fit (long sequences OOM), set
+#     tensor_parallel_size: 2 on stage 0 and move stage 1 to devices "2".
+#   * CUDA graphs stay ON for stage 0: the hybrid Mamba+MoE thinker
+#     captures cleanly and review benchmarks measured ~7x lower TTS RTF
+#     vs enforce_eager (0.30 vs 2.08 single-stream), on par with the 2B
+#     and NVIDIA's reference demo.
+#   * Classifier-free guidance quality (en-24, self-transcribed ASR CER):
+#     baseline cfg=1.0 -> 0.72%; guided cfg=1.5 at temperature 0.1 -> 0.72%
+#     (gate: guided <= baseline, PASS). Note the 30B's best guided
+#     temperature is 0.1 (the yaml default) — unlike the 2B's 0.05, which
+#     measured 0.79% here. The serving layer currently injects 0.05 for
+#     guided requests on all sizes; acceptable (still within the anchor
+#     +2pp standing gate) but a size-aware refinement is a known follow-up.
+# ALWAYS pass this yaml explicitly (--stage-configs-path): bare serve of the
+# 30B repo auto-detects the shared model_type and lands on the 2B-tuned yaml.
+pipeline: audex_tts
+async_chunk: true
+
+connectors:
+  connector_of_shared_memory:
+    name: SharedMemoryConnector
+    extra:
+      shm_threshold_bytes: 65536
+      codec_streaming: true
+      connector_get_sleep_s: 0.01
+      connector_get_max_wait_first_chunk: 3000
+      connector_get_max_wait: 300
+      codec_chunk_frames: 25
+      initial_codec_chunk_frames: 1
+      codec_token_offset: 131077
+      codec_vocab_size: 65536
+      decoder_emit_chunk_frames: 5
+
+stages:
+  - stage_id: 0
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.85
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: true
+    max_num_batched_tokens: 512
+    max_model_len: 4096
+    dtype: bfloat16
+    devices: "0"
+    logits_processors:
+      - "vllm_omni.model_executor.models.audex.cfg:AudexCFGLogitsProcessor"
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.1
+      top_k: 80
+      max_tokens: 2048
+      seed: 0
+
+  - stage_id: 1
+    max_num_seqs: 8
+    gpu_memory_utilization: 0.08
+    enforce_eager: true
+    trust_remote_code: true
+    enable_prefix_caching: false
+    async_scheduling: false
+    max_num_batched_tokens: 8192
+    max_model_len: 8192
+    devices: "0"
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+    default_sampling_params:
+      temperature: 0.0
+      max_tokens: 8192
+      seed: 0

vllm_omni/deploy/dreamzero_tp1_cfg2.yaml

@@ -11,6 +11,8 @@ stages:
     max_num_seqs: 1
     enforce_eager: false
     model_class_name: DreamZeroPipeline
+    # DreamZero runs on the AR-Diffusion engine (engine-level paged KV cache).
+    engine_backend: vllm_omni.experimental.ar_diffusion.engine.ARDiffusionEngine
     parallel_config:
       tensor_parallel_size: 1
       cfg_parallel_size: 2

vllm_omni/deploy/moss_tts_local.yaml

@@ -9,24 +9,25 @@ connectors:
   shm:
     name: SharedMemoryConnector
     extra:
-      codec_streaming: true
       connector_get_sleep_s: 0.005
       connector_get_max_wait_first_chunk: 1000
       connector_get_max_wait: 300
       initial_codec_chunk_frames: 1   # 1 frame * 80 ms/frame = 80 ms first audio chunk
       codec_chunk_frames: 15          # 15 frames * 80 ms/frame = 1.2 s steady chunks
       codec_left_context_frames: 0
-      codec_stream_slots: 8
 
 stages:
   - stage_id: 0
-    max_num_seqs: 8
+    max_num_seqs: 64
     gpu_memory_utilization: 0.60
     enable_prefix_caching: false
     async_scheduling: true
     max_num_batched_tokens: 65536
     max_model_len: 4096
     devices: "0"
+    # profiler_config:
+    #   profiler: torch
+    #   torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage0/batch_8
     output_connectors:
       to_stage_1: shm
     default_sampling_params:
@@ -38,17 +39,22 @@ stages:
       repetition_penalty: 1.1
 
   - stage_id: 1
-    max_num_seqs: 8
-    gpu_memory_utilization: 0.15
+    max_num_seqs: 64
+    gpu_memory_utilization: 0.30
     enforce_eager: false
     enable_prefix_caching: false
     async_scheduling: true
     compilation_config:
-      cudagraph_capture_sizes: [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]
+      # Reused by the codec as execution-batch buckets; T is derived from
+      # initial_codec_chunk_frames and codec_chunk_frames.
+      cudagraph_capture_sizes: [1, 2, 4, 8, 16, 32, 64]
       cudagraph_num_of_warmups: 1
+    # profiler_config:
+    #   profiler: torch
+    #   torch_profiler_dir: /root/vllm-omni-workspace/vllm-omni/results/moss_local_stage1/batch_8
     max_num_batched_tokens: 65536
     max_model_len: 65536
-    devices: "0"
+    devices: "1"
     input_connectors:
       from_stage_0: shm
     default_sampling_params:

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -634,7 +634,7 @@ class DiffusersPipelineLoader:
         # Discover pipeline components (DiT, encoders, VAEs) via
         # ModuleDiscovery, which consults SupportsComponentDiscovery
         # when available and falls back to well-known attribute names.
-        # This supports nested pipelines (e.g. LTX2TwoStagesPipeline
+        # This supports nested pipelines (e.g. LTX2DistilledPipeline
         # where the transformer lives at "pipe.transformer").
         discovered_modules = ModuleDiscovery.discover(model)
 

vllm_omni/diffusion/models/ltx2/__init__.py

@@ -9,28 +9,16 @@ from vllm_omni.diffusion.models.ltx2.ltx2_components import (
 from vllm_omni.diffusion.models.ltx2.ltx2_transformer import LTX2VideoTransformer3DModel
 from vllm_omni.diffusion.models.ltx2.pipeline_ltx2 import (
     LTX2I2VDMD2Pipeline,
-    LTX2ImageToVideoPipeline,
     LTX2Pipeline,
     LTX2T2VDMD2Pipeline,
-    LTX23ImageToVideoPipeline,
-    LTX23Pipeline,
-)
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
-from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import (
-    LTX2ImageToVideoTwoStagesPipeline,
-    LTX2TwoStagesPipeline,
 )
+from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_two_stage import LTX2DistilledPipeline
 
 __all__ = [
     "LTX2Pipeline",
     "LTX2T2VDMD2Pipeline",
-    "LTX2ImageToVideoPipeline",
     "LTX2I2VDMD2Pipeline",
-    "LTX2LatentUpsamplePipeline",
-    "LTX2TwoStagesPipeline",
-    "LTX2ImageToVideoTwoStagesPipeline",
-    "LTX23Pipeline",
-    "LTX23ImageToVideoPipeline",
+    "LTX2DistilledPipeline",
     "get_ltx2_post_process_func",
     "load_transformer_config",
     "create_transformer_from_config",

vllm_omni/diffusion/models/ltx2/ltx2_components.py

@@ -7,6 +7,7 @@ from __future__ import annotations
 
 import inspect
 import json
+import logging
 import os
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
@@ -14,11 +15,14 @@ from typing import TYPE_CHECKING, Any
 import torch
 from diffusers import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video, FlowMatchEulerDiscreteScheduler
 from diffusers.pipelines.ltx2 import LTX2TextConnectors
+from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
 from diffusers.pipelines.ltx2.vocoder import LTX2Vocoder
 from diffusers.video_processor import VideoProcessor
 from huggingface_hub import hf_hub_download
 from transformers import AutoTokenizer, Gemma3ForConditionalGeneration
 
+from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
+from vllm_omni.diffusion.attention.layer import Attention as OmniAttention
 from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import DistributedAutoencoderKLLTX2Video
 from vllm_omni.diffusion.distributed.utils import get_local_device
 from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
@@ -28,7 +32,12 @@ from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
 if TYPE_CHECKING:
     from vllm.model_executor.layers.quantization.base_config import QuantizationConfig
 
-from .ltx2_transformer import LTX2VideoTransformer3DModel
+from .ltx2_transformer import (
+    LTX2VideoTransformer3DModel,
+    apply_interleaved_rotary_emb,
+    apply_split_rotary_emb,
+    to_ltx_padding_mask,
+)
 
 try:
     from diffusers.pipelines.ltx2.vocoder import LTX2VocoderWithBWE
@@ -44,7 +53,9 @@ _LTX_COMPONENT_SUBFOLDERS = (
     "audio_vae",
     "vocoder",
     "scheduler",
+    "latent_upsampler",
 )
+logger = logging.getLogger(__name__)
 
 
 @dataclass(frozen=True)
@@ -81,6 +92,164 @@ LTX23_COMPONENT_PROFILE = LTXComponentProfile(
     vocoder_fallback_cls=LTX2Vocoder,
 )
 
+LTX2_DISTILLED_COMPONENT_PROFILE = LTXComponentProfile(
+    name="ltx2_distilled",
+    dit_modules=("transformer",),
+    encoder_modules=("text_encoder", "connectors"),
+    vae_modules=("vae", "audio_vae"),
+    resident_modules=("vocoder", "latent_upsampler"),
+    video_vae_cls=DistributedAutoencoderKLLTX2Video,
+)
+
+
+_COMPONENT_PROFILES: dict[tuple[str, str], LTXComponentProfile] = {
+    ("one_stage", "2"): LTX2_COMPONENT_PROFILE,
+    ("one_stage", "2.3"): LTX23_COMPONENT_PROFILE,
+    ("distilled_two_stage", "2"): LTX2_DISTILLED_COMPONENT_PROFILE,
+    ("dmd2", "2"): LTX2_COMPONENT_PROFILE,
+    ("dmd2", "2.3"): LTX23_COMPONENT_PROFILE,
+}
+
+
+def resolve_ltx_component_profile(pipeline_kind: str, model_version: str) -> LTXComponentProfile:
+    """Resolve component construction independently from execution recipes."""
+    try:
+        return _COMPONENT_PROFILES[(pipeline_kind, model_version)]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported LTX component kind/version: {pipeline_kind!r}/{model_version!r}.") from exc
+
+
+def _load_ltx_metadata_json(model: str, filename: str) -> dict[str, Any]:
+    """Load small checkpoint metadata without relying on repository names."""
+    if os.path.isdir(model):
+        path = os.path.join(model, filename)
+        if not os.path.isfile(path):
+            return {}
+    else:
+        try:
+            path = hf_hub_download(repo_id=model, filename=filename)
+        except Exception:
+            return {}
+    try:
+        with open(path) as config_file:
+            value = json.load(config_file)
+    except (OSError, TypeError, ValueError):
+        return {}
+    return value if isinstance(value, dict) else {}
+
+
+def detect_ltx_model_version(model: str) -> str:
+    """Detect LTX-2 versus LTX-2.3 from checkpoint component metadata.
+
+    Official checkpoints use ``model_version`` metadata. Diffusers repositories
+    expose the equivalent distinction through the BWE vocoder introduced by
+    LTX-2.3. Unknown conversions retain the official LTX-2 fallback.
+    """
+    model_index = _load_ltx_metadata_json(model, "model_index.json")
+    if str(model_index.get("model_version", "")).startswith("2.3"):
+        return "2.3"
+
+    vocoder_entry = model_index.get("vocoder")
+    if isinstance(vocoder_entry, (list, tuple)) and vocoder_entry:
+        vocoder_class = str(vocoder_entry[-1])
+    elif isinstance(vocoder_entry, dict):
+        vocoder_class = str(vocoder_entry.get("_class_name", ""))
+    else:
+        vocoder_class = ""
+    if vocoder_class == "LTX2VocoderWithBWE":
+        return "2.3"
+
+    vocoder_config = _load_ltx_metadata_json(model, "vocoder/config.json")
+    if str(vocoder_config.get("model_version", "")).startswith("2.3"):
+        return "2.3"
+    if vocoder_config.get("_class_name") == "LTX2VocoderWithBWE":
+        return "2.3"
+
+    logger.info("Using LTX-2 defaults for checkpoint %s", model)
+    return "2"
+
+
+class _LTXConnectorAttnProcessor:
+    """Preserve official connector math around Omni attention dispatch."""
+
+    def __init__(self, *, has_learned_registers: bool) -> None:
+        self.has_learned_registers = has_learned_registers
+
+    def __call__(
+        self,
+        attn: Any,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: torch.Tensor | None = None,
+        attention_mask: torch.Tensor | None = None,
+        query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+    ) -> torch.Tensor:
+        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states
+        gate_logits = attn.to_gate_logits(hidden_states) if attn.to_gate_logits is not None else None
+
+        query = attn.norm_q(attn.to_q(hidden_states))
+        key = attn.norm_k(attn.to_k(encoder_hidden_states))
+        value = attn.to_v(encoder_hidden_states)
+
+        if query_rotary_emb is not None:
+            key_rotary_emb = key_rotary_emb if key_rotary_emb is not None else query_rotary_emb
+            if attn.rope_type == "interleaved":
+                query = apply_interleaved_rotary_emb(query, query_rotary_emb)
+                key = apply_interleaved_rotary_emb(key, key_rotary_emb)
+            elif attn.rope_type == "split":
+                query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim)
+                key = apply_split_rotary_emb(key, key_rotary_emb, head_dim=attn.head_dim)
+            else:
+                raise ValueError(f"Unsupported LTX connector RoPE type: {attn.rope_type}")
+
+        batch_size, _, inner_dim = query.shape
+        head_dim = inner_dim // attn.heads
+        kv_heads = attn.inner_kv_dim // attn.head_dim
+        query = query.view(batch_size, -1, attn.heads, head_dim)
+        key = key.view(batch_size, -1, kv_heads, head_dim)
+        value = value.view(batch_size, -1, kv_heads, head_dim)
+
+        # The connector replaces padding tokens with learned registers before
+        # entering its blocks, so every key is valid and the old padding mask
+        # becomes an all-keep no-op.
+        if self.has_learned_registers:
+            attention_mask = None
+        elif attention_mask is not None and attn.omni_attention.attn_backend.get_name().upper() == "FLASH_ATTN":
+            attention_mask = to_ltx_padding_mask(attention_mask)
+        attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
+        hidden_states = attn.omni_attention(query, key, value, attn_metadata)
+        hidden_states = hidden_states.reshape(batch_size, -1, inner_dim)
+
+        if gate_logits is not None:
+            hidden_states = hidden_states.unflatten(2, (attn.heads, -1))
+            hidden_states = hidden_states * (2.0 * torch.sigmoid(gate_logits)).unsqueeze(-1)
+            hidden_states = hidden_states.flatten(2, 3)
+
+        hidden_states = attn.to_out[0](hidden_states)
+        return attn.to_out[1](hidden_states)
+
+
+def _install_connector_attention(connectors: LTX2TextConnectors) -> None:
+    for connector_name in ("video_connector", "audio_connector"):
+        connector = getattr(connectors, connector_name, None)
+        has_learned_registers = getattr(connector, "learnable_registers", None) is not None
+        for block_index, block in enumerate(getattr(connector, "transformer_blocks", ())):
+            attention = getattr(block, "attn1", None)
+            if attention is not None:
+                attention.omni_attention = OmniAttention(
+                    num_heads=attention.heads,
+                    head_size=attention.head_dim,
+                    num_kv_heads=attention.inner_kv_dim // attention.head_dim,
+                    softmax_scale=1.0 / (attention.head_dim**0.5),
+                    causal=False,
+                    prefix=f"connectors.{connector_name}.transformer_blocks.{block_index}.attn1",
+                    role="ltx2.connector",
+                    role_category="self",
+                    skip_sequence_parallel=True,
+                    disable_kv_quant=True,
+                )
+                attention.set_processor(_LTXConnectorAttnProcessor(has_learned_registers=has_learned_registers))
+
 
 def _detect_vocoder_output_sample_rate(model: str) -> int | None:
     """Read the generated waveform sample rate from the vocoder config."""
@@ -188,6 +357,7 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
         local_files_only=local_files_only,
         dtype=dtype,
     )
+    _install_connector_attention(pipeline.connectors)
     pipeline.vae = _load_component(
         profile.video_vae_cls,
         model,
@@ -221,6 +391,18 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
             dtype=dtype,
         )
 
+    if "latent_upsampler" in profile.resident_modules:
+        # BlurDownsample constructs an integer kernel that must be initialized
+        # on CPU; component placement is handled uniformly after construction.
+        with torch.device("cpu"):
+            pipeline.latent_upsampler = _load_component(
+                LTX2LatentUpsamplerModel,
+                model,
+                "latent_upsampler",
+                local_files_only=local_files_only,
+                dtype=dtype,
+            )
+
     transformer_config = load_transformer_config(model, "transformer", local_files_only)
     quant_config = getattr(od_config, "quantization_config", None)
     pipeline.transformer = create_transformer_from_config(transformer_config, quant_config=quant_config)
@@ -249,8 +431,6 @@ def initialize_pipeline_components(pipeline: Any, od_config: Any) -> None:
             tokenizer_max_length = getattr(encoder_config, "max_seq_len", None)
     pipeline.tokenizer_max_length = int(tokenizer_max_length or 1024)
 
-    pipeline._guidance_scale = None
-    pipeline._guidance_rescale = None
     pipeline._interrupt = False
 
 
@@ -262,22 +442,15 @@ def load_transformer_config(
     """Load an LTX transformer config from a local model or the HF Hub."""
     if local_files_only:
         config_path = os.path.join(model_path, subfolder, "config.json")
-        if os.path.exists(config_path):
-            with open(config_path) as config_file:
-                return json.load(config_file)
+        if not os.path.exists(config_path):
+            raise FileNotFoundError(f"LTX transformer config not found: {config_path}")
     else:
-        try:
-            from huggingface_hub import hf_hub_download
-
-            config_path = hf_hub_download(
-                repo_id=model_path,
-                filename=f"{subfolder}/config.json",
-            )
-            with open(config_path) as config_file:
-                return json.load(config_file)
-        except Exception:
-            pass
-    return {}
+        config_path = hf_hub_download(
+            repo_id=model_path,
+            filename=f"{subfolder}/config.json",
+        )
+    with open(config_path) as config_file:
+        return json.load(config_file)
 
 
 def create_transformer_from_config(

vllm_omni/diffusion/models/ltx2/ltx2_conditioning.py

@@ -5,17 +5,20 @@
 
 from __future__ import annotations
 
+import math
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any
 
+import numpy as np
 import PIL.Image
 import torch
+import torch.nn.functional as F
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
 from diffusers.utils.torch_utils import randn_tensor
 from diffusers.video_processor import VideoProcessor
 
 from . import ltx2_latents as latent_ops
-from .ltx2_denoise import I2VVideoAudioScheduler
+from .ltx2_guidance import euler_step_from_velocity
 
 if TYPE_CHECKING:
     from .ltx2_denoise import LTXDenoiseContext, LTXForwardContext
@@ -44,6 +47,41 @@ def _repeat_prompt_tensor_for_outputs(tensor: torch.Tensor, num_outputs: int) ->
     return tensor.repeat_interleave(num_outputs, dim=0)
 
 
+def _preprocess_i2v_pil_images(
+    images: PIL.Image.Image | list[PIL.Image.Image],
+    *,
+    height: int,
+    width: int,
+    device: torch.device | None = None,
+    dtype: torch.dtype | None = None,
+) -> torch.Tensor:
+    """Match official LTX aspect-preserving resize and center crop."""
+    image_list = images if isinstance(images, list) else [images]
+    processed = []
+    for image in image_list:
+        pixels = (
+            torch.from_numpy(np.array(image.convert("RGB"), copy=True))
+            .permute(2, 0, 1)
+            .unsqueeze(0)
+            .to(device=device, dtype=torch.float32)
+        )
+        source_height, source_width = pixels.shape[-2:]
+        scale = max(height / source_height, width / source_width)
+        resized_height = math.ceil(source_height * scale)
+        resized_width = math.ceil(source_width * scale)
+        pixels = F.interpolate(
+            pixels,
+            size=(resized_height, resized_width),
+            mode="bilinear",
+            align_corners=False,
+        )
+        crop_top = (resized_height - height) // 2
+        crop_left = (resized_width - width) // 2
+        pixels = pixels[:, :, crop_top : crop_top + height, crop_left : crop_left + width]
+        processed.append((pixels / 127.5 - 1.0).to(dtype=dtype))
+    return torch.cat(processed, dim=0)
+
+
 class LTXTextConditioningMixin:
     """Shared Gemma encoding and text connector orchestration."""
 
@@ -113,21 +151,64 @@ class LTXTextConditioningMixin:
         batch_size = len(prompt) if prompt is not None else prompt_embeds.shape[0]
         negative_prompt_embeds_provided = negative_prompt_embeds is not None
 
-        if prompt_embeds is None:
-            prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
-                prompt=prompt,
-                num_videos_per_prompt=num_videos_per_prompt,
+        if do_classifier_free_guidance and prompt_embeds is None and negative_prompt_embeds is None:
+            negative_prompt = negative_prompt or ""
+            negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt
+            if prompt is not None and type(prompt) is not type(negative_prompt):
+                raise TypeError(
+                    f"`negative_prompt` should be the same type as `prompt`, but got {type(negative_prompt)} !="
+                    f" {type(prompt)}."
+                )
+            if isinstance(negative_prompt, list) and batch_size != len(negative_prompt):
+                raise ValueError(
+                    f"`negative_prompt` has batch size {len(negative_prompt)}, but `prompt` has batch size"
+                    f" {batch_size}."
+                )
+
+            # Official LTX encodes positive and negative prompts in one Gemma batch;
+            # separate bf16 GEMMs produce a different conditioning trajectory.
+            combined_embeds, combined_mask = self._get_gemma_prompt_embeds(
+                prompt=prompt + negative_prompt,
+                num_videos_per_prompt=1,
                 max_sequence_length=max_sequence_length,
                 scale_factor=scale_factor,
                 device=device,
                 dtype=dtype,
             )
-        elif num_videos_per_prompt > 1:
-            prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
-            prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
-                prompt_attention_mask,
-                num_videos_per_prompt,
-            )
+            prompt_embeds = combined_embeds[:batch_size]
+            negative_prompt_embeds = combined_embeds[batch_size:]
+            prompt_attention_mask = combined_mask[:batch_size]
+            negative_prompt_attention_mask = combined_mask[batch_size:]
+            if num_videos_per_prompt > 1:
+                prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
+                negative_prompt_embeds = _repeat_prompt_tensor_for_outputs(
+                    negative_prompt_embeds,
+                    num_videos_per_prompt,
+                )
+                prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
+                    prompt_attention_mask,
+                    num_videos_per_prompt,
+                )
+                negative_prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
+                    negative_prompt_attention_mask,
+                    num_videos_per_prompt,
+                )
+        else:
+            if prompt_embeds is None:
+                prompt_embeds, prompt_attention_mask = self._get_gemma_prompt_embeds(
+                    prompt=prompt,
+                    num_videos_per_prompt=num_videos_per_prompt,
+                    max_sequence_length=max_sequence_length,
+                    scale_factor=scale_factor,
+                    device=device,
+                    dtype=dtype,
+                )
+            elif num_videos_per_prompt > 1:
+                prompt_embeds = _repeat_prompt_tensor_for_outputs(prompt_embeds, num_videos_per_prompt)
+                prompt_attention_mask = _repeat_prompt_tensor_for_outputs(
+                    prompt_attention_mask,
+                    num_videos_per_prompt,
+                )
 
         if do_classifier_free_guidance and negative_prompt_embeds is None:
             negative_prompt = negative_prompt or ""
@@ -243,9 +324,10 @@ class LTXTextConditioningMixin:
 
 
 class LTXI2VConditioningMixin:
-    """First-frame conditioning behavior common to LTX2 and LTX2.3."""
+    """First-frame conditioning with optional unified T2V/I2V dispatch."""
 
     support_image_input = True
+    unified_text_image_entry = False
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -254,10 +336,6 @@ class LTXI2VConditioningMixin:
             resample="bilinear",
         )
 
-    def forward(self, req: Any, image: Any | None = None, *args: Any, **kwargs: Any) -> Any:
-        """Preserve the public I2V positional API while sharing the T2V runner."""
-        return super().forward(req, *args, image=image, **kwargs)
-
     @staticmethod
     def _resolve_single_prompt_image(raw_image: Any) -> Any:
         if isinstance(raw_image, list):
@@ -283,11 +361,12 @@ class LTXI2VConditioningMixin:
         image: Any | None,
         request_inputs: LTXRequestInputs,
     ) -> Any | None:
-        if image is not None or not req.prompts:
+        prompts = getattr(req, "prompts", None)
+        if image is not None or not prompts:
             return image
 
         raw_images = []
-        for prompt_item in req.prompts:
+        for prompt_item in prompts:
             if isinstance(prompt_item, str):
                 raw_image = None
             else:
@@ -300,8 +379,12 @@ class LTXI2VConditioningMixin:
                 raw_image = PIL.Image.open(raw_image).convert("RGB")
             raw_images.append(raw_image)
 
+        if all(raw_image is None for raw_image in raw_images) and self.unified_text_image_entry:
+            return None
         if any(raw_image is None for raw_image in raw_images) and request_inputs.latents is None:
-            raise ValueError("Image is required for LTX I2V generation.")
+            if not self.unified_text_image_entry:
+                raise ValueError("Image is required for LTX I2V generation.")
+            raise ValueError("LTX request batches cannot mix text-to-video and image-to-video prompts.")
         if len(raw_images) == 1:
             return raw_images[0]
         return raw_images or image
@@ -311,75 +394,19 @@ class LTXI2VConditioningMixin:
         request_inputs: LTXRequestInputs,
         image: Any | None = None,
     ) -> None:
-        if image is None and request_inputs.latents is None:
+        if not self.unified_text_image_entry and image is None and request_inputs.latents is None:
             raise ValueError("Provide either `image` or `latents`. Cannot leave both undefined.")
         super()._check_forward_inputs(request_inputs, image=image)
 
-    def prepare_latents(
+    def _encode_i2v_image_latents(
         self,
-        image: torch.Tensor | None = None,
-        batch_size: int = 1,
-        num_channels_latents: int = 128,
-        height: int = 512,
-        width: int = 768,
-        num_frames: int = 121,
-        noise_scale: float = 0.0,
-        dtype: torch.dtype | None = None,
-        device: torch.device | None = None,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        num_frames, height, width = latent_ops.resolve_video_latent_shape(
-            height,
-            width,
-            num_frames,
-            vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
-            vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
-        )
-        shape = (batch_size, num_channels_latents, num_frames, height, width)
-        mask_shape = (batch_size, 1, num_frames, height, width)
-
-        if latents is not None:
-            if latents.ndim == 5:
-                batch_size, _, num_frames, height, width = latents.shape
-                mask_shape = (batch_size, 1, num_frames, height, width)
-                conditioning_mask = latents.new_zeros(mask_shape)
-                conditioning_mask[:, :, 0] = 1.0
-                latents = latent_ops.normalize_latents(
-                    latents,
-                    self.vae.latents_mean,
-                    self.vae.latents_std,
-                    self.vae.config.scaling_factor,
-                )
-                latents = latent_ops.create_noised_state(
-                    latents,
-                    noise_scale * (1 - conditioning_mask),
-                    generator,
-                )
-                latents = latent_ops.pack_latents(
-                    latents,
-                    self.transformer_spatial_patch_size,
-                    self.transformer_temporal_patch_size,
-                )
-            else:
-                conditioning_mask = latents.new_zeros(mask_shape)
-                conditioning_mask[:, :, 0] = 1.0
-
-            conditioning_mask = latent_ops.pack_latents(
-                conditioning_mask,
-                self.transformer_spatial_patch_size,
-                self.transformer_temporal_patch_size,
-            ).squeeze(-1)
-            if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
-                raise ValueError(
-                    "Provided `latents` tensor has shape"
-                    f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
-                )
-            return latents.to(device=device, dtype=dtype), conditioning_mask
-
-        if image is None:
-            raise ValueError("`image` must be provided when `latents` is None.")
-
+        image: torch.Tensor,
+        *,
+        batch_size: int,
+        generator: torch.Generator | list[torch.Generator] | None,
+        dtype: torch.dtype | None,
+    ) -> torch.Tensor:
+        """Encode one conditioning frame per prompt into normalized latents."""
         image_batch_size = image.shape[0]
         if image_batch_size == 0:
             raise ValueError("`image` batch is empty.")
@@ -415,38 +442,173 @@ class LTXI2VConditioningMixin:
                 for img in image
             ]
 
-        init_latents = torch.cat(init_latents, dim=0).to(dtype)
+        init_latents = torch.cat(init_latents, dim=0).to(dtype=dtype)
         if num_videos_per_prompt > 1:
             init_latents = init_latents.repeat_interleave(num_videos_per_prompt, dim=0)
-        init_latents = latent_ops.normalize_latents(
+        return latent_ops.normalize_latents(
             init_latents,
             self.vae.latents_mean,
             self.vae.latents_std,
         )
+
+    @staticmethod
+    def _replace_first_latent_frame(
+        latents: torch.Tensor,
+        image_latents: torch.Tensor,
+    ) -> torch.Tensor:
+        expected_shape = (latents.shape[0], latents.shape[1], 1, latents.shape[3], latents.shape[4])
+        if tuple(image_latents.shape) != expected_shape:
+            raise ValueError(
+                f"Encoded LTX conditioning image has shape {tuple(image_latents.shape)}, expected {expected_shape}."
+            )
+        latents = latents.clone()
+        latents[:, :, :1] = image_latents
+        return latents
+
+    def prepare_latents(
+        self,
+        image: torch.Tensor | None = None,
+        batch_size: int = 1,
+        num_channels_latents: int = 128,
+        height: int = 512,
+        width: int = 768,
+        num_frames: int = 121,
+        noise_scale: float = 0.0,
+        dtype: torch.dtype | None = None,
+        device: torch.device | None = None,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+    ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]:
+        if image is None and self.unified_text_image_entry:
+            return super().prepare_latents(
+                batch_size=batch_size,
+                num_channels_latents=num_channels_latents,
+                height=height,
+                width=width,
+                num_frames=num_frames,
+                noise_scale=noise_scale,
+                dtype=dtype,
+                device=device,
+                generator=generator,
+                latents=latents,
+            )
+
+        num_frames, height, width = latent_ops.resolve_video_latent_shape(
+            height,
+            width,
+            num_frames,
+            vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
+            vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
+        )
+        mask_shape = (batch_size, 1, num_frames, height, width)
+
+        if latents is not None:
+            unpacked_latents = latents.ndim == 5
+            if latents.ndim == 5:
+                expected_shape = (batch_size, num_channels_latents, num_frames, height, width)
+                if tuple(latents.shape) != expected_shape:
+                    raise ValueError(
+                        f"Provided unpacked `latents` have shape {tuple(latents.shape)}, expected {expected_shape}."
+                    )
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+                latents = latent_ops.normalize_latents(
+                    latents,
+                    self.vae.latents_mean,
+                    self.vae.latents_std,
+                    self.vae.config.scaling_factor,
+                )
+                if image is not None:
+                    image_latents = self._encode_i2v_image_latents(
+                        image,
+                        batch_size=batch_size,
+                        generator=generator,
+                        dtype=dtype,
+                    )
+                    # Official distilled Stage 2 starts from the upsampled Stage 1
+                    # state, then reapplies the source image at final resolution.
+                    latents = self._replace_first_latent_frame(latents, image_latents)
+                latents = latent_ops.pack_latents(
+                    latents,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+            else:
+                conditioning_mask = latents.new_zeros(mask_shape)
+                conditioning_mask[:, :, 0] = 1.0
+
+            conditioning_mask = latent_ops.pack_latents(
+                conditioning_mask,
+                self.transformer_spatial_patch_size,
+                self.transformer_temporal_patch_size,
+            ).squeeze(-1)
+            if latents.ndim != 3 or latents.shape[:2] != conditioning_mask.shape:
+                raise ValueError(
+                    "Provided `latents` tensor has shape"
+                    f" {latents.shape}, but the expected shape is {conditioning_mask.shape + (num_channels_latents,)}."
+                )
+            if unpacked_latents:
+                # Official LTX patchifies the latent state before drawing noise.
+                # Sampling the packed shape preserves its seed-to-output contract.
+                latents = latent_ops.create_noised_state(
+                    latents,
+                    noise_scale * (1 - conditioning_mask).unsqueeze(-1),
+                    generator,
+                )
+            elif image is not None:
+                image_latents = self._encode_i2v_image_latents(
+                    image,
+                    batch_size=batch_size,
+                    generator=generator,
+                    dtype=dtype,
+                )
+                unpacked = latent_ops.unpack_latents(
+                    latents,
+                    num_frames,
+                    height,
+                    width,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+                unpacked = self._replace_first_latent_frame(unpacked, image_latents)
+                latents = latent_ops.pack_latents(
+                    unpacked,
+                    self.transformer_spatial_patch_size,
+                    self.transformer_temporal_patch_size,
+                )
+            return latents.to(device=device, dtype=dtype), conditioning_mask
+
+        if image is None:
+            raise ValueError("`image` must be provided when `latents` is None.")
+        init_latents = self._encode_i2v_image_latents(
+            image,
+            batch_size=batch_size,
+            generator=generator,
+            dtype=dtype,
+        )
         init_latents = init_latents.repeat(1, 1, num_frames, 1, 1)
 
         conditioning_mask = torch.zeros(mask_shape, device=device, dtype=dtype)
         conditioning_mask[:, :, 0] = 1.0
-        noise = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-        latents = init_latents * conditioning_mask + noise * (1 - conditioning_mask)
-
+        init_latents = latent_ops.pack_latents(
+            init_latents,
+            self.transformer_spatial_patch_size,
+            self.transformer_temporal_patch_size,
+        )
         conditioning_mask = latent_ops.pack_latents(
             conditioning_mask,
             self.transformer_spatial_patch_size,
             self.transformer_temporal_patch_size,
         ).squeeze(-1)
-        latents = latent_ops.pack_latents(
-            latents,
-            self.transformer_spatial_patch_size,
-            self.transformer_temporal_patch_size,
-        )
+        noise = randn_tensor(init_latents.shape, generator=generator, device=device, dtype=dtype)
+        latents = init_latents * conditioning_mask.unsqueeze(-1) + noise * (1 - conditioning_mask).unsqueeze(-1)
         return latents, conditioning_mask
 
     def _step_video_latents_i2v(
         self,
         noise_pred_video: torch.Tensor,
         latents: torch.Tensor,
-        timestep: torch.Tensor,
+        step_index: int,
         latent_num_frames: int,
         latent_height: int,
         latent_width: int,
@@ -469,12 +631,14 @@ class LTXI2VConditioningMixin:
         )
         noise_pred_video = noise_pred_video[:, :, 1:]
         noise_latents = latents_unpacked[:, :, 1:]
-        pred_latents = self.scheduler.step(
-            noise_pred_video,
-            timestep,
+        # Official LTX advances the unconditioned frames directly on the
+        # sigma trajectory; the first frame remains fixed conditioning.
+        pred_latents = euler_step_from_velocity(
             noise_latents,
-            return_dict=False,
-        )[0]
+            noise_pred_video,
+            self.scheduler.sigmas,
+            step_index,
+        )
         latents_unpacked = torch.cat([latents_unpacked[:, :, :1], pred_latents], dim=2)
         return latent_ops.pack_latents(
             latents_unpacked,
@@ -491,12 +655,31 @@ class LTXI2VConditioningMixin:
         noise_scale: float,
         image: Any | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        if request_inputs.latents is None:
+        if image is None and self.unified_text_image_entry:
+            return super()._prepare_video_latents_stage(
+                request_inputs,
+                prompt_context,
+                device=device,
+                noise_scale=noise_scale,
+                image=None,
+            )
+
+        if image is not None:
             if isinstance(image, torch.Tensor):
                 if image.ndim == 3:
                     image = image.unsqueeze(0)
             elif isinstance(image, list) and image and isinstance(image[0], torch.Tensor):
                 image = torch.stack(image, dim=0)
+            elif isinstance(image, PIL.Image.Image) or (
+                isinstance(image, list) and image and isinstance(image[0], PIL.Image.Image)
+            ):
+                image = _preprocess_i2v_pil_images(
+                    image,
+                    height=request_inputs.height,
+                    width=request_inputs.width,
+                    device=device,
+                    dtype=prompt_context.positive_connector_prompt_embeds.dtype,
+                )
             else:
                 image = self.video_processor.preprocess(
                     image,
@@ -509,32 +692,17 @@ class LTXI2VConditioningMixin:
             )
 
         return self.prepare_latents(
-            image,
-            prompt_context.batch_size * request_inputs.num_videos_per_prompt,
-            self.transformer.config.in_channels,
-            request_inputs.height,
-            request_inputs.width,
-            request_inputs.num_frames,
-            noise_scale,
-            torch.float32,
-            device,
-            request_inputs.generator,
-            request_inputs.latents,
-        )
-
-    def _make_video_audio_scheduler(
-        self,
-        audio_scheduler: Any,
-        latent_num_frames: int,
-        latent_height: int,
-        latent_width: int,
-    ) -> Any:
-        return I2VVideoAudioScheduler(
-            self,
-            audio_scheduler,
-            latent_num_frames,
-            latent_height,
-            latent_width,
+            image=image,
+            batch_size=prompt_context.batch_size * request_inputs.num_videos_per_prompt,
+            num_channels_latents=self.transformer.config.in_channels,
+            height=request_inputs.height,
+            width=request_inputs.width,
+            num_frames=request_inputs.num_frames,
+            noise_scale=noise_scale,
+            dtype=prompt_context.positive_connector_prompt_embeds.dtype,
+            device=device,
+            generator=request_inputs.generator,
+            latents=request_inputs.latents,
         )
 
     def _prepare_denoise_context_for_cfg(
@@ -544,6 +712,8 @@ class LTXI2VConditioningMixin:
     ) -> LTXDenoiseContext:
         denoise_ctx = super()._prepare_denoise_context_for_cfg(forward_ctx, denoise_ctx)
         if denoise_ctx.conditioning_mask is None:
+            if self.unified_text_image_entry:
+                return denoise_ctx
             raise ValueError("LTX I2V denoising requires a conditioning mask.")
 
         mask_batch = denoise_ctx.conditioning_mask.shape[0]
@@ -564,15 +734,46 @@ class LTXI2VConditioningMixin:
         ts: torch.Tensor,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
+        *,
+        video_token_count: int,
+        audio_token_count: int,
     ) -> dict[str, torch.Tensor]:
-        kwargs = super()._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx)
+        if denoise_ctx.conditioning_mask is None:
+            if self.unified_text_image_entry:
+                return super()._denoise_timestep_kwargs(
+                    ts,
+                    forward_ctx,
+                    denoise_ctx,
+                    video_token_count=video_token_count,
+                    audio_token_count=audio_token_count,
+                )
+            raise ValueError("LTX I2V denoising requires a conditioning mask.")
+        kwargs = super()._denoise_timestep_kwargs(
+            ts,
+            forward_ctx,
+            denoise_ctx,
+            video_token_count=video_token_count,
+            audio_token_count=audio_token_count,
+        )
         conditioning_mask = (
             denoise_ctx.conditioning_mask if forward_ctx.cfg_parallel_ready else denoise_ctx.conditioning_mask_for_model
         )
         if conditioning_mask is None:
             raise ValueError("LTX I2V denoising requires a conditioning mask.")
         kwargs.update(
-            timestep=ts.unsqueeze(-1) * (1 - conditioning_mask),
-            audio_timestep=ts,
+            timestep=ts.reshape(-1, 1) * (1 - conditioning_mask),
         )
         return kwargs
+
+    def _video_guidance_model_sigma(
+        self,
+        sigma: torch.Tensor,
+        denoise_ctx: LTXDenoiseContext,
+    ) -> torch.Tensor:
+        if denoise_ctx.conditioning_mask is None:
+            if self.unified_text_image_entry:
+                return super()._video_guidance_model_sigma(sigma, denoise_ctx)
+            raise ValueError("LTX I2V guidance requires a conditioning mask.")
+        # The conditioned tokens are evaluated at timestep zero by the model,
+        # so velocity-to-x0 guidance must use that same per-token sigma.
+        return sigma.reshape(-1, 1, 1) * (1 - denoise_ctx.conditioning_mask).unsqueeze(-1)

vllm_omni/diffusion/models/ltx2/ltx2_denoise.py

@@ -6,20 +6,22 @@
 from __future__ import annotations
 
 import copy
+import math
 from collections.abc import Callable, Iterable
 from dataclasses import dataclass
 from typing import TYPE_CHECKING, Any, Protocol
 
-import numpy as np
 import torch
 from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import retrieve_timesteps
 
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
-from .ltx2_latents import LTXAVState
+from .ltx2_guidance import euler_step_from_velocity
+from .ltx2_latents import LTXAVState, unpack_audio_latents, unpad_audio_latents
 
 if TYPE_CHECKING:
     from .ltx2_conditioning import LTXPromptContext
+    from .ltx2_recipes import LTXPhaseRecipe
     from .ltx2_request import LTXRequestInputs
 
 
@@ -41,7 +43,7 @@ class LTXForwardContext:
     padded_audio_num_frames: int
     timesteps: torch.Tensor
     audio_scheduler: Any
-    video_audio_scheduler: Any
+    video_audio_step_adapter: Any
 
     @property
     def batch_size(self) -> int:
@@ -71,6 +73,7 @@ class LTXPhaseResult:
     forward_context: LTXForwardContext
     video: torch.Tensor
     audio: torch.Tensor
+    audio_for_next_phase: torch.Tensor | None = None
 
 
 class LTXDenoisePipeline(Protocol):
@@ -122,61 +125,81 @@ def calculate_shift(
     return image_seq_len * slope + intercept
 
 
-class VideoAudioScheduler:
-    """Composite scheduler dispatching video and audio updates."""
+class LTXVideoAudioStepAdapter:
+    """Expose the shared LTX Euler update through the distributed scheduler API."""
 
-    def __init__(self, video_scheduler, audio_scheduler):
-        self.video_scheduler = video_scheduler
-        self.audio_scheduler = audio_scheduler
-
-    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
-        video_out = self.video_scheduler.step(
-            noise_pred[0],
-            t[0],
-            latents[0],
-            return_dict=False,
-            generator=generator,
-        )[0]
-        audio_out = self.audio_scheduler.step(
-            noise_pred[1],
-            t[1],
-            latents[1],
-            return_dict=False,
-            generator=generator,
-        )[0]
-        return ((video_out, audio_out),)
-
-
-class I2VVideoAudioScheduler:
-    """Update the unconditioned video frames and the full audio state."""
-
-    def __init__(self, pipeline, audio_scheduler, latent_num_frames, latent_height, latent_width):
-        self.video_scheduler = pipeline.scheduler
-        self.audio_scheduler = audio_scheduler
+    def __init__(
+        self,
+        pipeline: Any,
+        audio_scheduler: Any,
+        latent_num_frames: int,
+        latent_height: int,
+        latent_width: int,
+        *,
+        image_conditioned: bool,
+    ) -> None:
         self._pipeline = pipeline
+        self._audio_scheduler = audio_scheduler
         self._latent_num_frames = latent_num_frames
         self._latent_height = latent_height
         self._latent_width = latent_width
+        self._image_conditioned = image_conditioned
+        self._step_index = 0
 
     def step(self, noise_pred, t, latents, return_dict=False, generator=None):
-        video_out = self._pipeline._step_video_latents_i2v(
-            noise_pred[0],
-            latents[0],
-            t[0],
-            self._latent_num_frames,
-            self._latent_height,
-            self._latent_width,
-        )
-        audio_out = self.audio_scheduler.step(
-            noise_pred[1],
-            t[1],
+        del t, return_dict, generator
+        if self._image_conditioned:
+            video_out = self._pipeline._step_video_latents_i2v(
+                noise_pred[0],
+                latents[0],
+                self._step_index,
+                self._latent_num_frames,
+                self._latent_height,
+                self._latent_width,
+            )
+        else:
+            video_out = euler_step_from_velocity(
+                latents[0],
+                noise_pred[0],
+                self._pipeline.scheduler.sigmas,
+                self._step_index,
+            )
+        audio_out = euler_step_from_velocity(
             latents[1],
-            return_dict=False,
-            generator=generator,
-        )[0]
+            noise_pred[1],
+            self._audio_scheduler.sigmas,
+            self._step_index,
+        )
+        self._step_index += 1
         return ((video_out, audio_out),)
 
 
+def _official_ltx_sigmas(scheduler: Any, steps: int, device: torch.device) -> torch.Tensor:
+    """Build the official LTX one-stage sigma schedule in torch fp32."""
+    config = scheduler.config
+    base_anchor = config.get("base_image_seq_len", 1024)
+    max_anchor = config.get("max_image_seq_len", 4096)
+    base_shift = config.get("base_shift", 0.95)
+    max_shift = config.get("max_shift", 2.05)
+
+    sigmas = torch.linspace(1.0, 0.0, steps + 1)
+    slope = (max_shift - base_shift) / (max_anchor - base_anchor)
+    # Official LTX2 one-stage intentionally uses the max sequence anchor, so the shift stays at max_shift.
+    sigma_shift = max_anchor * slope + (base_shift - slope * base_anchor)
+    exp_shift = math.exp(sigma_shift)
+    sigmas = torch.where(sigmas != 0, exp_shift / (exp_shift + (1 / sigmas - 1)), 0)
+
+    # Official non-distilled checkpoints explicitly set this to 0.1. Missing
+    # or None means the loaded scheduler disabled terminal stretching.
+    terminal = config.get("shift_terminal")
+    if terminal is not None:
+        non_zero = sigmas != 0
+        one_minus_sigmas = 1.0 - sigmas[non_zero]
+        scale = one_minus_sigmas[-1] / (1.0 - terminal)
+        sigmas[non_zero] = 1.0 - one_minus_sigmas / scale
+    return sigmas.to(dtype=torch.float32, device=device)
+
+
 def _set_scheduler_sigmas(scheduler: Any, sigmas: torch.Tensor) -> torch.Tensor:
     sigmas = sigmas.to(torch.float32)
     timesteps = sigmas[:-1] * scheduler.config.get("num_train_timesteps", 1000)
@@ -198,16 +221,20 @@ def prepare_scheduler_stage(
     latent_num_frames: int,
     latent_height: int,
     latent_width: int,
+    use_official_sigma_schedule: bool,
+    image_conditioned: bool = False,
 ) -> tuple[Any, Any, torch.Tensor]:
     if sigmas is not None and timesteps is not None:
         raise ValueError("Only one of `sigmas` or `timesteps` may be provided.")
 
     audio_scheduler = copy.deepcopy(pipeline.scheduler)
-    video_audio_scheduler = pipeline._make_video_audio_scheduler(
+    video_audio_step_adapter = LTXVideoAudioStepAdapter(
+        pipeline,
         audio_scheduler,
         latent_num_frames,
         latent_height,
         latent_width,
+        image_conditioned=image_conditioned,
     )
     if sigmas is not None:
         scheduler_sigmas = torch.as_tensor(sigmas, dtype=torch.float32, device=device)
@@ -217,12 +244,14 @@ def prepare_scheduler_stage(
             scheduler_sigmas = torch.cat([scheduler_sigmas, scheduler_sigmas.new_zeros(1)])
         timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas)
         _set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone())
-        return audio_scheduler, video_audio_scheduler, timesteps_tensor
+        return audio_scheduler, video_audio_step_adapter, timesteps_tensor
+
+    if sigmas is None and timesteps is None and use_official_sigma_schedule:
+        scheduler_sigmas = _official_ltx_sigmas(pipeline.scheduler, request_inputs.num_inference_steps, device)
+        timesteps_tensor = _set_scheduler_sigmas(pipeline.scheduler, scheduler_sigmas)
+        _set_scheduler_sigmas(audio_scheduler, scheduler_sigmas.clone())
+        return audio_scheduler, video_audio_step_adapter, timesteps_tensor
 
-    sigmas = np.linspace(1.0, 1 / request_inputs.num_inference_steps, request_inputs.num_inference_steps)
-    # Official LTX2 one-stage omits the latent, so mu stays at max_shift across resolutions.
-    # Mu varies with latent token count only when the scheduler receives a latent, as in official HQ two-stage;
-    # explicit sigma schedules bypass mu entirely.
     mu = calculate_shift(
         pipeline.scheduler.config.get("max_image_seq_len", 4096),
         pipeline.scheduler.config.get("base_image_seq_len", 1024),
@@ -246,7 +275,7 @@ def prepare_scheduler_stage(
         sigmas=sigmas,
         mu=mu,
     )
-    return audio_scheduler, video_audio_scheduler, timesteps_tensor
+    return audio_scheduler, video_audio_step_adapter, timesteps_tensor
 
 
 def prepare_rope_coords_stage(
@@ -283,15 +312,25 @@ def build_transformer_kwargs(
     encoder_attention_mask: torch.Tensor | None,
     audio_encoder_attention_mask: torch.Tensor | None,
     ts: torch.Tensor,
+    attention_kwargs: dict[str, Any] | None = None,
 ) -> dict[str, Any]:
+    del encoder_attention_mask, audio_encoder_attention_mask
     return {
         "hidden_states": hidden_states,
         "audio_hidden_states": audio_hidden_states,
         "encoder_hidden_states": encoder_hidden_states,
         "audio_encoder_hidden_states": audio_encoder_hidden_states,
-        **pipeline._denoise_timestep_kwargs(ts, forward_ctx, denoise_ctx),
-        "encoder_attention_mask": encoder_attention_mask,
-        "audio_encoder_attention_mask": audio_encoder_attention_mask,
+        **pipeline._denoise_timestep_kwargs(
+            ts,
+            forward_ctx,
+            denoise_ctx,
+            video_token_count=hidden_states.shape[1],
+            audio_token_count=audio_hidden_states.shape[1],
+        ),
+        # This is valid only because LTX connectors replace every padding token
+        # with a learned register, making all output context tokens valid.
+        "encoder_attention_mask": None,
+        "audio_encoder_attention_mask": None,
         "num_frames": forward_ctx.latent_num_frames,
         "height": forward_ctx.latent_height,
         "width": forward_ctx.latent_width,
@@ -299,7 +338,7 @@ def build_transformer_kwargs(
         "audio_num_frames": forward_ctx.padded_audio_num_frames,
         "video_coords": denoise_ctx.video_coords,
         "audio_coords": denoise_ctx.audio_coords,
-        "attention_kwargs": forward_ctx.attention_kwargs,
+        "attention_kwargs": forward_ctx.attention_kwargs if attention_kwargs is None else attention_kwargs,
         "return_dict": False,
     }
 
@@ -317,7 +356,7 @@ def step_denoised_latents(
         (timestep, timestep),
         (denoise_ctx.latents, denoise_ctx.audio_latents),
         do_true_cfg=pipeline.do_classifier_free_guidance,
-        per_request_scheduler=forward_ctx.video_audio_scheduler,
+        per_request_scheduler=forward_ctx.video_audio_step_adapter,
     )
     return pipeline._synchronize_cfg_parallel_step_output(
         latents,
@@ -338,11 +377,12 @@ class LTXPhaseExecutor:
         sigmas: list[float] | None,
         timesteps: list[int] | None,
         attention_kwargs: dict[str, Any] | None,
+        phase_recipe: LTXPhaseRecipe,
         image: Any | None = None,
         prompt_context: LTXPromptContext | None = None,
     ) -> LTXPhaseResult:
         pipeline._check_forward_inputs(request_inputs, image=image)
-        cfg_parallel_ready = pipeline._setup_forward_runtime(request_inputs, attention_kwargs)
+        cfg_parallel_ready = pipeline._setup_forward_runtime(req, request_inputs, attention_kwargs)
         device = pipeline.device
         if prompt_context is None:
             prompt_context = pipeline._prepare_prompt_context(
@@ -372,7 +412,7 @@ class LTXPhaseExecutor:
                 noise_scale=noise_scale,
             )
         )
-        audio_scheduler, video_audio_scheduler, timesteps_tensor = prepare_scheduler_stage(
+        audio_scheduler, video_audio_step_adapter, timesteps_tensor = prepare_scheduler_stage(
             pipeline,
             request_inputs,
             device=device,
@@ -381,6 +421,8 @@ class LTXPhaseExecutor:
             latent_num_frames=latent_num_frames,
             latent_height=latent_height,
             latent_width=latent_width,
+            use_official_sigma_schedule=phase_recipe.use_official_sigma_schedule,
+            image_conditioned=conditioning_mask is not None,
         )
         forward_ctx = LTXForwardContext(
             req=req,
@@ -397,7 +439,7 @@ class LTXPhaseExecutor:
             padded_audio_num_frames=padded_audio_num_frames,
             timesteps=timesteps_tensor,
             audio_scheduler=audio_scheduler,
-            video_audio_scheduler=video_audio_scheduler,
+            video_audio_step_adapter=video_audio_step_adapter,
         )
         video_coords, audio_coords = prepare_rope_coords_stage(pipeline, forward_ctx, latents, audio_latents)
         denoise_ctx = LTXDenoiseContext(
@@ -427,4 +469,13 @@ class LTXPhaseExecutor:
             state.video,
             state.audio,
         )
-        return LTXPhaseResult(forward_context=forward_ctx, video=latents, audio=audio_latents)
+        normalized_audio = unpack_audio_latents(
+            unpad_audio_latents(state.audio, forward_ctx.original_audio_num_frames),
+            num_mel_bins=forward_ctx.latent_mel_bins,
+        )
+        return LTXPhaseResult(
+            forward_context=forward_ctx,
+            video=latents,
+            audio=audio_latents,
+            audio_for_next_phase=normalized_audio,
+        )

vllm_omni/diffusion/models/ltx2/ltx2_guidance.py

@@ -1,16 +1,16 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Guidance strategies for the LTX model family."""
+"""Official multi-modal guidance for the LTX model family."""
 
 from __future__ import annotations
 
-from typing import TYPE_CHECKING, Any, Protocol
+import math
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
 
 import torch
-from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion import rescale_noise_cfg
 
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
 from vllm_omni.diffusion.distributed.parallel_state import (
     get_cfg_group,
     get_classifier_free_guidance_rank,
@@ -22,358 +22,409 @@ if TYPE_CHECKING:
     from .ltx2_latents import LTXAVState
 
 
-class LTXGuidanceStrategy(Protocol):
-    """Model-version guidance behavior consumed by the shared pipeline."""
+@dataclass(frozen=True)
+class LTXModalityGuidance:
+    """Guidance parameters applied to one LTX output modality."""
 
-    def validate_cfg_world_size(self, cfg_world_size: int) -> None: ...
+    cfg_scale: float = 1.0
+    stg_scale: float = 0.0
+    modality_scale: float = 1.0
+    rescale_scale: float = 0.0
+    stg_blocks: tuple[int, ...] = ()
 
-    def prepare_denoise_context(
-        self,
-        pipeline: Any,
-        forward_ctx: LTXForwardContext,
-        denoise_ctx: LTXDenoiseContext,
-    ) -> LTXDenoiseContext: ...
+    @property
+    def do_cfg(self) -> bool:
+        return not math.isclose(self.cfg_scale, 1.0)
 
-    def timestep_kwargs(
-        self,
-        ts: torch.Tensor,
-        forward_ctx: LTXForwardContext,
-        denoise_ctx: LTXDenoiseContext,
-    ) -> dict[str, torch.Tensor]: ...
+    @property
+    def do_stg(self) -> bool:
+        return not math.isclose(self.stg_scale, 0.0) and bool(self.stg_blocks)
 
-    def combine_cfg_noise(
-        self,
-        pipeline: Any,
-        positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
-        negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
-        true_cfg_scale: float,
-        cfg_normalize: bool,
-        context: dict[str, Any],
-    ) -> tuple[torch.Tensor, torch.Tensor]: ...
-
-    def predict_noise(
-        self,
-        pipeline: Any,
-        index: int,
-        timestep: torch.Tensor,
-        state: LTXAVState,
-        forward_ctx: LTXForwardContext,
-        denoise_ctx: LTXDenoiseContext,
-    ) -> tuple[torch.Tensor, torch.Tensor]: ...
+    @property
+    def do_modality_guidance(self) -> bool:
+        return not math.isclose(self.modality_scale, 1.0)
 
 
-class _LTXGuidanceStrategyBase:
-    def validate_cfg_world_size(self, cfg_world_size: int) -> None:
-        del cfg_world_size
+@dataclass(frozen=True)
+class LTXGuidanceSpec:
+    """Video and audio guidance requested for one denoise phase."""
 
-    def prepare_denoise_context(
-        self,
-        pipeline: Any,
-        forward_ctx: LTXForwardContext,
-        denoise_ctx: LTXDenoiseContext,
-    ) -> LTXDenoiseContext:
-        del pipeline, forward_ctx
-        return denoise_ctx
+    video: LTXModalityGuidance = LTXModalityGuidance()
+    audio: LTXModalityGuidance = LTXModalityGuidance()
 
-    def timestep_kwargs(
-        self,
-        ts: torch.Tensor,
-        forward_ctx: LTXForwardContext,
-        denoise_ctx: LTXDenoiseContext,
-    ) -> dict[str, torch.Tensor]:
-        del forward_ctx, denoise_ctx
-        return {"timestep": ts}
+    @classmethod
+    def positive_only(cls) -> LTXGuidanceSpec:
+        return cls()
+
+    @property
+    def do_cfg(self) -> bool:
+        return self.video.do_cfg or self.audio.do_cfg
+
+    @property
+    def do_stg(self) -> bool:
+        return self.video.do_stg or self.audio.do_stg
+
+    @property
+    def do_modality_guidance(self) -> bool:
+        return self.video.do_modality_guidance or self.audio.do_modality_guidance
+
+    @property
+    def do_rescale(self) -> bool:
+        return self.video.rescale_scale != 0.0 or self.audio.rescale_scale != 0.0
 
 
-def combine_velocity_via_x0(
+@dataclass(frozen=True)
+class LTXDenoisePass:
+    """One conditioned Transformer evaluation in a guidance batch."""
+
+    name: str
+    negative_video_context: bool = False
+    negative_audio_context: bool = False
+
+
+@dataclass(frozen=True)
+class LTXGuidancePlan:
+    """Concrete Transformer passes derived from a guidance specification."""
+
+    spec: LTXGuidanceSpec
+    passes: tuple[LTXDenoisePass, ...]
+
+    @classmethod
+    def build(cls, spec: LTXGuidanceSpec) -> LTXGuidancePlan:
+        passes = [LTXDenoisePass("cond")]
+        if spec.do_cfg:
+            passes.append(
+                LTXDenoisePass(
+                    "uncond",
+                    negative_video_context=spec.video.do_cfg,
+                    negative_audio_context=spec.audio.do_cfg,
+                )
+            )
+        if spec.do_stg:
+            passes.append(LTXDenoisePass("ptb"))
+        if spec.do_modality_guidance:
+            passes.append(LTXDenoisePass("mod"))
+        return cls(spec=spec, passes=tuple(passes))
+
+    @property
+    def names(self) -> tuple[str, ...]:
+        return tuple(item.name for item in self.passes)
+
+    @property
+    def cfg_parallel_compatible(self) -> bool:
+        return self.names == ("cond", "uncond") and not self.spec.do_rescale
+
+
+def x0_from_velocity(sample: torch.Tensor, velocity: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
+    """Convert velocity to x0 using the official fp32 arithmetic order."""
+    sigma = sigma.to(torch.float32)
+    return (sample.to(torch.float32) - velocity.to(torch.float32) * sigma).to(sample.dtype)
+
+
+def velocity_from_x0(sample: torch.Tensor, x0: torch.Tensor, sigma: torch.Tensor) -> torch.Tensor:
+    """Convert x0 back to velocity using the official fp32 arithmetic order."""
+    sigma = sigma.to(torch.float32)
+    return ((sample.to(torch.float32) - x0.to(torch.float32)) / sigma).to(sample.dtype)
+
+
+def euler_step_from_velocity(
     sample: torch.Tensor,
-    positive_velocity: torch.Tensor,
-    negative_velocity: torch.Tensor,
-    sigma: torch.Tensor,
-    guidance_scale: float,
+    velocity: torch.Tensor,
+    sigmas: torch.Tensor,
+    step_index: int,
 ) -> torch.Tensor:
-    """Apply CFG in x0 space and convert the result back to velocity."""
-    x0_cond = sample - positive_velocity * sigma
-    x0_uncond = sample - negative_velocity * sigma
-    x0_guided = x0_cond + (guidance_scale - 1) * (x0_cond - x0_uncond)
-    return (sample - x0_guided) / sigma
+    sigma = sigmas[step_index].to(torch.float32)
+    sigma_next = sigmas[step_index + 1].to(torch.float32)
+    return (sample.float() + velocity.float() * (sigma_next - sigma)).to(sample.dtype)
 
 
-class LTXLegacyVelocityGuidance(_LTXGuidanceStrategyBase):
-    """Existing LTX2 Diffusers-style velocity-space CFG behavior."""
-
-    def combine_cfg_noise(
-        self,
-        pipeline: Any,
-        positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
-        negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
-        true_cfg_scale: float,
-        cfg_normalize: bool,
-        context: dict[str, Any],
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        del context
-        video_pos, audio_pos = positive_noise_pred
-        video_neg, audio_neg = negative_noise_pred
-        video_combined = CFGParallelMixin.combine_cfg_noise(
-            pipeline,
-            video_pos,
-            video_neg,
-            true_cfg_scale,
-            cfg_normalize,
-        )
-        audio_combined = CFGParallelMixin.combine_cfg_noise(
-            pipeline,
-            audio_pos,
-            audio_neg,
-            true_cfg_scale,
-            cfg_normalize,
-        )
-        if pipeline.guidance_rescale and pipeline.guidance_rescale > 0:
-            video_combined = rescale_noise_cfg(
-                video_combined,
-                video_pos,
-                guidance_rescale=pipeline.guidance_rescale,
-            )
-            audio_combined = rescale_noise_cfg(
-                audio_combined,
-                audio_pos,
-                guidance_rescale=pipeline.guidance_rescale,
-            )
-        return video_combined, audio_combined
-
-    def predict_noise(
-        self,
-        pipeline: Any,
-        index: int,
-        timestep: torch.Tensor,
-        state: LTXAVState,
-        forward_ctx: LTXForwardContext,
-        denoise_ctx: LTXDenoiseContext,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        del index
-        prompt_context = forward_ctx.prompt_context
-        video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
-        audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
-        expanded_timestep = timestep.expand(video_input.shape[0])
-        positive_kwargs = pipeline._build_transformer_kwargs(
-            forward_ctx,
-            denoise_ctx,
-            hidden_states=video_input,
-            audio_hidden_states=audio_input,
-            encoder_hidden_states=prompt_context.positive_connector_prompt_embeds,
-            audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds,
-            encoder_attention_mask=prompt_context.positive_connector_attention_mask,
-            audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask,
-            ts=expanded_timestep,
-        )
-        negative_kwargs = (
-            {
-                **positive_kwargs,
-                "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds,
-                "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds,
-                "encoder_attention_mask": prompt_context.negative_connector_attention_mask,
-                "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask,
-            }
-            if pipeline.do_classifier_free_guidance
-            else None
-        )
-        return pipeline.predict_noise_maybe_with_cfg(
-            do_true_cfg=pipeline.do_classifier_free_guidance,
-            true_cfg_scale=forward_ctx.request_inputs.guidance_scale,
-            positive_kwargs=positive_kwargs,
-            negative_kwargs=negative_kwargs,
-            cfg_normalize=False,
-        )
+def combine_guided_x0(
+    *,
+    cond: torch.Tensor,
+    uncond_text: torch.Tensor | float,
+    uncond_perturbed: torch.Tensor | float,
+    uncond_modality: torch.Tensor | float,
+    guidance: LTXModalityGuidance,
+) -> torch.Tensor:
+    dtype = cond.dtype
+    cond = cond.float()
+    uncond_text = uncond_text.float() if isinstance(uncond_text, torch.Tensor) else uncond_text
+    uncond_perturbed = uncond_perturbed.float() if isinstance(uncond_perturbed, torch.Tensor) else uncond_perturbed
+    uncond_modality = uncond_modality.float() if isinstance(uncond_modality, torch.Tensor) else uncond_modality
+    pred = cond + (guidance.cfg_scale - 1) * (cond - uncond_text)
+    if guidance.do_stg:
+        pred = pred + guidance.stg_scale * (cond - uncond_perturbed)
+    pred = pred + (guidance.modality_scale - 1) * (cond - uncond_modality)
+    if guidance.rescale_scale != 0.0:
+        # The pinned official one-stage entry runs one generated sample and
+        # uses its full-tensor standard deviation. In Omni, dimension 0 may
+        # instead contain independently scheduled requests. Reduce each item
+        # separately so its result cannot depend on unrelated co-batched work.
+        reduce_dims = tuple(range(1, pred.ndim))
+        cond_std = cond.std(dim=reduce_dims, keepdim=True)
+        pred_std = pred.std(dim=reduce_dims, keepdim=True)
+        eps = torch.finfo(pred.dtype).eps
+        factor = cond_std / pred_std.clamp_min(eps)
+        factor = torch.where(pred_std > eps, factor, torch.ones_like(factor))
+        factor = guidance.rescale_scale * factor + (1 - guidance.rescale_scale)
+        pred = pred * factor
+    return pred.to(dtype)
 
 
-class LTXOfficialX0Guidance(_LTXGuidanceStrategyBase):
-    """Official LTX x0-space CFG behavior currently used by LTX2.3."""
+def _repeat_batch(tensor: torch.Tensor, repeats: int) -> torch.Tensor:
+    return tensor.repeat((repeats,) + (1,) * (tensor.ndim - 1))
 
-    def validate_cfg_world_size(self, cfg_world_size: int) -> None:
-        if cfg_world_size not in (1, 2):
-            raise ValueError(
-                f"LTX x0 guidance supports CFG parallelism with cfg_parallel_size 1 or 2, but got {cfg_world_size}."
-            )
 
+def _mask_for_pass(
+    pass_count: int,
+    batch_size: int,
+    pass_index: int,
+    reference: torch.Tensor,
+) -> torch.Tensor:
+    mask = reference.new_ones((pass_count * batch_size, 1, 1))
+    start = pass_index * batch_size
+    mask[start : start + batch_size] = 0
+    return mask
+
+
+def build_perturbation_kwargs(
+    plan: LTXGuidancePlan,
+    batch_size: int,
+    reference: torch.Tensor,
+) -> dict[str, Any]:
+    kwargs: dict[str, Any] = {}
+    if "ptb" in plan.names:
+        mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("ptb"), reference)
+        if plan.spec.video.do_stg:
+            kwargs["video_self_attention_mask"] = mask
+            kwargs["video_self_attention_blocks"] = plan.spec.video.stg_blocks
+        if plan.spec.audio.do_stg:
+            kwargs["audio_self_attention_mask"] = mask
+            kwargs["audio_self_attention_blocks"] = plan.spec.audio.stg_blocks
+    if "mod" in plan.names:
+        mask = _mask_for_pass(len(plan.passes), batch_size, plan.names.index("mod"), reference)
+        kwargs["a2v_cross_attention_mask"] = mask
+        kwargs["v2a_cross_attention_mask"] = mask
+    return kwargs
+
+
+class LTXGuidanceExecutor:
+    """Build and execute official LTX multi-guidance Transformer passes."""
+
+    @staticmethod
+    def validate_cfg_world_size(plan: LTXGuidancePlan, cfg_world_size: int) -> None:
+        if cfg_world_size == 1:
+            return
+        if cfg_world_size != 2:
+            raise ValueError(f"LTX CFG parallelism supports cfg_parallel_size 1 or 2, got {cfg_world_size}.")
+        if not plan.cfg_parallel_compatible:
+            raise ValueError("LTX CFG parallelism only supports CFG-only guidance without rescale.")
+
+    @staticmethod
     def prepare_denoise_context(
-        self,
-        pipeline: Any,
-        forward_ctx: LTXForwardContext,
+        plan: LTXGuidancePlan,
+        cfg_parallel_ready: bool,
         denoise_ctx: LTXDenoiseContext,
     ) -> LTXDenoiseContext:
-        if pipeline.do_classifier_free_guidance and not forward_ctx.cfg_parallel_ready:
-            denoise_ctx.video_coords = denoise_ctx.video_coords.repeat(
-                (2,) + (1,) * (denoise_ctx.video_coords.ndim - 1)
-            )
-            denoise_ctx.audio_coords = denoise_ctx.audio_coords.repeat(
-                (2,) + (1,) * (denoise_ctx.audio_coords.ndim - 1)
-            )
+        if len(plan.passes) > 1 and not cfg_parallel_ready:
+            denoise_ctx.video_coords = _repeat_batch(denoise_ctx.video_coords, len(plan.passes))
+            denoise_ctx.audio_coords = _repeat_batch(denoise_ctx.audio_coords, len(plan.passes))
         return denoise_ctx
 
+    @staticmethod
     def timestep_kwargs(
-        self,
         ts: torch.Tensor,
-        forward_ctx: LTXForwardContext,
-        denoise_ctx: LTXDenoiseContext,
+        video_token_count: int,
+        audio_token_count: int,
     ) -> dict[str, torch.Tensor]:
-        del forward_ctx, denoise_ctx
-        return {"timestep": ts, "sigma": ts}
+        return {
+            "timestep": ts.reshape(-1, 1).expand(-1, video_token_count),
+            "audio_timestep": ts.reshape(-1, 1).expand(-1, audio_token_count),
+            "sigma": ts,
+            "audio_sigma": ts,
+        }
 
-    def combine_cfg_noise(
-        self,
-        pipeline: Any,
-        positive_noise_pred: tuple[torch.Tensor, torch.Tensor],
-        negative_noise_pred: tuple[torch.Tensor, torch.Tensor],
-        true_cfg_scale: float,
-        cfg_normalize: bool,
-        context: dict[str, Any],
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma")
-        if any(context.get(name) is None for name in required):
-            raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.")
-
-        video_pos, audio_pos = positive_noise_pred
-        video_neg, audio_neg = negative_noise_pred
-        video_combined = combine_velocity_via_x0(
-            context["video_latents"],
-            video_pos,
-            video_neg,
-            context["video_sigma"],
-            true_cfg_scale,
+    @staticmethod
+    def combine_cfg_velocity(
+        sample: torch.Tensor,
+        positive_velocity: torch.Tensor,
+        negative_velocity: torch.Tensor,
+        sigma: torch.Tensor,
+        guidance_scale: float,
+        *,
+        model_sigma: torch.Tensor | None = None,
+    ) -> torch.Tensor:
+        guidance = LTXModalityGuidance(cfg_scale=guidance_scale)
+        model_sigma = sigma if model_sigma is None else model_sigma
+        guided_x0 = combine_guided_x0(
+            cond=x0_from_velocity(sample, positive_velocity, model_sigma),
+            uncond_text=x0_from_velocity(sample, negative_velocity, model_sigma),
+            uncond_perturbed=0.0,
+            uncond_modality=0.0,
+            guidance=guidance,
         )
-        audio_combined = combine_velocity_via_x0(
-            context["audio_latents"],
-            audio_pos,
-            audio_neg,
-            context["audio_sigma"],
-            true_cfg_scale,
-        )
-        if cfg_normalize:
-            video_combined = pipeline.cfg_normalize_function(video_pos, video_combined)
-            audio_combined = pipeline.cfg_normalize_function(audio_pos, audio_combined)
-        return video_combined, audio_combined
+        return velocity_from_x0(sample, guided_x0, sigma)
 
     def predict_parallel_cfg(
         self,
         pipeline: Any,
-        true_cfg_scale: float,
-        positive_kwargs: dict[str, Any],
-        negative_kwargs: dict[str, Any],
-        cfg_normalize: bool = True,
-        output_slice: int | None = None,
-        **context: Any,
-    ) -> tuple[torch.Tensor, torch.Tensor]:
-        def maybe_slice(pred: tuple[torch.Tensor, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]:
-            if output_slice is None:
-                return pred
-            return pred[0][:, :output_slice], pred[1][:, :output_slice]
-
-        cfg_world_size = get_classifier_free_guidance_world_size()
-        self.validate_cfg_world_size(cfg_world_size)
-        if cfg_world_size != 2:
-            raise ValueError(f"Parallel LTX x0 guidance requires cfg_parallel_size 2, but got {cfg_world_size}.")
-
-        cfg_rank = get_classifier_free_guidance_rank()
-        branch_kwargs = positive_kwargs if cfg_rank == 0 else negative_kwargs
-        local_video_pred, local_audio_pred = maybe_slice(pipeline.predict_noise(**branch_kwargs))
-        cfg_group = get_cfg_group()
-        gathered_video = cfg_group.all_gather(local_video_pred, separate_tensors=True)
-        gathered_audio = cfg_group.all_gather(local_audio_pred, separate_tensors=True)
-        return self.combine_cfg_noise(
-            pipeline,
-            (gathered_video[0], gathered_audio[0]),
-            (gathered_video[1], gathered_audio[1]),
-            true_cfg_scale,
-            cfg_normalize,
-            context,
-        )
-
-    def predict_noise(
-        self,
-        pipeline: Any,
+        plan: LTXGuidancePlan,
         index: int,
         timestep: torch.Tensor,
         state: LTXAVState,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        prompt_context = forward_ctx.prompt_context
-        guidance_scale = forward_ctx.request_inputs.guidance_scale
-        audio_scheduler = forward_ctx.audio_scheduler
-        if forward_ctx.cfg_parallel_ready:
-            video_input = state.video.to(prompt_context.positive_connector_prompt_embeds.dtype)
-            audio_input = state.audio.to(prompt_context.positive_connector_prompt_embeds.dtype)
-            ts = timestep.expand(video_input.shape[0])
-            positive_kwargs = pipeline._build_transformer_kwargs(
-                forward_ctx,
-                denoise_ctx,
-                hidden_states=video_input,
-                audio_hidden_states=audio_input,
-                encoder_hidden_states=prompt_context.positive_connector_prompt_embeds,
-                audio_encoder_hidden_states=prompt_context.positive_connector_audio_prompt_embeds,
-                encoder_attention_mask=prompt_context.positive_connector_attention_mask,
-                audio_encoder_attention_mask=prompt_context.positive_connector_attention_mask,
-                ts=ts,
-            )
-            negative_kwargs = {
-                **positive_kwargs,
-                "encoder_hidden_states": prompt_context.negative_connector_prompt_embeds,
-                "audio_encoder_hidden_states": prompt_context.negative_connector_audio_prompt_embeds,
-                "encoder_attention_mask": prompt_context.negative_connector_attention_mask,
-                "audio_encoder_attention_mask": prompt_context.negative_connector_attention_mask,
-            }
-            return self.predict_parallel_cfg(
-                pipeline,
-                true_cfg_scale=guidance_scale,
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-                cfg_normalize=False,
-                video_latents=state.video,
-                audio_latents=state.audio,
-                video_sigma=pipeline.scheduler.sigmas[index],
-                audio_sigma=audio_scheduler.sigmas[index],
-            )
-
-        video_input = torch.cat([state.video] * 2) if pipeline.do_classifier_free_guidance else state.video
-        video_input = video_input.to(prompt_context.connector_prompt_embeds.dtype)
-        audio_input = torch.cat([state.audio] * 2) if pipeline.do_classifier_free_guidance else state.audio
-        audio_input = audio_input.to(prompt_context.connector_prompt_embeds.dtype)
-        ts = timestep.expand(video_input.shape[0])
-        transformer_kwargs = pipeline._build_transformer_kwargs(
+        self.validate_cfg_world_size(plan, get_classifier_free_guidance_world_size())
+        prompt = forward_ctx.prompt_context
+        negative = get_classifier_free_guidance_rank() == 1
+        video_context = (
+            prompt.negative_connector_prompt_embeds
+            if negative and plan.spec.video.do_cfg
+            else prompt.positive_connector_prompt_embeds
+        )
+        audio_context = (
+            prompt.negative_connector_audio_prompt_embeds
+            if negative and plan.spec.audio.do_cfg
+            else prompt.positive_connector_audio_prompt_embeds
+        )
+        if video_context is None or audio_context is None:
+            raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
+        video_input = state.video.to(prompt.positive_connector_prompt_embeds.dtype)
+        audio_input = state.audio.to(prompt.positive_connector_audio_prompt_embeds.dtype)
+        kwargs = pipeline._build_transformer_kwargs(
             forward_ctx,
             denoise_ctx,
             hidden_states=video_input,
             audio_hidden_states=audio_input,
-            encoder_hidden_states=prompt_context.connector_prompt_embeds,
-            audio_encoder_hidden_states=prompt_context.connector_audio_prompt_embeds,
-            encoder_attention_mask=prompt_context.connector_attention_mask,
-            audio_encoder_attention_mask=prompt_context.connector_attention_mask,
-            ts=ts,
+            encoder_hidden_states=video_context,
+            audio_encoder_hidden_states=audio_context,
+            encoder_attention_mask=None,
+            audio_encoder_attention_mask=None,
+            ts=timestep.expand(video_input.shape[0]),
         )
-        with pipeline._transformer_cache_context("cond_uncond"):
-            video_pred, audio_pred = pipeline.transformer(**transformer_kwargs)
-        video_pred = video_pred.float()
-        audio_pred = audio_pred.float()
-
-        if pipeline.do_classifier_free_guidance:
-            video_uncond, video_cond = video_pred.chunk(2)
-            video_pred = combine_velocity_via_x0(
+        local_video, local_audio = pipeline.predict_noise(**kwargs)
+        group = get_cfg_group()
+        video = group.all_gather(local_video, separate_tensors=True)
+        audio = group.all_gather(local_audio, separate_tensors=True)
+        video_sigma = pipeline.scheduler.sigmas[index]
+        return (
+            self.combine_cfg_velocity(
                 state.video,
-                video_cond,
-                video_uncond,
-                pipeline.scheduler.sigmas[index],
-                guidance_scale,
-            )
-            audio_uncond, audio_cond = audio_pred.chunk(2)
-            audio_pred = combine_velocity_via_x0(
+                video[0],
+                video[1],
+                video_sigma,
+                plan.spec.video.cfg_scale,
+                model_sigma=pipeline._video_guidance_model_sigma(video_sigma, denoise_ctx),
+            ),
+            self.combine_cfg_velocity(
                 state.audio,
-                audio_cond,
-                audio_uncond,
-                audio_scheduler.sigmas[index],
-                guidance_scale,
+                audio[0],
+                audio[1],
+                forward_ctx.audio_scheduler.sigmas[index],
+                plan.spec.audio.cfg_scale,
+            ),
+        )
+
+    def predict_noise(
+        self,
+        pipeline: Any,
+        plan: LTXGuidancePlan,
+        index: int,
+        timestep: torch.Tensor,
+        state: LTXAVState,
+        forward_ctx: LTXForwardContext,
+        denoise_ctx: LTXDenoiseContext,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if forward_ctx.cfg_parallel_ready:
+            return self.predict_parallel_cfg(pipeline, plan, index, timestep, state, forward_ctx, denoise_ctx)
+
+        prompt = forward_ctx.prompt_context
+        video_contexts: list[torch.Tensor] = []
+        audio_contexts: list[torch.Tensor] = []
+        for denoise_pass in plan.passes:
+            video_context = (
+                prompt.negative_connector_prompt_embeds
+                if denoise_pass.negative_video_context
+                else prompt.positive_connector_prompt_embeds
             )
-        return video_pred, audio_pred
+            audio_context = (
+                prompt.negative_connector_audio_prompt_embeds
+                if denoise_pass.negative_audio_context
+                else prompt.positive_connector_audio_prompt_embeds
+            )
+            if video_context is None or audio_context is None:
+                raise ValueError("Negative prompt context is required when LTX CFG is enabled.")
+            video_contexts.append(video_context)
+            audio_contexts.append(audio_context)
+
+        pass_count = len(plan.passes)
+        video_input = _repeat_batch(state.video, pass_count).to(prompt.positive_connector_prompt_embeds.dtype)
+        audio_input = _repeat_batch(state.audio, pass_count).to(prompt.positive_connector_audio_prompt_embeds.dtype)
+        attention_kwargs = dict(forward_ctx.attention_kwargs or {})
+        perturbations = build_perturbation_kwargs(plan, state.video.shape[0], video_input)
+        if perturbations:
+            attention_kwargs["ltx_perturbation_kwargs"] = perturbations
+        kwargs = pipeline._build_transformer_kwargs(
+            forward_ctx,
+            denoise_ctx,
+            hidden_states=video_input,
+            audio_hidden_states=audio_input,
+            encoder_hidden_states=torch.cat(video_contexts),
+            audio_encoder_hidden_states=torch.cat(audio_contexts),
+            encoder_attention_mask=None,
+            audio_encoder_attention_mask=None,
+            ts=timestep.expand(video_input.shape[0]),
+            attention_kwargs=attention_kwargs,
+        )
+        with pipeline._transformer_cache_context("guided"):
+            video_velocity, audio_velocity = pipeline.transformer(**kwargs)
+        video_splits = dict(zip(plan.names, video_velocity.chunk(pass_count), strict=True))
+        audio_splits = dict(zip(plan.names, audio_velocity.chunk(pass_count), strict=True))
+
+        def guide_modality(
+            sample: torch.Tensor,
+            splits: dict[str, torch.Tensor],
+            sigma: torch.Tensor,
+            guidance: LTXModalityGuidance,
+            *,
+            model_sigma: torch.Tensor | None = None,
+        ) -> torch.Tensor:
+            model_sigma = sigma if model_sigma is None else model_sigma
+            # Official guidance reduces contiguous BSC tensors. The fused
+            # transformer output may be channel-major after splitting, which
+            # changes fp32 reduction order and can cross bf16 rounding bounds.
+            x0 = {name: x0_from_velocity(sample, value, model_sigma).contiguous() for name, value in splits.items()}
+            guided = combine_guided_x0(
+                cond=x0["cond"],
+                uncond_text=x0.get("uncond", 0.0),
+                uncond_perturbed=x0.get("ptb", 0.0),
+                uncond_modality=x0.get("mod", 0.0),
+                guidance=guidance,
+            )
+            return velocity_from_x0(sample, guided, sigma)
+
+        return (
+            guide_modality(
+                state.video,
+                video_splits,
+                pipeline.scheduler.sigmas[index],
+                plan.spec.video,
+                model_sigma=pipeline._video_guidance_model_sigma(
+                    pipeline.scheduler.sigmas[index],
+                    denoise_ctx,
+                ),
+            ),
+            guide_modality(
+                state.audio,
+                audio_splits,
+                forward_ctx.audio_scheduler.sigmas[index],
+                plan.spec.audio,
+            ),
+        )
 
 
-LTX_LEGACY_VELOCITY_GUIDANCE = LTXLegacyVelocityGuidance()
-LTX_OFFICIAL_X0_GUIDANCE = LTXOfficialX0Guidance()
+LTX_GUIDANCE_EXECUTOR = LTXGuidanceExecutor()
+
+# Kept as a narrow compatibility helper for CFG-parallel tests and extensions.
+combine_velocity_via_x0 = LTXGuidanceExecutor.combine_cfg_velocity

vllm_omni/diffusion/models/ltx2/ltx2_latents.py

@@ -120,45 +120,14 @@ def create_noised_state(
     return noise_scale * noise + (1 - noise_scale) * latents
 
 
-def pack_audio_latents(
-    latents: torch.Tensor,
-    patch_size: int | None = None,
-    patch_size_t: int | None = None,
-) -> torch.Tensor:
-    if patch_size is not None and patch_size_t is not None:
-        batch_size, _, latent_length, latent_mel_bins = latents.shape
-        post_patch_latent_length = latent_length / patch_size_t
-        post_patch_mel_bins = latent_mel_bins / patch_size
-        latents = latents.reshape(
-            batch_size,
-            -1,
-            post_patch_latent_length,
-            patch_size_t,
-            post_patch_mel_bins,
-            patch_size,
-        )
-        return latents.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2)
+def pack_audio_latents(latents: torch.Tensor) -> torch.Tensor:
     return latents.transpose(1, 2).flatten(2, 3)
 
 
 def unpack_audio_latents(
     latents: torch.Tensor,
-    latent_length: int,
     num_mel_bins: int,
-    patch_size: int | None = None,
-    patch_size_t: int | None = None,
 ) -> torch.Tensor:
-    if patch_size is not None and patch_size_t is not None:
-        batch_size = latents.size(0)
-        latents = latents.reshape(
-            batch_size,
-            latent_length,
-            num_mel_bins,
-            -1,
-            patch_size_t,
-            patch_size,
-        )
-        return latents.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3)
     return latents.unflatten(2, (-1, num_mel_bins)).transpose(1, 2)
 
 
@@ -224,18 +193,19 @@ def prepare_video_latents(
         vae_spatial_compression_ratio=pipeline.vae_spatial_compression_ratio,
         vae_temporal_compression_ratio=pipeline.vae_temporal_compression_ratio,
     )
-    shape = (batch_size, num_channels_latents, num_frames, height, width)
+    spatial_patch_size = pipeline.transformer_spatial_patch_size
+    temporal_patch_size = pipeline.transformer_temporal_patch_size
+    shape = (
+        batch_size,
+        (num_frames // temporal_patch_size) * (height // spatial_patch_size) * (width // spatial_patch_size),
+        num_channels_latents * temporal_patch_size * spatial_patch_size * spatial_patch_size,
+    )
     if isinstance(generator, list) and len(generator) != batch_size:
         raise ValueError(
             f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
             f" size of {batch_size}. Make sure the batch size matches the length of the generators."
         )
-    latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-    return pack_latents(
-        latents,
-        pipeline.transformer_spatial_patch_size,
-        pipeline.transformer_temporal_patch_size,
-    )
+    return randn_tensor(shape, generator=generator, device=device, dtype=dtype)
 
 
 def prepare_audio_latents(
@@ -249,6 +219,7 @@ def prepare_audio_latents(
     device: torch.device | None = None,
     generator: torch.Generator | list[torch.Generator] | None = None,
     latents: torch.Tensor | None = None,
+    latents_normalized: bool = False,
 ) -> tuple[torch.Tensor, int, int]:
     original_latent_length = audio_latent_length
     latent_mel_bins = num_mel_bins // pipeline.audio_vae_mel_compression_ratio
@@ -260,11 +231,12 @@ def prepare_audio_latents(
             latents = pack_audio_latents(latents)
         if latents.ndim != 3:
             raise ValueError(f"Provided `latents` has shape {latents.shape}, expected [batch, seq, features].")
-        latents = normalize_audio_latents(
-            latents,
-            pipeline.audio_vae.latents_mean,
-            pipeline.audio_vae.latents_std,
-        )
+        if not latents_normalized:
+            latents = normalize_audio_latents(
+                latents,
+                pipeline.audio_vae.latents_mean,
+                pipeline.audio_vae.latents_std,
+            )
         latents = create_noised_state(latents, noise_scale, generator)
 
         if latents.shape[1] not in {original_latent_length, padded_latent_length}:
@@ -283,11 +255,11 @@ def prepare_audio_latents(
             latents = torch.cat([latents, padding], dim=1)
         return latents.to(device=device, dtype=dtype), original_latent_length, padded_latent_length
 
-    shape = (batch_size, num_channels_latents, padded_latent_length, latent_mel_bins)
+    shape = (batch_size, padded_latent_length, num_channels_latents * latent_mel_bins)
     if isinstance(generator, list) and len(generator) != batch_size:
         raise ValueError(
             f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
             f" size of {batch_size}. Make sure the batch size matches the length of the generators."
         )
     latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype)
-    return pack_audio_latents(latents), original_latent_length, padded_latent_length
+    return latents, original_latent_length, padded_latent_length

vllm_omni/diffusion/models/ltx2/ltx2_recipes.py

@@ -1,20 +1,175 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""User-visible one-stage defaults for the LTX model family."""
+"""Declarative execution recipes for the LTX model family."""
 
 from dataclasses import dataclass
+from typing import Literal
+
+from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance
+
+LTX_DEFAULT_NEGATIVE_PROMPT = (
+    "blurry, out of focus, overexposed, underexposed, low contrast, washed out colors, excessive noise, "
+    "grainy texture, poor lighting, flickering, motion blur, distorted proportions, unnatural skin tones, "
+    "deformed facial features, asymmetrical face, missing facial features, extra limbs, disfigured hands, "
+    "wrong hand count, artifacts around text, inconsistent perspective, camera shake, incorrect depth of "
+    "field, background too sharp, background clutter, distracting reflections, harsh shadows, inconsistent "
+    "lighting direction, color banding, cartoonish rendering, 3D CGI look, unrealistic materials, uncanny "
+    "valley effect, incorrect ethnicity, wrong gender, exaggerated expressions, wrong gaze direction, "
+    "mismatched lip sync, silent or muted audio, distorted voice, robotic voice, echo, background noise, "
+    "off-sync audio, incorrect dialogue, added dialogue, repetitive speech, jittery movement, awkward "
+    "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, "
+    "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts."
+)
+
+LTX_DISTILLED_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0)
+LTX_STAGE_2_DISTILLED_SIGMAS = (0.909375, 0.725, 0.421875, 0.0)
 
 
 @dataclass(frozen=True)
-class LTXOneStageRecipe:
+class LTXPhaseRecipe:
+    """One denoise phase and the transition used to construct its input."""
+
+    name: str
+    guidance: LTXGuidanceSpec
+    spatial_downscale: int = 1
+    sigmas: tuple[float, ...] | None = None
+    noise_scale: float = 0.0
+    input_transform: Literal["initial", "spatial_upsample"] = "initial"
+    allow_guidance_override: bool = True
+    use_official_sigma_schedule: bool = True
+
+    def __post_init__(self) -> None:
+        if self.spatial_downscale < 1:
+            raise ValueError("LTX phase spatial_downscale must be positive.")
+        if self.sigmas is not None and len(self.sigmas) < 2:
+            raise ValueError("An explicit LTX sigma schedule must contain at least two denoise sigmas.")
+
+    @property
+    def num_inference_steps(self) -> int | None:
+        return None if self.sigmas is None else len(self.sigmas) - 1
+
+
+@dataclass(frozen=True)
+class LTXPipelineRecipe:
+    """Request defaults, ordered phases, output routing, and request capabilities."""
+
+    phases: tuple[LTXPhaseRecipe, ...]
     height: int = 512
     width: int = 768
     num_frames: int = 121
     frame_rate: float = 24.0
     num_inference_steps: int = 40
-    guidance_scale: float = 4.0
+    negative_prompt: str = LTX_DEFAULT_NEGATIVE_PROMPT
+    video_output_phase: int = -1
+    audio_output_phase: int = -1
+    allow_request_sigmas: bool = True
+    allow_request_latents: bool = True
+    allow_negative_prompt: bool = True
+    fixed_num_inference_steps: bool = False
+
+    def __post_init__(self) -> None:
+        if not self.phases:
+            raise ValueError("An LTX pipeline recipe must contain at least one phase.")
+        if self.phases[0].input_transform != "initial":
+            raise ValueError("The first LTX phase must use the initial request input.")
+        phase_count = len(self.phases)
+        for output_phase in (self.video_output_phase, self.audio_output_phase):
+            if not -phase_count <= output_phase < phase_count:
+                raise ValueError(f"LTX output phase {output_phase} is outside a {phase_count}-phase recipe.")
+
+    @property
+    def request_guidance(self) -> LTXGuidanceSpec:
+        return self.phases[0].guidance
+
+    @property
+    def max_spatial_downscale(self) -> int:
+        return max(phase.spatial_downscale for phase in self.phases)
 
 
-LTX2_ONE_STAGE_RECIPE = LTXOneStageRecipe()
-LTX23_ONE_STAGE_RECIPE = LTXOneStageRecipe()
+def _official_guidance(stg_block: int) -> LTXGuidanceSpec:
+    return LTXGuidanceSpec(
+        video=LTXModalityGuidance(
+            cfg_scale=3.0,
+            stg_scale=1.0,
+            modality_scale=3.0,
+            rescale_scale=0.7,
+            stg_blocks=(stg_block,),
+        ),
+        audio=LTXModalityGuidance(
+            cfg_scale=7.0,
+            stg_scale=1.0,
+            modality_scale=3.0,
+            rescale_scale=0.7,
+            stg_blocks=(stg_block,),
+        ),
+    )
+
+
+LTX2_ONE_STAGE_RECIPE = LTXPipelineRecipe(
+    phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(29)),),
+)
+LTX23_ONE_STAGE_RECIPE = LTXPipelineRecipe(
+    num_inference_steps=30,
+    phases=(LTXPhaseRecipe(name="generate", guidance=_official_guidance(28)),),
+)
+LTX_POSITIVE_ONLY_RECIPE = LTXPipelineRecipe(
+    phases=(
+        LTXPhaseRecipe(
+            name="generate",
+            guidance=LTXGuidanceSpec.positive_only(),
+            use_official_sigma_schedule=False,
+        ),
+    ),
+)
+LTX2_DISTILLED_TWO_STAGE_RECIPE = LTXPipelineRecipe(
+    # Official distilled arguments describe the final output. Stage 1 applies
+    # spatial_downscale=2 and therefore runs at 512x768.
+    height=1024,
+    width=1536,
+    num_inference_steps=len(LTX_DISTILLED_SIGMAS) - 1,
+    negative_prompt="",
+    phases=(
+        LTXPhaseRecipe(
+            name="generate_lowres",
+            guidance=LTXGuidanceSpec.positive_only(),
+            spatial_downscale=2,
+            sigmas=LTX_DISTILLED_SIGMAS,
+            noise_scale=1.0,
+            allow_guidance_override=False,
+            use_official_sigma_schedule=False,
+        ),
+        LTXPhaseRecipe(
+            name="refine",
+            guidance=LTXGuidanceSpec.positive_only(),
+            sigmas=LTX_STAGE_2_DISTILLED_SIGMAS,
+            noise_scale=LTX_STAGE_2_DISTILLED_SIGMAS[0],
+            input_transform="spatial_upsample",
+            allow_guidance_override=False,
+            use_official_sigma_schedule=False,
+        ),
+    ),
+    video_output_phase=1,
+    audio_output_phase=1,
+    allow_request_sigmas=False,
+    allow_request_latents=False,
+    allow_negative_prompt=False,
+    fixed_num_inference_steps=True,
+)
+
+
+_PIPELINE_RECIPES: dict[tuple[str, str], LTXPipelineRecipe] = {
+    ("one_stage", "2"): LTX2_ONE_STAGE_RECIPE,
+    ("one_stage", "2.3"): LTX23_ONE_STAGE_RECIPE,
+    ("distilled_two_stage", "2"): LTX2_DISTILLED_TWO_STAGE_RECIPE,
+    ("dmd2", "2"): LTX_POSITIVE_ONLY_RECIPE,
+    ("dmd2", "2.3"): LTX_POSITIVE_ONLY_RECIPE,
+}
+
+
+def resolve_ltx_pipeline_recipe(pipeline_kind: str, model_version: str) -> LTXPipelineRecipe:
+    """Resolve execution independently from component loading."""
+    try:
+        return _PIPELINE_RECIPES[(pipeline_kind, model_version)]
+    except KeyError as exc:
+        raise ValueError(f"Unsupported LTX pipeline kind/version: {pipeline_kind!r}/{model_version!r}.") from exc

vllm_omni/diffusion/models/ltx2/ltx2_request.py

@@ -1,17 +1,21 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Request normalization shared by LTX pipeline variants."""
+"""Request normalization and validation shared by LTX pipeline variants."""
 
 from __future__ import annotations
 
-from dataclasses import dataclass
+from dataclasses import dataclass, replace
 from typing import Any
 
 import torch
 
+from vllm_omni.diffusion.data import DiffusionOutput
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
+from .ltx2_guidance import LTXGuidanceSpec, LTXModalityGuidance
+from .ltx2_recipes import LTXPipelineRecipe
+
 
 @dataclass
 class LTXRequestInputs:
@@ -24,8 +28,7 @@ class LTXRequestInputs:
     num_frames: int
     frame_rate: float
     num_inference_steps: int
-    guidance_scale: float
-    guidance_rescale: float
+    guidance: LTXGuidanceSpec
     num_videos_per_prompt: int
     generator: torch.Generator | list[torch.Generator] | None
     latents: torch.Tensor | None
@@ -38,6 +41,65 @@ class LTXRequestInputs:
     decode_noise_scale: float | list[float] | None
     output_type: str
     max_sequence_length: int
+    audio_latents_normalized: bool = False
+
+    @property
+    def guidance_scale(self) -> float:
+        """Compatibility view for callers that only understand video CFG."""
+        return self.guidance.video.cfg_scale
+
+
+def validate_pipeline_request(
+    request_inputs: LTXRequestInputs,
+    *,
+    pipeline_recipe: LTXPipelineRecipe,
+    vae_spatial_compression_ratio: int,
+    vae_temporal_compression_ratio: int,
+    pipeline_name: str,
+    request_sigmas: list[float] | None = None,
+) -> None:
+    """Validate the request capabilities declared by an execution recipe."""
+    alignment = vae_spatial_compression_ratio * pipeline_recipe.max_spatial_downscale
+    if request_inputs.height % alignment != 0 or request_inputs.width % alignment != 0:
+        raise ValueError(
+            f"{pipeline_name} resolution must be divisible by {alignment}, got "
+            f"{request_inputs.width}x{request_inputs.height}."
+        )
+
+    if request_inputs.num_frames < 1 or (request_inputs.num_frames - 1) % vae_temporal_compression_ratio != 0:
+        raise ValueError(
+            f"{pipeline_name} `num_frames` must be {vae_temporal_compression_ratio} * k + 1, "
+            f"got {request_inputs.num_frames}."
+        )
+
+    if not pipeline_recipe.allow_request_latents and (
+        request_inputs.latents is not None or request_inputs.audio_latents is not None
+    ):
+        raise ValueError(f"{pipeline_name} does not accept request-provided video or audio latents.")
+
+    if not pipeline_recipe.allow_request_sigmas and request_sigmas is not None:
+        raise ValueError(f"{pipeline_name} uses fixed phase sigma schedules.")
+
+    if pipeline_recipe.fixed_num_inference_steps and (
+        request_inputs.num_inference_steps != pipeline_recipe.num_inference_steps
+    ):
+        raise ValueError(
+            f"{pipeline_name} uses {pipeline_recipe.num_inference_steps} fixed Stage 1 denoise steps; "
+            f"got {request_inputs.num_inference_steps}."
+        )
+
+    negative_prompt = request_inputs.negative_prompt
+    has_negative_prompt = (
+        any(bool(value) for value in negative_prompt) if isinstance(negative_prompt, list) else bool(negative_prompt)
+    )
+    if not pipeline_recipe.allow_negative_prompt and (
+        has_negative_prompt or request_inputs.negative_prompt_embeds is not None
+    ):
+        raise ValueError(f"{pipeline_name} uses positive-only guidance and does not accept negative conditioning.")
+
+    first_phase = pipeline_recipe.phases[0]
+    if not first_phase.allow_guidance_override and request_inputs.guidance != first_phase.guidance:
+        raise ValueError(f"{pipeline_name} uses fixed positive-only guidance and cannot override it.")
 
 
 def _unwrap_request_tensor(value: Any) -> Any:
@@ -65,11 +127,187 @@ def _get_audio_latents_from_sampling(sampling: Any) -> torch.Tensor | None:
     return sampling.extra_args.get("audio_latents")
 
 
+def _get_extra_arg(sampling: Any, *names: str) -> Any:
+    extra_args = sampling.extra_args or {}
+    for name in names:
+        value = extra_args.get(name)
+        if value is not None:
+            return value
+    return None
+
+
+def _normalize_stg_blocks(value: Any, default: tuple[int, ...]) -> tuple[int, ...]:
+    if value is None:
+        return default
+    if isinstance(value, int):
+        return (value,)
+    if not isinstance(value, (list, tuple)):
+        raise ValueError(f"STG blocks must be an integer or a list of integers, got {value!r}.")
+    return tuple(int(item) for item in value)
+
+
+def _resolve_negative_prompts(
+    prompts: list[Any],
+    fallback: str | list[str],
+) -> str | list[str]:
+    if not prompts:
+        return fallback
+    if isinstance(fallback, list) and len(fallback) != len(prompts):
+        raise ValueError(
+            f"`negative_prompt` has batch size {len(fallback)}, but the request batch has size {len(prompts)}."
+        )
+
+    resolved = []
+    for index, item in enumerate(prompts):
+        request_value = None if isinstance(item, str) else item.get("negative_prompt")
+        if request_value is not None:
+            resolved.append(request_value)
+        else:
+            resolved.append(fallback[index] if isinstance(fallback, list) else fallback)
+    return resolved
+
+
+def _resolve_modality_guidance(
+    sampling: Any,
+    modality: str,
+    default: LTXModalityGuidance,
+) -> LTXModalityGuidance:
+    cfg_aliases = (f"{modality}_cfg_scale", f"{modality}_cfg_guidance_scale")
+    stg_aliases = (f"{modality}_stg_scale", f"{modality}_stg_guidance_scale")
+    modality_aliases = (
+        ("video_modality_scale", "a2v_guidance_scale")
+        if modality == "video"
+        else ("audio_modality_scale", "v2a_guidance_scale")
+    )
+    cfg_scale = _get_extra_arg(sampling, *cfg_aliases)
+    stg_scale = _get_extra_arg(sampling, *stg_aliases)
+    modality_scale = _get_extra_arg(sampling, *modality_aliases)
+    rescale_scale = _get_extra_arg(sampling, f"{modality}_rescale_scale")
+    stg_blocks = _get_extra_arg(sampling, f"{modality}_stg_blocks")
+    return replace(
+        default,
+        cfg_scale=default.cfg_scale if cfg_scale is None else float(cfg_scale),
+        stg_scale=default.stg_scale if stg_scale is None else float(stg_scale),
+        modality_scale=default.modality_scale if modality_scale is None else float(modality_scale),
+        rescale_scale=default.rescale_scale if rescale_scale is None else float(rescale_scale),
+        stg_blocks=_normalize_stg_blocks(stg_blocks, default.stg_blocks),
+    )
+
+
+def _resolve_guidance_spec(
+    sampling: Any,
+    default: LTXGuidanceSpec,
+    *,
+    guidance_scale: float | None,
+) -> LTXGuidanceSpec:
+    video = _resolve_modality_guidance(sampling, "video", default.video)
+    audio = _resolve_modality_guidance(sampling, "audio", default.audio)
+
+    common_cfg = sampling.guidance_scale if sampling.guidance_scale_provided else guidance_scale
+    if common_cfg is not None:
+        video = replace(video, cfg_scale=float(common_cfg))
+        audio = replace(audio, cfg_scale=float(common_cfg))
+
+    return LTXGuidanceSpec(video=video, audio=audio)
+
+
 class LTXRequestMixin:
     """Normalize serving requests without coupling them to a model version."""
 
     supports_request_batch = False
-    supports_guidance_rescale = False
+
+    @staticmethod
+    def _validate_forward_options(
+        *,
+        timesteps: list[int] | None,
+        guidance_rescale: float | None,
+        noise_scale: float,
+        return_dict: bool,
+        attention_kwargs: dict[str, Any] | None,
+    ) -> None:
+        if timesteps is not None:
+            raise ValueError("LTX does not support `timesteps`; use `sigmas` for a custom schedule.")
+        if guidance_rescale not in (None, 0.0):
+            raise ValueError(
+                "LTX does not support the common `guidance_rescale`; use `video_rescale_scale` and "
+                "`audio_rescale_scale`."
+            )
+        if noise_scale != 0.0:
+            raise ValueError("LTX only supports `noise_scale=0.0`.")
+        if return_dict is not True:
+            raise ValueError("LTX only supports `return_dict=True`.")
+        if attention_kwargs is not None:
+            raise ValueError("LTX does not support public `attention_kwargs`.")
+
+    @torch.no_grad()
+    def forward(
+        self,
+        req: DiffusionRequestBatch,
+        *,
+        image: Any | None = None,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_frames: int | None = None,
+        frame_rate: float | None = None,
+        num_inference_steps: int | None = None,
+        sigmas: list[float] | None = None,
+        timesteps: list[int] | None = None,
+        guidance_scale: float | None = None,
+        guidance_rescale: float | None = None,
+        noise_scale: float = 0.0,
+        num_videos_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        audio_latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        negative_prompt_attention_mask: torch.Tensor | None = None,
+        decode_timestep: float | list[float] = 0.0,
+        decode_noise_scale: float | list[float] | None = None,
+        output_type: str = "np",
+        return_dict: bool = True,
+        attention_kwargs: dict[str, Any] | None = None,
+        max_sequence_length: int | None = None,
+    ) -> DiffusionOutput | list[DiffusionOutput]:
+        """Validate and normalize the public LTX request surface."""
+        self._validate_forward_options(
+            timesteps=timesteps,
+            guidance_rescale=guidance_rescale,
+            noise_scale=noise_scale,
+            return_dict=return_dict,
+            attention_kwargs=attention_kwargs,
+        )
+        if image is not None and not self.support_image_input:
+            raise ValueError(f"{self.__class__.__name__} does not support `image` input.")
+
+        return self._forward_request(
+            req,
+            image=image,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_inference_steps=num_inference_steps,
+            sigmas=sigmas,
+            guidance_scale=guidance_scale,
+            num_videos_per_prompt=num_videos_per_prompt,
+            generator=generator,
+            latents=latents,
+            audio_latents=audio_latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+            output_type=output_type,
+            max_sequence_length=max_sequence_length,
+        )
 
     @staticmethod
     def _resolve_request_sigmas(
@@ -128,8 +366,7 @@ class LTXRequestMixin:
         num_frames: int | None,
         frame_rate: float | None,
         num_inference_steps: int | None,
-        timesteps: list[int] | None,
-        guidance_scale: float,
+        guidance_scale: float | None,
         num_videos_per_prompt: int | None,
         generator: torch.Generator | list[torch.Generator] | None,
         latents: torch.Tensor | None,
@@ -142,38 +379,55 @@ class LTXRequestMixin:
         decode_noise_scale: float | list[float] | None,
         output_type: str,
         max_sequence_length: int | None,
-        guidance_rescale: float = 0.0,
     ) -> LTXRequestInputs:
         sampling_params_list = req.sampling_params_list
-        sampling = sampling_params_list[0]
-        prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
-        if all(isinstance(item, str) or item.get("negative_prompt") is None for item in req.prompts):
-            negative_prompt = None
-        elif req.prompts:
-            negative_prompt = [
-                "" if isinstance(item, str) else (item.get("negative_prompt") or "") for item in req.prompts
-            ]
+        for sampling in sampling_params_list:
+            if sampling.timesteps is not None:
+                raise ValueError("LTX does not support `timesteps`; use `sigmas` for a custom schedule.")
+            if sampling.guidance_rescale:
+                raise ValueError(
+                    "LTX does not support the common `guidance_rescale`; use `video_rescale_scale` and "
+                    "`audio_rescale_scale`."
+                )
+            if _get_extra_arg(sampling, "flow_shift") is not None:
+                raise ValueError("LTX does not support `flow_shift`; use `sigmas` for a custom schedule.")
 
-        height = sampling.height or height or self.one_stage_recipe.height
-        width = sampling.width or width or self.one_stage_recipe.width
-        num_frames = sampling.num_frames or num_frames or self.one_stage_recipe.num_frames
-        frame_rate = sampling.resolved_frame_rate or frame_rate or self.one_stage_recipe.frame_rate
-        num_inference_steps = (
-            sampling.num_inference_steps or num_inference_steps or self.one_stage_recipe.num_inference_steps
+        sampling = sampling_params_list[0]
+        is_dummy_run = req.is_dummy_run()
+        prompt = [item if isinstance(item, str) else (item.get("prompt") or "") for item in req.prompts] or prompt
+        negative_prompt = _resolve_negative_prompts(
+            req.prompts,
+            negative_prompt if negative_prompt is not None else self.pipeline_recipe.negative_prompt,
         )
-        if timesteps is None:
-            num_inference_steps = max(int(num_inference_steps), 2)
-        elif len(timesteps) < 2:
-            raise ValueError("`timesteps` must contain at least 2 values for FlowMatchEulerDiscreteScheduler.")
+
+        height = sampling.height or height or self.pipeline_recipe.height
+        width = sampling.width or width or self.pipeline_recipe.width
+        num_frames = sampling.num_frames or num_frames or self.pipeline_recipe.num_frames
+        frame_rate = sampling.resolved_frame_rate or frame_rate or self.pipeline_recipe.frame_rate
+        num_inference_steps = (
+            sampling.num_inference_steps or num_inference_steps or self.pipeline_recipe.num_inference_steps
+        )
+        if is_dummy_run and self.pipeline_recipe.fixed_num_inference_steps:
+            num_inference_steps = self.pipeline_recipe.num_inference_steps
+        num_inference_steps = max(int(num_inference_steps), 2)
 
         num_videos_per_prompt = (
             sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else num_videos_per_prompt or 1
         )
         max_sequence_length = sampling.max_sequence_length or max_sequence_length or self.tokenizer_max_length
-        if sampling.guidance_scale_provided:
-            guidance_scale = sampling.guidance_scale
-        if self.supports_guidance_rescale and sampling.guidance_rescale is not None:
-            guidance_rescale = sampling.guidance_rescale
+        guidance_specs = [
+            _resolve_guidance_spec(
+                item,
+                self.pipeline_recipe.request_guidance,
+                guidance_scale=guidance_scale,
+            )
+            for item in sampling_params_list
+        ]
+        guidance = guidance_specs[0]
+        if is_dummy_run and not self.pipeline_recipe.phases[0].allow_guidance_override:
+            guidance = self.pipeline_recipe.request_guidance
+        if any(item != guidance for item in guidance_specs[1:]):
+            raise ValueError("Batched LTX requests must use identical video/audio guidance parameters.")
 
         if self.supports_request_batch:
             if generator is None:
@@ -235,6 +489,10 @@ class LTXRequestMixin:
             ]
             if any(value is not None for value in request_negative_masks):
                 negative_prompt_attention_mask = torch.stack(request_negative_masks)  # type: ignore[arg-type]
+            if prompt_embeds is not None:
+                prompt = None
+            if negative_prompt_embeds is not None:
+                negative_prompt = None
 
         if sampling.decode_timestep is not None:
             decode_timestep = sampling.decode_timestep
@@ -251,8 +509,7 @@ class LTXRequestMixin:
             num_frames=int(num_frames),
             frame_rate=float(frame_rate),
             num_inference_steps=int(num_inference_steps),
-            guidance_scale=guidance_scale,
-            guidance_rescale=guidance_rescale,
+            guidance=guidance,
             num_videos_per_prompt=int(num_videos_per_prompt),
             generator=generator,
             latents=latents,

vllm_omni/diffusion/models/ltx2/ltx2_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.interface import SupportsComponentDiscovery
 from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
 from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
 from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch, split_diffusion_output_by_request
+from vllm_omni.platforms import current_omni_platform
 
 from . import ltx2_latents as latent_ops
-from .ltx2_components import LTXComponentProfile, initialize_pipeline_components
+from .ltx2_components import (
+    LTXComponentProfile,
+    detect_ltx_model_version,
+    initialize_pipeline_components,
+    resolve_ltx_component_profile,
+)
 from .ltx2_conditioning import LTXPromptContext, LTXTextConditioningMixin
 from .ltx2_denoise import (
     LTXDenoiseContext,
     LTXForwardContext,
     LTXPhaseExecutor,
     LTXPhaseResult,
-    VideoAudioScheduler,
     build_transformer_kwargs,
     step_denoised_latents,
 )
-from .ltx2_guidance import LTXGuidanceStrategy
-from .ltx2_recipes import LTXOneStageRecipe
-from .ltx2_request import LTXRequestInputs, LTXRequestMixin
+from .ltx2_guidance import (
+    LTX_GUIDANCE_EXECUTOR,
+    LTXGuidanceExecutor,
+    LTXGuidancePlan,
+    LTXGuidanceSpec,
+    LTXModalityGuidance,
+)
+from .ltx2_recipes import (
+    LTXPhaseRecipe,
+    LTXPipelineRecipe,
+    resolve_ltx_pipeline_recipe,
+)
+from .ltx2_request import (
+    LTXRequestInputs,
+    LTXRequestMixin,
+    validate_pipeline_request,
+)
 
 
 def _expand_per_prompt_decode_value(
@@ -92,7 +112,7 @@ def _prepare_decode_timestep_conditioning(
     )
 
 
-class LTXPipelineRuntime(
+class LTXRuntime(
     LTXRequestMixin,
     LTXTextConditioningMixin,
     nn.Module,
@@ -101,27 +121,146 @@ class LTXPipelineRuntime(
     SupportsComponentDiscovery,
     DiffusionPipelineProfilerMixin,
 ):
-    """Shared Omni runtime for explicitly composed LTX denoise phases."""
+    """Shared Omni runtime for recipe-driven LTX denoise phases."""
 
+    pipeline_kind: ClassVar[str] = "one_stage"
     component_profile: ClassVar[LTXComponentProfile]
-    guidance_strategy: ClassVar[LTXGuidanceStrategy]
-    one_stage_recipe: ClassVar[LTXOneStageRecipe]
+    pipeline_recipe: ClassVar[LTXPipelineRecipe]
+    guidance_executor: ClassVar[LTXGuidanceExecutor] = LTX_GUIDANCE_EXECUTOR
     supports_request_batch = False
     connector_batches_cfg = False
     distributed_video_decode = True
     support_image_input = False
-    dummy_run_num_frames = 2
+    dummy_run_num_frames = 1
     preserve_sp_padded_audio_duration = False
     reports_stage_durations = False
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = "") -> None:
         del prefix
+        self.model_version = detect_ltx_model_version(od_config.model)
+        self.component_profile = resolve_ltx_component_profile(self.pipeline_kind, self.model_version)
+        self.pipeline_recipe = resolve_ltx_pipeline_recipe(self.pipeline_kind, self.model_version)
+        self._dit_modules = list(self.component_profile.dit_modules)
+        self._encoder_modules = list(self.component_profile.encoder_modules)
+        self._vae_modules = list(self.component_profile.vae_modules)
+        self._resident_modules = list(self.component_profile.resident_modules)
+        if self.model_version == "2.3":
+            self.preserve_sp_padded_audio_duration = True
+            self.reports_stage_durations = True
         super().__init__()
+        self._guidance_plan = LTXGuidancePlan.build(self.pipeline_recipe.request_guidance)
         initialize_pipeline_components(self, od_config)
         self.setup_diffusion_pipeline_profiler(
             enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler
         )
 
+    def _forward_request(
+        self,
+        req: DiffusionRequestBatch,
+        *,
+        image: Any | None = None,
+        prompt: str | list[str] | None = None,
+        negative_prompt: str | list[str] | None = None,
+        height: int | None = None,
+        width: int | None = None,
+        num_frames: int | None = None,
+        frame_rate: float | None = None,
+        num_inference_steps: int | None = None,
+        sigmas: list[float] | None = None,
+        guidance_scale: float | None = None,
+        num_videos_per_prompt: int | None = 1,
+        generator: torch.Generator | list[torch.Generator] | None = None,
+        latents: torch.Tensor | None = None,
+        audio_latents: torch.Tensor | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        negative_prompt_embeds: torch.Tensor | None = None,
+        prompt_attention_mask: torch.Tensor | None = None,
+        negative_prompt_attention_mask: torch.Tensor | None = None,
+        decode_timestep: float | list[float] = 0.0,
+        decode_noise_scale: float | list[float] | None = None,
+        output_type: str = "np",
+        max_sequence_length: int | None = None,
+    ) -> DiffusionOutput | list[DiffusionOutput]:
+        request_inputs = self._resolve_request_inputs(
+            req,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            frame_rate=frame_rate,
+            num_inference_steps=num_inference_steps,
+            guidance_scale=guidance_scale,
+            num_videos_per_prompt=num_videos_per_prompt,
+            generator=generator,
+            latents=latents,
+            audio_latents=audio_latents,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            prompt_attention_mask=prompt_attention_mask,
+            negative_prompt_attention_mask=negative_prompt_attention_mask,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+            output_type=output_type,
+            max_sequence_length=max_sequence_length,
+        )
+        image = self._resolve_request_image(req, image, request_inputs)
+        request_sigmas = self._resolve_request_sigmas(req, sigmas)
+        validate_pipeline_request(
+            request_inputs,
+            pipeline_recipe=self.pipeline_recipe,
+            vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
+            vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
+            pipeline_name=self.__class__.__name__,
+            request_sigmas=request_sigmas,
+        )
+        return self._run_recipe(req, request_inputs, request_sigmas=request_sigmas, image=image)
+
+    def _run_recipe(
+        self,
+        req: DiffusionRequestBatch,
+        request_inputs: LTXRequestInputs,
+        *,
+        request_sigmas: list[float] | None,
+        image: Any | None = None,
+    ) -> DiffusionOutput | list[DiffusionOutput]:
+        """Execute one- and multi-phase recipes through the same control flow."""
+        phase_results: list[LTXPhaseResult] = []
+        prompt_context = None
+
+        for phase_recipe in self.pipeline_recipe.phases:
+            self._enter_phase(phase_recipe)
+            phase_inputs = self._build_phase_inputs(
+                request_inputs,
+                phase_recipe,
+                phase_results[-1] if phase_results else None,
+            )
+            phase_result = self.run_phase(
+                req,
+                phase_inputs,
+                noise_scale=phase_recipe.noise_scale,
+                sigmas=list(phase_recipe.sigmas) if phase_recipe.sigmas is not None else request_sigmas,
+                timesteps=None,
+                attention_kwargs=None,
+                phase_recipe=phase_recipe,
+                image=image,
+                prompt_context=prompt_context,
+            )
+            phase_results.append(phase_result)
+            prompt_context = phase_result.forward_context.prompt_context
+
+        final_context = phase_results[-1].forward_context
+        output_phase = LTXPhaseResult(
+            forward_context=final_context,
+            video=phase_results[self.pipeline_recipe.video_output_phase].video,
+            audio=phase_results[self.pipeline_recipe.audio_output_phase].audio,
+        )
+        return self.decode_phase(output_phase)
+
+    def _enter_phase(self, phase: LTXPhaseRecipe) -> None:
+        """Hook for a future phase-weight strategy."""
+        del phase
+
     def prepare_latents(
         self,
         batch_size: int = 1,
@@ -160,6 +299,7 @@ class LTXPipelineRuntime(
         device: torch.device | None = None,
         generator: torch.Generator | list[torch.Generator] | None = None,
         latents: torch.Tensor | None = None,
+        latents_normalized: bool = False,
     ) -> tuple[torch.Tensor, int, int]:
         return latent_ops.prepare_audio_latents(
             self,
@@ -172,19 +312,16 @@ class LTXPipelineRuntime(
             device,
             generator,
             latents,
+            latents_normalized,
         )
 
     @property
     def guidance_scale(self):
-        return self._guidance_scale
-
-    @property
-    def guidance_rescale(self):
-        return self._guidance_rescale
+        return self._guidance_plan.spec.video.cfg_scale
 
     @property
     def do_classifier_free_guidance(self):
-        return self._guidance_scale is not None and self._guidance_scale > 1.0
+        return self._guidance_plan.spec.do_cfg
 
     @property
     def interrupt(self):
@@ -212,21 +349,29 @@ class LTXPipelineRuntime(
     ):
         if kwargs is not None:
             context = {**kwargs, **context}
-        return self.guidance_strategy.combine_cfg_noise(
-            self,
-            positive_noise_pred,
-            negative_noise_pred,
-            true_cfg_scale,
-            cfg_normalize,
-            context,
+        del cfg_normalize
+        required = ("video_latents", "audio_latents", "video_sigma", "audio_sigma")
+        if any(context.get(name) is None for name in required):
+            raise ValueError("LTX x0-space CFG requires video/audio latents and sigmas.")
+        video_pos, audio_pos = positive_noise_pred
+        video_neg, audio_neg = negative_noise_pred
+        return (
+            self.guidance_executor.combine_cfg_velocity(
+                context["video_latents"],
+                video_pos,
+                video_neg,
+                context["video_sigma"],
+                true_cfg_scale,
+            ),
+            self.guidance_executor.combine_cfg_velocity(
+                context["audio_latents"],
+                audio_pos,
+                audio_neg,
+                context["audio_sigma"],
+                true_cfg_scale,
+            ),
         )
 
-    def predict_noise_with_parallel_cfg(self, *args, **kwargs):
-        predict_parallel_cfg = getattr(self.guidance_strategy, "predict_parallel_cfg", None)
-        if predict_parallel_cfg is None:
-            raise NotImplementedError("The selected LTX guidance strategy does not implement parallel CFG.")
-        return predict_parallel_cfg(self, *args, **kwargs)
-
     def _synchronize_cfg_parallel_step_output(
         self,
         latents: tuple[torch.Tensor, torch.Tensor],
@@ -238,23 +383,35 @@ class LTXPipelineRuntime(
         # CUDA async execution otherwise permits numerical drift to accumulate
         # across CFG-parallel denoise steps.
         latents = tuple(tensor.contiguous() for tensor in latents)
-        device = next((tensor.device for tensor in latents if tensor.is_cuda), None)
-        if device is not None:
-            torch.cuda.current_stream(device).synchronize()
+        current_omni_platform.synchronize()
         return latents
 
     def _setup_forward_runtime(
         self,
+        req: DiffusionRequestBatch,
         request_inputs: LTXRequestInputs,
         attention_kwargs: dict[str, Any] | None,
     ) -> bool:
-        self._guidance_scale = request_inputs.guidance_scale
-        self._guidance_rescale = request_inputs.guidance_rescale
+        self._guidance_plan = LTXGuidancePlan.build(request_inputs.guidance)
         del attention_kwargs
         self._interrupt = False
         cfg_world_size = get_classifier_free_guidance_world_size()
+        if (
+            req.is_dummy_run()
+            and cfg_world_size == 2
+            and self.do_classifier_free_guidance
+            and not self._guidance_plan.cfg_parallel_compatible
+        ):
+            # The single-device warmup follows the full recipe. CFG2 executes
+            # the same per-modality CFG scales without STG/modality passes.
+            self._guidance_plan = LTXGuidancePlan.build(
+                LTXGuidanceSpec(
+                    video=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.video.cfg_scale),
+                    audio=LTXModalityGuidance(cfg_scale=self._guidance_plan.spec.audio.cfg_scale),
+                )
+            )
         if self.do_classifier_free_guidance:
-            self.guidance_strategy.validate_cfg_world_size(cfg_world_size)
+            self.guidance_executor.validate_cfg_world_size(self._guidance_plan, cfg_world_size)
         return self.do_classifier_free_guidance and cfg_world_size > 1
 
     def _check_forward_inputs(
@@ -325,11 +482,8 @@ class LTXPipelineRuntime(
         vae_decode_needs_all_ranks = False
         is_distributed_vae_enabled = getattr(self.vae, "is_distributed_enabled", None)
         if self.distributed_video_decode and dist_initialized and callable(is_distributed_vae_enabled):
-            try:
-                # Distributed tiled decode is collective, so every rank must enter it.
-                vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled())
-            except Exception:
-                pass
+            # Distributed tiled decode is collective, so every rank must enter it.
+            vae_decode_needs_all_ranks = bool(is_distributed_vae_enabled())
 
         should_decode_video = not self.distributed_video_decode or is_output_rank or vae_decode_needs_all_ranks
         if should_decode_video:
@@ -361,16 +515,16 @@ class LTXPipelineRuntime(
         image: Any | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         latents = self.prepare_latents(
-            prompt_context.batch_size * request_inputs.num_videos_per_prompt,
-            self.transformer.config.in_channels,
-            request_inputs.height,
-            request_inputs.width,
-            request_inputs.num_frames,
-            noise_scale,
-            torch.float32,
-            device,
-            request_inputs.generator,
-            request_inputs.latents,
+            batch_size=prompt_context.batch_size * request_inputs.num_videos_per_prompt,
+            num_channels_latents=self.transformer.config.in_channels,
+            height=request_inputs.height,
+            width=request_inputs.width,
+            num_frames=request_inputs.num_frames,
+            noise_scale=noise_scale,
+            dtype=prompt_context.positive_connector_prompt_embeds.dtype,
+            device=device,
+            generator=request_inputs.generator,
+            latents=request_inputs.latents,
         )
         return latents, None
 
@@ -416,10 +570,11 @@ class LTXPipelineRuntime(
             audio_latent_length=audio_num_frames,
             num_mel_bins=num_mel_bins,
             noise_scale=noise_scale,
-            dtype=torch.float32,
+            dtype=prompt_context.positive_connector_audio_prompt_embeds.dtype,
             device=device,
             generator=request_inputs.generator,
             latents=request_inputs.audio_latents,
+            latents_normalized=request_inputs.audio_latents_normalized,
         )
         return audio_latents, original_num_frames, padded_num_frames, latent_mel_bins
 
@@ -439,29 +594,37 @@ class LTXPipelineRuntime(
         padded_length = latent_ops.get_sp_padded_audio_latent_length(requested_length, int(sp_size))
         return requested_length if provided_length in {requested_length, padded_length} else provided_length
 
-    def _make_video_audio_scheduler(
-        self,
-        audio_scheduler: Any,
-        latent_num_frames: int,
-        latent_height: int,
-        latent_width: int,
-    ) -> Any:
-        return VideoAudioScheduler(self.scheduler, audio_scheduler)
-
     def _prepare_denoise_context_for_cfg(
         self,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
     ) -> LTXDenoiseContext:
-        return self.guidance_strategy.prepare_denoise_context(self, forward_ctx, denoise_ctx)
+        return self.guidance_executor.prepare_denoise_context(
+            self._guidance_plan,
+            forward_ctx.cfg_parallel_ready,
+            denoise_ctx,
+        )
 
     def _denoise_timestep_kwargs(
         self,
         ts: torch.Tensor,
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
+        *,
+        video_token_count: int,
+        audio_token_count: int,
     ) -> dict[str, torch.Tensor]:
-        return self.guidance_strategy.timestep_kwargs(ts, forward_ctx, denoise_ctx)
+        del forward_ctx, denoise_ctx
+        return self.guidance_executor.timestep_kwargs(ts, video_token_count, audio_token_count)
+
+    def _video_guidance_model_sigma(
+        self,
+        sigma: torch.Tensor,
+        denoise_ctx: LTXDenoiseContext,
+    ) -> torch.Tensor:
+        """Return the video timestep used to convert velocity predictions to x0."""
+        del denoise_ctx
+        return sigma
 
     def _build_transformer_kwargs(
         self,
@@ -475,6 +638,7 @@ class LTXPipelineRuntime(
         encoder_attention_mask: torch.Tensor | None,
         audio_encoder_attention_mask: torch.Tensor | None,
         ts: torch.Tensor,
+        attention_kwargs: dict[str, Any] | None = None,
     ) -> dict[str, Any]:
         return build_transformer_kwargs(
             self,
@@ -487,6 +651,7 @@ class LTXPipelineRuntime(
             encoder_attention_mask=encoder_attention_mask,
             audio_encoder_attention_mask=audio_encoder_attention_mask,
             ts=ts,
+            attention_kwargs=attention_kwargs,
         )
 
     def _predict_noise_for_step(
@@ -497,8 +662,9 @@ class LTXPipelineRuntime(
         forward_ctx: LTXForwardContext,
         denoise_ctx: LTXDenoiseContext,
     ) -> tuple[torch.Tensor, torch.Tensor]:
-        return self.guidance_strategy.predict_noise(
+        return self.guidance_executor.predict_noise(
             self,
+            self._guidance_plan,
             index,
             timestep,
             state,
@@ -562,7 +728,6 @@ class LTXPipelineRuntime(
         )
         audio_latents = latent_ops.unpack_audio_latents(
             audio_latents,
-            forward_ctx.original_audio_num_frames,
             num_mel_bins=forward_ctx.latent_mel_bins,
         )
         return latents, audio_latents
@@ -576,6 +741,7 @@ class LTXPipelineRuntime(
         sigmas: list[float] | None,
         timesteps: list[int] | None,
         attention_kwargs: dict[str, Any] | None,
+        phase_recipe: LTXPhaseRecipe,
         image: Any | None = None,
         prompt_context: LTXPromptContext | None = None,
     ) -> LTXPhaseResult:
@@ -588,10 +754,84 @@ class LTXPipelineRuntime(
             sigmas=sigmas,
             timesteps=timesteps,
             attention_kwargs=attention_kwargs,
+            phase_recipe=phase_recipe,
             image=image,
             prompt_context=prompt_context,
         )
 
+    def _build_phase_inputs(
+        self,
+        request_inputs: LTXRequestInputs,
+        phase: LTXPhaseRecipe,
+        previous_phase: LTXPhaseResult | None,
+    ) -> LTXRequestInputs:
+        """Resolve one phase from immutable request inputs and prior AV state."""
+        divisor = phase.spatial_downscale
+        if request_inputs.height % divisor != 0 or request_inputs.width % divisor != 0:
+            raise ValueError(
+                f"LTX phase {phase.name!r} cannot scale resolution "
+                f"{request_inputs.width}x{request_inputs.height} by {divisor}."
+            )
+        height = request_inputs.height // divisor
+        width = request_inputs.width // divisor
+        latents = request_inputs.latents
+        audio_latents = request_inputs.audio_latents
+        audio_latents_normalized = request_inputs.audio_latents_normalized
+        decode_timestep = request_inputs.decode_timestep
+        decode_noise_scale = request_inputs.decode_noise_scale
+
+        if phase.input_transform == "spatial_upsample":
+            if previous_phase is None:
+                raise ValueError(f"LTX phase {phase.name!r} requires a previous phase to upsample.")
+            if previous_phase.video.ndim != 5:
+                raise ValueError(f"LTX spatial upsampling expects a 5D video latent, got {previous_phase.video.shape}.")
+            latents = self._spatial_upsample_phase(previous_phase.video)
+            if previous_phase.audio_for_next_phase is not None:
+                audio_latents = previous_phase.audio_for_next_phase
+                audio_latents_normalized = True
+            else:
+                audio_latents = previous_phase.audio
+                audio_latents_normalized = False
+            decode_timestep = 0.0
+            decode_noise_scale = None
+
+            expected_shape = latent_ops.resolve_video_latent_shape(
+                height,
+                width,
+                request_inputs.num_frames,
+                vae_spatial_compression_ratio=self.vae_spatial_compression_ratio,
+                vae_temporal_compression_ratio=self.vae_temporal_compression_ratio,
+            )
+            if latents.shape[2:] != expected_shape:
+                raise ValueError(
+                    f"LTX phase {phase.name!r} upsampler produced {tuple(latents.shape[2:])}, expected "
+                    f"{expected_shape} for {width}x{height}."
+                )
+        elif phase.input_transform != "initial":
+            raise ValueError(f"Unsupported LTX phase input transform: {phase.input_transform!r}.")
+
+        guidance = request_inputs.guidance if phase.allow_guidance_override else phase.guidance
+        return replace(
+            request_inputs,
+            height=height,
+            width=width,
+            num_inference_steps=phase.num_inference_steps or request_inputs.num_inference_steps,
+            guidance=guidance,
+            latents=latents,
+            audio_latents=audio_latents,
+            audio_latents_normalized=audio_latents_normalized,
+            decode_timestep=decode_timestep,
+            decode_noise_scale=decode_noise_scale,
+        )
+
+    def _spatial_upsample_phase(self, latents: torch.Tensor) -> torch.Tensor:
+        """Apply the recipe-managed spatial upsampler to unpacked latents."""
+        latent_upsampler = getattr(self, "latent_upsampler", None)
+        if latent_upsampler is None:
+            raise RuntimeError("This LTX pipeline recipe requires a spatial latent upsampler.")
+        dtype = getattr(latent_upsampler, "dtype", latents.dtype)
+        return latent_upsampler(latents.to(device=self.device, dtype=dtype))
+
     def decode_phase(self, phase: LTXPhaseResult) -> DiffusionOutput | list[DiffusionOutput]:
         """Decode one completed phase and restore per-request outputs."""
         forward_ctx = phase.forward_context

vllm_omni/diffusion/models/ltx2/ltx2_transformer.py

@@ -57,6 +57,9 @@ _RMSNORM_INIT_PARAMS = inspect.signature(RMSNorm.__init__).parameters
 
 def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) -> nn.Module:
     """Bridge diffusers' RMSNorm API onto vLLM's `has_weight` variant."""
+    if not elementwise_affine:
+        return nn.RMSNorm(hidden_size, eps=eps, elementwise_affine=False)
+
     kwargs: dict[str, Any] = {"eps": eps}
     if "elementwise_affine" in _RMSNORM_INIT_PARAMS:
         kwargs["elementwise_affine"] = elementwise_affine
@@ -79,10 +82,14 @@ def _make_rms_norm(hidden_size: int, *, eps: float, elementwise_affine: bool) ->
 
 def apply_interleaved_rotary_emb(x: torch.Tensor, freqs: tuple[torch.Tensor, torch.Tensor]) -> torch.Tensor:
     cos, sin = freqs
+    if cos.dtype != x.dtype:
+        cos = cos.to(dtype=x.dtype)
+    if sin.dtype != x.dtype:
+        sin = sin.to(dtype=x.dtype)
     # Concrete pair count instead of -1 keeps SDPA shape static under torch.compile.
     x_real, x_imag = x.unflatten(2, (x.shape[2] // 2, 2)).unbind(-1)  # [B, S, C // 2]
     x_rotated = torch.stack([-x_imag, x_real], dim=-1).flatten(2)
-    out = (x.float() * cos + x_rotated.float() * sin).to(x.dtype)
+    out = x * cos + x_rotated * sin
     return out
 
 
@@ -94,6 +101,10 @@ def apply_split_rotary_emb(
 ) -> torch.Tensor:
     # `head_dim` is plumbed in (not inferred via `-1`) so SDPA shape stays static under torch.compile.
     cos, sin = freqs
+    if cos.dtype != x.dtype:
+        cos = cos.to(dtype=x.dtype)
+    if sin.dtype != x.dtype:
+        sin = sin.to(dtype=x.dtype)
 
     x_dtype = x.dtype
     needs_reshape = False
@@ -112,7 +123,7 @@ def apply_split_rotary_emb(
     r = last // 2
 
     # (..., 2, r)
-    split_x = x.reshape(*x.shape[:-1], 2, r).float()  # Explicitly upcast to float
+    split_x = x.reshape(*x.shape[:-1], 2, r)
     first_x = split_x[..., :1, :]  # (..., 1, r)
     second_x = split_x[..., 1:, :]  # (..., 1, r)
 
@@ -327,6 +338,23 @@ class TensorParallelRMSNorm(nn.Module):
         return out.to(dtype=x_dtype)
 
 
+def to_ltx_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor:
+    """Convert an additive/expanded LTX mask to a 2D padding mask."""
+    if attention_mask.ndim > 2:
+        if attention_mask.is_floating_point():
+            valid = attention_mask >= 0
+        else:
+            valid = attention_mask.to(torch.bool)
+        batch_size = valid.shape[0]
+        key_length = valid.shape[-1]
+        attention_mask = valid.reshape(batch_size, -1, key_length).all(dim=1)
+    if attention_mask.is_floating_point():
+        attention_mask = attention_mask >= 0
+    if attention_mask.dtype != torch.bool:
+        attention_mask = attention_mask.to(torch.bool)
+    return attention_mask
+
+
 class LTX2AudioVideoAttnProcessor:
     r"""
     Processor for implementing attention (SDPA is used by default if you're using PyTorch 2.0) for the LTX-2.0 model.
@@ -344,24 +372,6 @@ class LTX2AudioVideoAttnProcessor:
                 "Please upgrade your PyTorch installation."
             )
 
-    @staticmethod
-    def _to_padding_mask(attention_mask: torch.Tensor) -> torch.Tensor:
-        # Convert additive/expanded masks into a 2D padding mask for flash-attn.
-        if attention_mask.ndim > 2:
-            if attention_mask.is_floating_point():
-                valid = attention_mask >= 0
-            else:
-                valid = attention_mask.to(torch.bool)
-            b = valid.shape[0]
-            key_len = valid.shape[-1]
-            valid = valid.reshape(b, -1, key_len).all(dim=1)
-            attention_mask = valid
-        if attention_mask.is_floating_point():
-            attention_mask = attention_mask >= 0
-        if attention_mask.dtype != torch.bool:
-            attention_mask = attention_mask.to(torch.bool)
-        return attention_mask
-
     @staticmethod
     def _is_sp_enabled() -> bool:
         if not is_forward_context_available():
@@ -390,12 +400,12 @@ class LTX2AudioVideoAttnProcessor:
             # For cross-attention, encoder sequence length != query length, so drop the mask.
             if encoder_hidden_states is not None and encoder_hidden_states.shape[1] != hidden_states.shape[1]:
                 return None
-            return self._to_padding_mask(attention_mask)
+            return to_ltx_padding_mask(attention_mask)
 
         attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
         attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])
         if attn.attn.attn_backend.get_name().upper() == "FLASH_ATTN":
-            attention_mask = self._to_padding_mask(attention_mask)
+            attention_mask = to_ltx_padding_mask(attention_mask)
         return attention_mask
 
     @staticmethod
@@ -467,6 +477,7 @@ class LTX2AudioVideoAttnProcessor:
         attention_mask: torch.Tensor | None = None,
         query_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
         key_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None,
+        perturbation_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         is_self_attention = encoder_hidden_states is None
         batch_size, sequence_length, _ = hidden_states.shape if is_self_attention else encoder_hidden_states.shape
@@ -509,17 +520,21 @@ class LTX2AudioVideoAttnProcessor:
             elif attn.rope_type == "split":
                 query = apply_split_rotary_emb(query, query_rotary_emb, head_dim=attn.head_dim)
                 key = apply_split_rotary_emb(
-                    key, key_rotary_emb if key_rotary_emb is not None else query_rotary_emb, head_dim=attn.head_dim
+                    key,
+                    key_rotary_emb if key_rotary_emb is not None else query_rotary_emb,
+                    head_dim=attn.head_dim,
                 )
 
         query = query.unflatten(2, (attn.heads, attn.head_dim))
         key = key.unflatten(2, (attn.heads, attn.head_dim))
-        value = value.unflatten(2, (attn.heads, attn.head_dim))
+        value_heads = value.unflatten(2, (attn.heads, attn.head_dim))
 
         attn_metadata = AttentionMetadata(attn_mask=attention_mask) if attention_mask is not None else None
-        hidden_states = attn.attn(query, key, value, attn_metadata)
+        hidden_states = attn.attn(query, key, value_heads, attn_metadata)
         hidden_states = hidden_states.flatten(2, 3)
         hidden_states = hidden_states.to(query.dtype)
+        if perturbation_mask is not None:
+            hidden_states = hidden_states * perturbation_mask + value * (1 - perturbation_mask)
 
         # LTX-2.3: per-head gated attention
         if attn.to_gate_logits is not None:
@@ -560,6 +575,7 @@ class LTX2Attention(torch.nn.Module):
         norm_elementwise_affine: bool = True,
         rope_type: str = "interleaved",
         apply_gated_attention: bool = False,
+        pack_qkv: bool = True,
         processor=None,
         quant_config: "QuantizationConfig | None" = None,
         prefix: str = "",
@@ -591,7 +607,7 @@ class LTX2Attention(torch.nn.Module):
         self.to_q = None
         self.to_k = None
         self.to_v = None
-        if cross_attention_dim is None:
+        if cross_attention_dim is None and pack_qkv:
             self.to_qkv = QKVParallelLinear(
                 hidden_size=query_dim,
                 head_size=self.head_dim,
@@ -847,6 +863,9 @@ class LTX2VideoTransformerBlock(nn.Module):
             qk_norm=qk_norm,
             rope_type=rope_type,
             apply_gated_attention=audio_gated_attn,
+            # The official checkpoint evaluates audio Q/K/V separately. Fusing
+            # the bf16 GEMMs changes denoise numerics enough to accumulate.
+            pack_qkv=False,
             quant_config=quant_config,
             prefix=f"{prefix}.audio_attn1" if prefix else "audio_attn1",
         )
@@ -992,8 +1011,10 @@ class LTX2VideoTransformerBlock(nn.Module):
         v2a_cross_attention_mask: torch.Tensor | None = None,
         use_a2v_cross_attention: bool = True,
         use_v2a_cross_attention: bool = True,
-        perturbation_mask: torch.Tensor | None = None,
-        all_perturbed: bool | None = None,
+        video_self_attention_perturbation_mask: torch.Tensor | None = None,
+        audio_self_attention_perturbation_mask: torch.Tensor | None = None,
+        a2v_cross_attention_perturbation_mask: torch.Tensor | None = None,
+        v2a_cross_attention_perturbation_mask: torch.Tensor | None = None,
     ) -> torch.Tensor:
         batch_size = hidden_states.size(0)
 
@@ -1012,6 +1033,7 @@ class LTX2VideoTransformerBlock(nn.Module):
             encoder_hidden_states=None,
             query_rotary_emb=video_rotary_emb,
             attention_mask=self_attention_mask,
+            perturbation_mask=video_self_attention_perturbation_mask,
         )
         hidden_states = hidden_states + attn_hidden_states * gate_msa
 
@@ -1031,6 +1053,7 @@ class LTX2VideoTransformerBlock(nn.Module):
             encoder_hidden_states=None,
             query_rotary_emb=audio_rotary_emb,
             attention_mask=audio_self_attention_mask,
+            perturbation_mask=audio_self_attention_perturbation_mask,
         )
         audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * audio_gate_msa
 
@@ -1123,6 +1146,8 @@ class LTX2VideoTransformerBlock(nn.Module):
                     key_rotary_emb=ca_audio_rotary_emb,
                     attention_mask=a2v_cross_attention_mask,
                 )
+                if a2v_cross_attention_perturbation_mask is not None:
+                    a2v_attn_hidden_states = a2v_attn_hidden_states * a2v_cross_attention_perturbation_mask
                 hidden_states = hidden_states + a2v_gate * a2v_attn_hidden_states
 
             # 3.3. Video-to-Audio Cross Attention: Q: Audio; K,V: Video
@@ -1141,6 +1166,8 @@ class LTX2VideoTransformerBlock(nn.Module):
                     key_rotary_emb=ca_video_rotary_emb,
                     attention_mask=v2a_cross_attention_mask,
                 )
+                if v2a_cross_attention_perturbation_mask is not None:
+                    v2a_attn_hidden_states = v2a_attn_hidden_states * v2a_cross_attention_perturbation_mask
                 audio_hidden_states = audio_hidden_states + v2a_gate * v2a_attn_hidden_states
 
         # 4. Feedforward
@@ -1345,7 +1372,10 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
             return self.prepare_audio_coords(*args, **kwargs)
 
     def forward(
-        self, coords: torch.Tensor, device: str | torch.device | None = None
+        self,
+        coords: torch.Tensor,
+        device: str | torch.device | None = None,
+        out_dtype: torch.dtype | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor]:
         device = device or coords.device
 
@@ -1421,6 +1451,9 @@ class LTX2AudioVideoRotaryPosEmbed(nn.Module):
             cos_freqs = torch.swapaxes(cos_freq, 1, 2)  # (B,H,T,D//2)
             sin_freqs = torch.swapaxes(sin_freq, 1, 2)  # (B,H,T,D//2)
 
+        if out_dtype is not None:
+            cos_freqs = cos_freqs.to(out_dtype)
+            sin_freqs = sin_freqs.to(out_dtype)
         return cos_freqs, sin_freqs
 
 
@@ -1488,8 +1521,11 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 # Shard prompt embeds across sequence
                 "encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
                 "audio_encoder_hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3, split_output=False),
-                # Shard video timestep when provided as (B, seq_len)
+                # Keep per-token video/audio timestep embeddings aligned with
+                # their corresponding latent sequence shards. Official LTX
+                # multi-guidance supplies both tensors as (B, seq_len).
                 "timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False),
+                "audio_timestep": SequenceParallelInput(split_dim=1, expected_dims=2, split_output=False),
             },
             "rope": {
                 0: SequenceParallelInput(split_dim=rope_split_dim, expected_dims=rope_expected_dims, split_output=True),
@@ -1856,6 +1892,12 @@ class LTX2VideoTransformer3DModel(nn.Module):
         # Determine timestep for audio.
         audio_timestep = audio_timestep if audio_timestep is not None else timestep
 
+        def batch_step(step: torch.Tensor) -> torch.Tensor:
+            return step[:, 0] if step.ndim > 1 else step
+
+        sigma = sigma if sigma is not None else batch_step(timestep)
+        audio_sigma = audio_sigma if audio_sigma is not None else batch_step(audio_timestep)
+
         # convert encoder_attention_mask to a bias the same way we do for attention_mask
         if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
             encoder_attention_mask = (1 - encoder_attention_mask.to(hidden_states.dtype)) * -10000.0
@@ -1877,12 +1919,22 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 batch_size, audio_num_frames, audio_hidden_states.device
             )
 
-        video_rotary_emb = self.rope(video_coords, device=hidden_states.device)
-        audio_rotary_emb = self.audio_rope(audio_coords, device=audio_hidden_states.device)
+        video_rotary_emb = self.rope(video_coords, device=hidden_states.device, out_dtype=hidden_states.dtype)
+        audio_rotary_emb = self.audio_rope(
+            audio_coords,
+            device=audio_hidden_states.device,
+            out_dtype=audio_hidden_states.dtype,
+        )
 
-        video_cross_attn_rotary_emb = self.cross_attn_rope(video_coords[:, 0:1, :], device=hidden_states.device)
+        video_cross_attn_rotary_emb = self.cross_attn_rope(
+            video_coords[:, 0:1, :],
+            device=hidden_states.device,
+            out_dtype=hidden_states.dtype,
+        )
         audio_cross_attn_rotary_emb = self.cross_attn_audio_rope(
-            audio_coords[:, 0:1, :], device=audio_hidden_states.device
+            audio_coords[:, 0:1, :],
+            device=audio_hidden_states.device,
+            out_dtype=audio_hidden_states.dtype,
         )
 
         # 2. Patchify input projections
@@ -1920,7 +1972,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
             hidden_dtype=hidden_states.dtype,
         )
         video_cross_attn_a2v_gate, _ = self.av_cross_attn_video_a2v_gate(
-            timestep.flatten() * timestep_cross_attn_gate_scale_factor,
+            audio_sigma.flatten() * timestep_cross_attn_gate_scale_factor,
             batch_size=batch_size,
             hidden_dtype=hidden_states.dtype,
         )
@@ -1935,7 +1987,7 @@ class LTX2VideoTransformer3DModel(nn.Module):
             hidden_dtype=audio_hidden_states.dtype,
         )
         audio_cross_attn_v2a_gate, _ = self.av_cross_attn_audio_v2a_gate(
-            audio_timestep.flatten() * timestep_cross_attn_gate_scale_factor,
+            sigma.flatten() * timestep_cross_attn_gate_scale_factor,
             batch_size=batch_size,
             hidden_dtype=audio_hidden_states.dtype,
         )
@@ -1945,7 +1997,6 @@ class LTX2VideoTransformer3DModel(nn.Module):
         audio_cross_attn_v2a_gate = audio_cross_attn_v2a_gate.view(batch_size, -1, audio_cross_attn_v2a_gate.shape[-1])
 
         # 3.3. LTX-2.3: Compute prompt modulation from sigma
-        audio_sigma = audio_sigma if audio_sigma is not None else sigma
         if self.prompt_modulation and sigma is not None:
             temb_prompt, _ = self.prompt_adaln(sigma.flatten(), batch_size=batch_size, hidden_dtype=hidden_states.dtype)
             temb_prompt_audio, _ = self.audio_prompt_adaln(
@@ -1967,8 +2018,17 @@ class LTX2VideoTransformer3DModel(nn.Module):
             audio_encoder_hidden_states = self.audio_caption_projection(audio_encoder_hidden_states)
             audio_encoder_hidden_states = audio_encoder_hidden_states.view(batch_size, -1, audio_hidden_states.size(-1))
 
+        perturbation_kwargs = (attention_kwargs or {}).get("ltx_perturbation_kwargs", {})
+
+        def perturbation_mask_for(name: str, block_idx: int) -> torch.Tensor | None:
+            mask = perturbation_kwargs.get(f"{name}_mask")
+            blocks = perturbation_kwargs.get(f"{name}_blocks")
+            if mask is None or (blocks is not None and block_idx not in blocks):
+                return None
+            return mask
+
         # 5. Run transformer blocks
-        for block in self.transformer_blocks:
+        for block_idx, block in enumerate(self.transformer_blocks):
             block_kwargs = {
                 "temb": temb,
                 "temb_audio": temb_audio,
@@ -1984,6 +2044,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
                 "ca_audio_rotary_emb": audio_cross_attn_rotary_emb,
                 "encoder_attention_mask": encoder_attention_mask,
                 "audio_encoder_attention_mask": audio_encoder_attention_mask,
+                "video_self_attention_perturbation_mask": perturbation_mask_for("video_self_attention", block_idx),
+                "audio_self_attention_perturbation_mask": perturbation_mask_for("audio_self_attention", block_idx),
+                "a2v_cross_attention_perturbation_mask": perturbation_mask_for("a2v_cross_attention", block_idx),
+                "v2a_cross_attention_perturbation_mask": perturbation_mask_for("v2a_cross_attention", block_idx),
             }
             # NOTE: We intentionally pass the first 4 args here positionally,
             # because CacheDiT expects the first 2 args to be hidden_states
@@ -2066,7 +2130,10 @@ class LTX2VideoTransformer3DModel(nn.Module):
             for param_name, weight_name, shard_id in stacked_params_mapping:
                 if weight_name not in name:
                     continue
-                name = name.replace(weight_name, param_name)
+                packed_name = name.replace(weight_name, param_name)
+                if packed_name not in params_dict:
+                    continue
+                name = packed_name
                 param = params_dict[name]
                 weight_loader = param.weight_loader
                 weight_loader(param, loaded_weight, shard_id)

vllm_omni/diffusion/models/ltx2/pipeline_ltx2.py

@@ -5,250 +5,56 @@
 
 from __future__ import annotations
 
-from typing import Any, ClassVar
+from typing import ClassVar
 
-import torch
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 from vllm_omni.diffusion.models.dmd2 import DMD2PipelineMixin
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
 
-from .ltx2_components import (
-    LTX2_COMPONENT_PROFILE,
-    LTX23_COMPONENT_PROFILE,
-    LTXComponentProfile,
-)
+from .ltx2_components import LTX2_COMPONENT_PROFILE
 from .ltx2_components import (
     get_ltx2_post_process_func as get_ltx2_post_process_func,  # noqa: F401
 )
 from .ltx2_conditioning import LTXI2VConditioningMixin
-from .ltx2_guidance import (
-    LTX_LEGACY_VELOCITY_GUIDANCE,
-    LTX_OFFICIAL_X0_GUIDANCE,
+from .ltx2_recipes import (
+    LTX2_ONE_STAGE_RECIPE,
+    LTX_POSITIVE_ONLY_RECIPE,
 )
-from .ltx2_pipeline_runtime import LTXPipelineRuntime
-from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE, LTX23_ONE_STAGE_RECIPE, LTXOneStageRecipe
-from .ltx2_request import LTXRequestInputs
+from .ltx2_runtime import LTXRuntime
 
 
-class LTXOneStagePipeline(LTXPipelineRuntime):
-    """Single execution path configured by model-version and task entries."""
+class LTX2Pipeline(LTXI2VConditioningMixin, LTXRuntime):
+    """LTX-2 family one-stage entry, configured from checkpoint metadata."""
 
-    component_profile: ClassVar[LTXComponentProfile]
-    one_stage_recipe: ClassVar[LTXOneStageRecipe]
-
-    supports_request_batch = True
-    supports_guidance_rescale = False
-
-    def _forward_impl(
-        self,
-        req: DiffusionRequestBatch,
-        request_inputs: LTXRequestInputs,
-        *,
-        noise_scale: float,
-        sigmas: list[float] | None,
-        timesteps: list[int] | None,
-        attention_kwargs: dict[str, Any] | None,
-        image: Any | None = None,
-    ) -> DiffusionOutput | list[DiffusionOutput]:
-        phase = self.run_phase(
-            req,
-            request_inputs,
-            noise_scale=noise_scale,
-            sigmas=sigmas,
-            timesteps=timesteps,
-            attention_kwargs=attention_kwargs,
-            image=image,
-        )
-        return self.decode_phase(phase)
-
-    @torch.no_grad()
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_frames: int | None = None,
-        frame_rate: float | None = None,
-        num_inference_steps: int | None = None,
-        sigmas: list[float] | None = None,
-        timesteps: list[int] | None = None,
-        guidance_scale: float = 4.0,
-        guidance_rescale: float = 0.0,
-        noise_scale: float = 0.0,
-        num_videos_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        audio_latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        prompt_attention_mask: torch.Tensor | None = None,
-        negative_prompt_attention_mask: torch.Tensor | None = None,
-        decode_timestep: float | list[float] = 0.0,
-        decode_noise_scale: float | list[float] | None = None,
-        output_type: str = "np",
-        return_dict: bool = True,
-        attention_kwargs: dict[str, Any] | None = None,
-        max_sequence_length: int | None = None,
-        *,
-        image: Any | None = None,
-    ) -> DiffusionOutput | list[DiffusionOutput]:
-        del return_dict
-        sigmas = self._resolve_request_sigmas(req, sigmas)
-        request_inputs = self._resolve_request_inputs(
-            req,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            frame_rate=frame_rate,
-            num_inference_steps=num_inference_steps,
-            timesteps=timesteps,
-            guidance_scale=self.one_stage_recipe.guidance_scale if guidance_scale is None else guidance_scale,
-            guidance_rescale=guidance_rescale,
-            num_videos_per_prompt=num_videos_per_prompt,
-            generator=generator,
-            latents=latents,
-            audio_latents=audio_latents,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-            decode_timestep=decode_timestep,
-            decode_noise_scale=decode_noise_scale,
-            output_type=output_type,
-            max_sequence_length=max_sequence_length,
-        )
-        image = self._resolve_request_image(req, image, request_inputs)
-        forward_kwargs = {
-            "noise_scale": noise_scale,
-            "sigmas": sigmas,
-            "timesteps": timesteps,
-            "attention_kwargs": attention_kwargs,
-        }
-        if self.support_image_input:
-            forward_kwargs["image"] = image
-        return self._forward_impl(
-            req,
-            request_inputs,
-            **forward_kwargs,
-        )
-
-
-class LTX2Pipeline(LTXOneStagePipeline):
-    """LTX2 one-stage text-to-video entry."""
-
-    supports_guidance_rescale = True
+    pipeline_kind = "one_stage"
+    unified_text_image_entry = True
     component_profile = LTX2_COMPONENT_PROFILE
-    guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE
-    one_stage_recipe = LTX2_ONE_STAGE_RECIPE
+    pipeline_recipe = LTX2_ONE_STAGE_RECIPE
     _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules)
     _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules)
     _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules)
     _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules)
-
-
-class LTX23Pipeline(LTXOneStagePipeline):
-    """LTX2.3 one-stage text-to-video entry."""
-
-    connector_batches_cfg = True
-    preserve_sp_padded_audio_duration = True
-    reports_stage_durations = True
-    component_profile = LTX23_COMPONENT_PROFILE
-    guidance_strategy = LTX_OFFICIAL_X0_GUIDANCE
-    one_stage_recipe = LTX23_ONE_STAGE_RECIPE
-    _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules)
-    _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules)
-    _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules)
-    _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules)
-
-    @torch.no_grad()
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_frames: int | None = None,
-        frame_rate: float | None = None,
-        num_inference_steps: int | None = None,
-        sigmas: list[float] | None = None,
-        timesteps: list[int] | None = None,
-        guidance_scale: float = 4.0,
-        noise_scale: float = 0.0,
-        num_videos_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        audio_latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        prompt_attention_mask: torch.Tensor | None = None,
-        negative_prompt_attention_mask: torch.Tensor | None = None,
-        decode_timestep: float | list[float] = 0.0,
-        decode_noise_scale: float | list[float] | None = None,
-        output_type: str = "np",
-        return_dict: bool = True,
-        attention_kwargs: dict[str, Any] | None = None,
-        max_sequence_length: int | None = None,
-        *,
-        image: Any | None = None,
-    ) -> DiffusionOutput | list[DiffusionOutput]:
-        """Preserve the pre-refactor LTX2.3 positional argument order."""
-        return super().forward(
-            req,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            frame_rate=frame_rate,
-            num_inference_steps=num_inference_steps,
-            sigmas=sigmas,
-            timesteps=timesteps,
-            guidance_scale=guidance_scale,
-            guidance_rescale=0.0,
-            noise_scale=noise_scale,
-            num_videos_per_prompt=num_videos_per_prompt,
-            generator=generator,
-            latents=latents,
-            audio_latents=audio_latents,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-            decode_timestep=decode_timestep,
-            decode_noise_scale=decode_noise_scale,
-            output_type=output_type,
-            return_dict=return_dict,
-            attention_kwargs=attention_kwargs,
-            max_sequence_length=max_sequence_length,
-            image=image,
-        )
-
-
-class LTX2ImageToVideoPipeline(LTXI2VConditioningMixin, LTX2Pipeline):
-    """LTX2 one-stage image-to-video entry."""
-
-
-class LTX23ImageToVideoPipeline(LTXI2VConditioningMixin, LTX23Pipeline):
-    """LTX2.3 one-stage image-to-video entry."""
+    supports_request_batch = True
 
 
 class LTX2T2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline):
     """LTX2 T2V entry for FastGen DMD2-distilled models."""
 
+    pipeline_kind = "dmd2"
+    support_image_input = False
+    pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__(od_config=od_config, prefix=prefix)
         self.__init_dmd2__()
 
 
-class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2ImageToVideoPipeline):
+class LTX2I2VDMD2Pipeline(DMD2PipelineMixin, LTX2Pipeline):
     """LTX2 I2V entry for FastGen DMD2-distilled models."""
 
+    pipeline_kind = "dmd2"
+    unified_text_image_entry = False
+    pipeline_recipe = LTX_POSITIVE_ONLY_RECIPE
+
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
         super().__init__(od_config=od_config, prefix=prefix)
         self.__init_dmd2__()

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_latent_upsample.py

@@ -1,262 +0,0 @@
-import os
-
-import torch
-import torch.nn as nn
-from diffusers import AutoencoderKLLTX2Video
-from diffusers.image_processor import PipelineImageInput
-from diffusers.pipelines.ltx2.latent_upsampler import LTX2LatentUpsamplerModel
-from diffusers.pipelines.stable_diffusion.pipeline_stable_diffusion_img2img import retrieve_latents
-from diffusers.utils.torch_utils import randn_tensor
-from diffusers.video_processor import VideoProcessor
-from vllm.logger import init_logger
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.distributed.utils import get_local_device
-
-logger = init_logger(__name__)
-
-
-class LTX2LatentUpsamplePipeline(nn.Module):
-    def __init__(
-        self,
-        od_config: OmniDiffusionConfig,
-        vae: AutoencoderKLLTX2Video,
-        latent_upsampler: LTX2LatentUpsamplerModel = None,
-    ) -> None:
-        super().__init__()
-
-        if vae is None:
-            raise ValueError("vae must be provided")
-        self.vae = vae
-
-        self.device = get_local_device()
-        model = od_config.model
-        local_files_only = os.path.exists(model)
-
-        if latent_upsampler is None:
-            # Use cpu context to create latent upsampler. The code k[:, None] @ k[None, :] in
-            # diffuser's BlurDownsample is not supported on GPU as k is type of torch.Int64
-            with torch.device("cpu"):
-                latent_upsampler = LTX2LatentUpsamplerModel.from_pretrained(
-                    model,
-                    subfolder="latent_upsampler",
-                    torch_dtype=torch.bfloat16,
-                    local_files_only=local_files_only,
-                ).to(self.device)
-        self.latent_upsampler = latent_upsampler
-
-        self.vae_spatial_compression_ratio = (
-            self.vae.spatial_compression_ratio if getattr(self, "vae", None) is not None else 32
-        )
-        self.vae_temporal_compression_ratio = (
-            self.vae.temporal_compression_ratio if getattr(self, "vae", None) is not None else 8
-        )
-        self.video_processor = VideoProcessor(vae_scale_factor=self.vae_spatial_compression_ratio)
-
-    def prepare_latents(
-        self,
-        video: torch.Tensor | None = None,
-        batch_size: int = 1,
-        num_frames: int = 121,
-        height: int = 512,
-        width: int = 768,
-        spatial_patch_size: int = 1,
-        temporal_patch_size: int = 1,
-        dtype: torch.dtype | None = None,
-        device: torch.device | None = None,
-        generator: torch.Generator | None = None,
-        latents: torch.Tensor | None = None,
-    ) -> torch.Tensor:
-        if latents is not None:
-            if latents.ndim == 3:
-                # Convert token seq [B, S, D] to latent video [B, C, F, H, W]
-                latent_num_frames = (num_frames - 1) // self.vae_temporal_compression_ratio + 1
-                latent_height = height // self.vae_spatial_compression_ratio
-                latent_width = width // self.vae_spatial_compression_ratio
-                latents = self._unpack_latents(
-                    latents, latent_num_frames, latent_height, latent_width, spatial_patch_size, temporal_patch_size
-                )
-            return latents.to(device=device, dtype=dtype)
-
-        video = video.to(device=device, dtype=self.vae.dtype)
-        if isinstance(generator, list):
-            if len(generator) != batch_size:
-                raise ValueError(
-                    f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
-                    f" size of {batch_size}. Make sure the batch size matches the length of the generators."
-                )
-
-            init_latents = [
-                retrieve_latents(self.vae.encode(video[i].unsqueeze(0)), generator[i]) for i in range(batch_size)
-            ]
-        else:
-            init_latents = [retrieve_latents(self.vae.encode(vid.unsqueeze(0)), generator) for vid in video]
-
-        init_latents = torch.cat(init_latents, dim=0).to(dtype)
-        # NOTE: latent upsampler operates on the unnormalized latents, so don't normalize here
-        # init_latents = self._normalize_latents(init_latents, self.vae.latents_mean, self.vae.latents_std)
-        return init_latents
-
-    def adain_filter_latent(self, latents: torch.Tensor, reference_latents: torch.Tensor, factor: float = 1.0):
-        result = latents.clone()
-
-        for i in range(latents.size(0)):
-            for c in range(latents.size(1)):
-                r_sd, r_mean = torch.std_mean(reference_latents[i, c], dim=None)  # index by original dim order
-                i_sd, i_mean = torch.std_mean(result[i, c], dim=None)
-
-                result[i, c] = ((result[i, c] - i_mean) / i_sd) * r_sd + r_mean
-
-        result = torch.lerp(latents, result, factor)
-        return result
-
-    def tone_map_latents(self, latents: torch.Tensor, compression: float) -> torch.Tensor:
-        # Remap [0-1] to [0-0.75] and apply sigmoid compression in one shot
-        scale_factor = compression * 0.75
-        abs_latents = torch.abs(latents)
-
-        # Sigmoid compression: sigmoid shifts large values toward 0.2, small values stay ~1.0
-        # When scale_factor=0, sigmoid term vanishes, when scale_factor=0.75, full effect
-        sigmoid_term = torch.sigmoid(4.0 * scale_factor * (abs_latents - 1.0))
-        scales = 1.0 - 0.8 * scale_factor * sigmoid_term
-
-        filtered = latents * scales
-        return filtered
-
-    @staticmethod
-    # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._denormalize_latents
-    def _denormalize_latents(
-        latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
-    ) -> torch.Tensor:
-        # Denormalize latents across the channel dimension [B, C, F, H, W]
-        latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
-        latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
-        latents = latents * latents_std / scaling_factor + latents_mean
-        return latents
-
-    @staticmethod
-    # Copied from diffusers.pipelines.ltx2.pipeline_ltx2.LTX2Pipeline._unpack_latents
-    def _unpack_latents(
-        latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
-    ) -> torch.Tensor:
-        # Packed latents of shape [B, S, D] (S is the effective video sequence length, D is the effective feature dimensions) # noqa
-        # are unpacked and reshaped into a video tensor of shape [B, C, F, H, W]. This is the inverse operation of
-        # what happens in the `_pack_latents` method.
-        batch_size = latents.size(0)
-        latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
-        latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
-        return latents
-
-    def check_inputs(self, video, height, width, latents, tone_map_compression_ratio):
-        if height % self.vae_spatial_compression_ratio != 0 or width % self.vae_spatial_compression_ratio != 0:
-            raise ValueError(f"`height` and `width` have to be divisible by 32 but are {height} and {width}.")
-
-        if video is not None and latents is not None:
-            raise ValueError("Only one of `video` or `latents` can be provided.")
-        if video is None and latents is None:
-            raise ValueError("One of `video` or `latents` has to be provided.")
-
-        if not (0 <= tone_map_compression_ratio <= 1):
-            raise ValueError("`tone_map_compression_ratio` must be in the range [0, 1]")
-
-    def forward(
-        self,
-        video: list[PipelineImageInput] | None = None,
-        height: int = 512,
-        width: int = 768,
-        num_frames: int = 121,
-        spatial_patch_size: int = 1,
-        temporal_patch_size: int = 1,
-        latents: torch.Tensor | None = None,
-        latents_normalized: bool = False,
-        decode_timestep: float | list[float] = 0.0,
-        decode_noise_scale: float | list[float] | None = None,
-        adain_factor: float = 0.0,
-        tone_map_compression_ratio: float = 0.0,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        output_type: str | None = "pil",
-        return_dict: bool = True,
-    ):
-        self.check_inputs(
-            video=video,
-            height=height,
-            width=width,
-            latents=latents,
-            tone_map_compression_ratio=tone_map_compression_ratio,
-        )
-
-        if video is not None:
-            # Batched video input is not yet tested/supported. TODO: take a look later
-            batch_size = 1
-        else:
-            batch_size = latents.shape[0]
-        device = self.device
-
-        if video is not None:
-            num_frames = len(video)
-            if num_frames % self.vae_temporal_compression_ratio != 1:
-                num_frames = num_frames // self.vae_temporal_compression_ratio * self.vae_temporal_compression_ratio + 1
-                video = video[:num_frames]
-                logger.warning(
-                    f"Video length expected to be of the form `k * {self.vae_temporal_compression_ratio} + 1` but is {len(video)}. Truncating to {num_frames} frames."  # noqa
-                )
-            video = self.video_processor.preprocess_video(video, height=height, width=width)
-            video = video.to(device=device, dtype=torch.float32)
-
-        latents_supplied = latents is not None
-        latents = self.prepare_latents(
-            video=video,
-            batch_size=batch_size,
-            num_frames=num_frames,
-            height=height,
-            width=width,
-            spatial_patch_size=spatial_patch_size,
-            temporal_patch_size=temporal_patch_size,
-            dtype=torch.float32,
-            device=device,
-            generator=generator,
-            latents=latents,
-        )
-
-        if latents_supplied and latents_normalized:
-            latents = self._denormalize_latents(
-                latents, self.vae.latents_mean, self.vae.latents_std, self.vae.config.scaling_factor
-            )
-        latents = latents.to(self.latent_upsampler.dtype)
-        latents_upsampled = self.latent_upsampler(latents)
-
-        if adain_factor > 0.0:
-            latents = self.adain_filter_latent(latents_upsampled, latents, adain_factor)
-        else:
-            latents = latents_upsampled
-
-        if tone_map_compression_ratio > 0.0:
-            latents = self.tone_map_latents(latents, tone_map_compression_ratio)
-
-        if output_type == "latent":
-            video = latents
-        else:
-            if not self.vae.config.timestep_conditioning:
-                timestep = None
-            else:
-                noise = randn_tensor(latents.shape, generator=generator, device=device, dtype=latents.dtype)
-                if not isinstance(decode_timestep, list):
-                    decode_timestep = [decode_timestep] * batch_size
-                if decode_noise_scale is None:
-                    decode_noise_scale = decode_timestep
-                elif not isinstance(decode_noise_scale, list):
-                    decode_noise_scale = [decode_noise_scale] * batch_size
-
-                timestep = torch.tensor(decode_timestep, device=device, dtype=latents.dtype)
-                decode_noise_scale = torch.tensor(decode_noise_scale, device=device, dtype=latents.dtype)[
-                    :, None, None, None, None
-                ]
-                latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
-
-            video = self.vae.decode(latents, timestep, return_dict=False)[0]
-            video = self.video_processor.postprocess_video(video, output_type=output_type)
-
-        if not return_dict:
-            return (video,)
-
-        return DiffusionOutput(output=(video,))

vllm_omni/diffusion/models/ltx2/pipeline_ltx2_two_stage.py

@@ -1,207 +1,27 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
-"""Two-stage entry points for the LTX model family."""
+"""Distilled two-stage entry points for the LTX model family."""
 
-from __future__ import annotations
+from vllm_omni.diffusion.data import OmniDiffusionConfig
 
-import os
-from dataclasses import replace
-from typing import Any, ClassVar
-
-import torch
-from diffusers.pipelines.ltx2.utils import DISTILLED_SIGMA_VALUES, STAGE_2_DISTILLED_SIGMA_VALUES
-
-from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig
-from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch
-
-from .ltx2_components import (
-    LTX2_COMPONENT_PROFILE,
-)
+from .ltx2_components import LTX2_DISTILLED_COMPONENT_PROFILE
 from .ltx2_components import (
     get_ltx2_post_process_func as get_ltx2_post_process_func,  # noqa: F401
 )
-from .ltx2_conditioning import LTXI2VConditioningMixin
-from .ltx2_guidance import LTX_LEGACY_VELOCITY_GUIDANCE
-from .ltx2_pipeline_runtime import LTXPipelineRuntime
-from .ltx2_recipes import LTX2_ONE_STAGE_RECIPE
-from .ltx2_request import LTXRequestInputs
-from .pipeline_ltx2_latent_upsample import LTX2LatentUpsamplePipeline
+from .ltx2_recipes import LTX2_DISTILLED_TWO_STAGE_RECIPE
+from .pipeline_ltx2 import LTX2Pipeline
 
 
-class LTX2TwoStagesPipeline(LTXPipelineRuntime):
-    """Legacy distilled-only LTX2 two-stage compatibility entry."""
+class LTX2DistilledPipeline(LTX2Pipeline):
+    """Unified LTX-2 distilled two-stage T2V/I2V entry."""
 
-    component_profile = LTX2_COMPONENT_PROFILE
-    guidance_strategy = LTX_LEGACY_VELOCITY_GUIDANCE
-    one_stage_recipe = LTX2_ONE_STAGE_RECIPE
+    pipeline_kind = "distilled_two_stage"
+    component_profile = LTX2_DISTILLED_COMPONENT_PROFILE
+    pipeline_recipe = LTX2_DISTILLED_TWO_STAGE_RECIPE
     supports_request_batch = False
-    support_image_input = False
-
-    _dit_modules: ClassVar[list[str]] = list(component_profile.dit_modules)
-    _encoder_modules: ClassVar[list[str]] = list(component_profile.encoder_modules)
-    _vae_modules: ClassVar[list[str]] = list(component_profile.vae_modules)
-    _resident_modules: ClassVar[list[str]] = list(component_profile.resident_modules)
+    support_image_input = True
+    unified_text_image_entry = True
 
     def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        model_path = od_config.model
-        self.distilled = "distilled" in os.path.basename(os.path.normpath(model_path))
-        if not self.distilled:
-            raise NotImplementedError(f"{model_path} is not supported for {self.__class__.__name__}.")
-
         super().__init__(od_config=od_config, prefix=prefix)
-        self.upsample_pipe = LTX2LatentUpsamplePipeline(vae=self.vae, od_config=od_config)
-
-    def _run_two_stage(
-        self,
-        req: DiffusionRequestBatch,
-        request_inputs: LTXRequestInputs,
-        *,
-        noise_scale: float,
-        timesteps: list[int] | None,
-        attention_kwargs: dict[str, Any] | None,
-        image: Any | None = None,
-    ) -> DiffusionOutput | list[DiffusionOutput]:
-        stage1 = self.run_phase(
-            req,
-            request_inputs,
-            noise_scale=noise_scale,
-            sigmas=DISTILLED_SIGMA_VALUES if self.distilled else None,
-            timesteps=timesteps,
-            attention_kwargs=attention_kwargs,
-            image=image,
-        )
-        upscaled_video_latent = self.upsample_pipe(
-            latents=stage1.video,
-            output_type="latent",
-            return_dict=False,
-        )[0]
-
-        stage2_inputs = replace(
-            request_inputs,
-            num_inference_steps=3,
-            guidance_scale=1.0,
-            latents=upscaled_video_latent,
-            audio_latents=stage1.audio,
-            decode_timestep=0.0,
-            decode_noise_scale=None,
-            output_type="np",
-        )
-        stage2 = self.run_phase(
-            req,
-            stage2_inputs,
-            noise_scale=STAGE_2_DISTILLED_SIGMA_VALUES[0],
-            sigmas=STAGE_2_DISTILLED_SIGMA_VALUES,
-            timesteps=None,
-            attention_kwargs=attention_kwargs,
-            prompt_context=stage1.forward_context.prompt_context,
-        )
-        return self.decode_phase(stage2)
-
-    @torch.no_grad()
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_frames: int | None = None,
-        frame_rate: float | None = None,
-        num_inference_steps: int | None = None,
-        timesteps: list[int] | None = None,
-        guidance_scale: float = 4.0,
-        guidance_rescale: float = 0.0,
-        noise_scale: float = 0.0,
-        num_videos_per_prompt: int | None = 1,
-        generator: torch.Generator | list[torch.Generator] | None = None,
-        latents: torch.Tensor | None = None,
-        audio_latents: torch.Tensor | None = None,
-        prompt_embeds: torch.Tensor | None = None,
-        negative_prompt_embeds: torch.Tensor | None = None,
-        prompt_attention_mask: torch.Tensor | None = None,
-        negative_prompt_attention_mask: torch.Tensor | None = None,
-        decode_timestep: float | list[float] = 0.0,
-        decode_noise_scale: float | list[float] | None = None,
-        output_type: str = "np",
-        return_dict: bool = True,
-        attention_kwargs: dict[str, Any] | None = None,
-        max_sequence_length: int | None = None,
-        *,
-        image: Any | None = None,
-    ) -> DiffusionOutput | list[DiffusionOutput]:
-        del return_dict
-        request_inputs = self._resolve_request_inputs(
-            req,
-            prompt=prompt,
-            negative_prompt=negative_prompt,
-            height=height,
-            width=width,
-            num_frames=num_frames,
-            frame_rate=frame_rate,
-            num_inference_steps=num_inference_steps,
-            timesteps=timesteps,
-            guidance_scale=(
-                getattr(getattr(self, "one_stage_recipe", None), "guidance_scale", 4.0)
-                if guidance_scale is None
-                else guidance_scale
-            ),
-            guidance_rescale=guidance_rescale,
-            num_videos_per_prompt=num_videos_per_prompt,
-            generator=generator,
-            latents=latents,
-            audio_latents=audio_latents,
-            prompt_embeds=prompt_embeds,
-            negative_prompt_embeds=negative_prompt_embeds,
-            prompt_attention_mask=prompt_attention_mask,
-            negative_prompt_attention_mask=negative_prompt_attention_mask,
-            decode_timestep=decode_timestep,
-            decode_noise_scale=decode_noise_scale,
-            output_type=output_type,
-            max_sequence_length=max_sequence_length,
-        )
-        if self.support_image_input:
-            image = self._resolve_request_image(req, image, request_inputs)
-        return self._run_two_stage(
-            req,
-            request_inputs,
-            noise_scale=noise_scale,
-            timesteps=timesteps,
-            attention_kwargs=attention_kwargs,
-            image=image,
-        )
-
-
-class LTX2ImageToVideoTwoStagesPipeline(LTXI2VConditioningMixin, LTX2TwoStagesPipeline):
-    """LTX2 two-stage image-to-video entry."""
-
-    def forward(
-        self,
-        req: DiffusionRequestBatch,
-        image: Any | None = None,
-        prompt: str | list[str] | None = None,
-        negative_prompt: str | list[str] | None = None,
-        height: int | None = None,
-        width: int | None = None,
-        num_frames: int | None = None,
-        frame_rate: float | None = None,
-        num_inference_steps: int | None = None,
-        sigmas: list[float] | None = None,
-        *args: Any,
-        **kwargs: Any,
-    ) -> DiffusionOutput | list[DiffusionOutput]:
-        """Preserve the legacy I2V sigma slot, which was ignored by this pipeline."""
-        del sigmas
-        return super().forward(
-            req,
-            image,
-            prompt,
-            negative_prompt,
-            height,
-            width,
-            num_frames,
-            frame_rate,
-            num_inference_steps,
-            *args,
-            **kwargs,
-        )

vllm_omni/diffusion/registry.py

@@ -71,20 +71,10 @@ _DIFFUSION_MODELS = {
         "pipeline_ltx2",
         "LTX2Pipeline",
     ),
-    "LTX2ImageToVideoPipeline": (
-        "ltx2",
-        "pipeline_ltx2",
-        "LTX2ImageToVideoPipeline",
-    ),
-    "LTX2TwoStagesPipeline": (
+    "LTX2DistilledPipeline": (
         "ltx2",
         "pipeline_ltx2_two_stage",
-        "LTX2TwoStagesPipeline",
-    ),
-    "LTX2ImageToVideoTwoStagesPipeline": (
-        "ltx2",
-        "pipeline_ltx2_two_stage",
-        "LTX2ImageToVideoTwoStagesPipeline",
+        "LTX2DistilledPipeline",
     ),
     "LTX2T2VDMD2Pipeline": (
         "ltx2",
@@ -96,16 +86,6 @@ _DIFFUSION_MODELS = {
         "pipeline_ltx2",
         "LTX2I2VDMD2Pipeline",
     ),
-    "LTX23Pipeline": (
-        "ltx2",
-        "pipeline_ltx2",
-        "LTX23Pipeline",
-    ),
-    "LTX23ImageToVideoPipeline": (
-        "ltx2",
-        "pipeline_ltx2",
-        "LTX23ImageToVideoPipeline",
-    ),
     "StableAudioPipeline": (
         "stable_audio",
         "pipeline_stable_audio",
@@ -449,7 +429,7 @@ def _apply_sequence_parallel_if_enabled(model, od_config: OmniDiffusionConfig) -
 
         for attr in transformer_attrs:
             if not hasattr(model, attr):
-                # Some pipeline like LTX2TwoStagesPipeline have recursive
+                # Some pipelines have recursive
                 # modules that have the transformer
                 module = find_module_with_attr(model, attr)
                 if module is None:
@@ -519,13 +499,9 @@ _DIFFUSION_POST_PROCESS_FUNCS = {
     "WanPipeline": "get_wan22_post_process_func",
     "WanVACEPipeline": "get_wan22_vace_post_process_func",
     "LTX2Pipeline": "get_ltx2_post_process_func",
-    "LTX2TwoStagesPipeline": "get_ltx2_post_process_func",
-    "LTX2ImageToVideoPipeline": "get_ltx2_post_process_func",
-    "LTX2ImageToVideoTwoStagesPipeline": "get_ltx2_post_process_func",
+    "LTX2DistilledPipeline": "get_ltx2_post_process_func",
     "LTX2T2VDMD2Pipeline": "get_ltx2_post_process_func",
     "LTX2I2VDMD2Pipeline": "get_ltx2_post_process_func",
-    "LTX23Pipeline": "get_ltx2_post_process_func",
-    "LTX23ImageToVideoPipeline": "get_ltx2_post_process_func",
     "StableAudioPipeline": "get_stable_audio_post_process_func",
     "SoulXSingerPipeline": "get_soulxsinger_post_process_func",
     "SoulXSingerSVCPipeline": "get_soulxsinger_post_process_func",