Revert "[BugFix] Fix 'GPUARWorker' attribute error (#5478)"

This reverts commit f4640d9a2e.

.buildkite/cuda/test-merge.yml

@@ -49,13 +49,17 @@ steps:
     depends_on: upload-merge-pipeline
     steps:
       - label: "Entrypoint Test with H100"
-        timeout_in_minutes: 60
+        # TEMPORARY 60 -> 90: engine startup/sleep-wake is ~13-15 min slower per
+        # run on the vLLM 0.26 image (model loading ~2x, init ~2x; see builds
+        # 2915 vs 2916), which pushed this suite past 60 min. Revert once the
+        # 0.26 startup regression is understood and fixed.
+        timeout_in_minutes: 90
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and H100" --run-level "advanced_model"
         mirror_hardwares: h100_2
 
       - label: "Entrypoint Test with L4"
-        timeout_in_minutes: 30
+        timeout_in_minutes: 60
         commands:
           - pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
         mirror_hardwares: l4_1

.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.ci

@@ -1,7 +1,7 @@
-# Target vllm version: v0.25.0 (releases/v0.25.0 branch).
+# Target vllm version: v0.26.0 (releases/v0.26.0 branch).
 # Keep VLLM_BASE_TAG aligned with the target vLLM release on every rebase.
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.25.0
+ARG VLLM_BASE_TAG=v0.26.0
 
 # We use registry-based layer caching for the dependencies,
 # where the tag is determined based on the hash of the contents

docker/Dockerfile.cuda

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.25.0
+ARG BASE_IMAGE=vllm/vllm-openai:v0.26.0
 FROM ${BASE_IMAGE}
 
 ARG COMMON_WORKDIR=/app

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" \

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.25.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.26.0
 FROM ${BASE_IMAGE} AS base
 
 # Declare a variable to know if we want to use the nightly build or the stable build.
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
 #    we should swap over to use stable release ASAP.
 #    We should avoid relying on custom commits.
 ARG USE_NIGHTLY_BUILD=0
-ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
+ARG VLLM_VERSION_OR_COMMIT_HASH=v0.26.0
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 

docker/Dockerfile.xpu

@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
 ENV UV_INDEX_STRATEGY="unsafe-best-match"
 ENV UV_LINK_MODE="copy"
 
-ARG VLLM_VERSION=v0.25.0
+ARG VLLM_VERSION=v0.26.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/getting_started/installation/gpu.md

@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
-Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, and 0.25.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, 0.23.0, 0.24.0, 0.25.0, and 0.26.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
 
 === "NVIDIA CUDA"
 

docs/getting_started/installation/gpu/cuda.inc.md

@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.25.0 --torch-backend=auto
+uv pip install vllm==0.26.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
 
 ```bash
-uv pip install vllm==0.25.0 --torch-backend=auto
+uv pip install vllm==0.26.0 --torch-backend=auto
 ```
 
-The 0.25.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.26.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.25.0
+git checkout v0.26.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0 (the default for v0.25.0; the wheel filename has no `+cu130` suffix)
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.0-cp38-abi3-manylinux_2_28_x86_64.whl
+#For CUDA 13.0 (the default for v0.26.0; the wheel filename has no `+cu130` suffix)
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.26.0/vllm-0.26.0-cp38-abi3-manylinux_2_28_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.25.0 \
+    vllm/vllm-omni:v0.26.0 \
     vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 

docs/getting_started/installation/gpu/rocm.inc.md

@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
+uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.26.0 release version of the library
 
 ```bash
-uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
+uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
 ```
 
-The pre-built 0.25.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The pre-built 0.26.0 vLLM wheel targets ROCm 7.2.3. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.25.0
+git checkout v0.26.0
 python3 -m pip install -r requirements/rocm.txt
 python3 setup.py develop
 ```
@@ -130,7 +130,7 @@ docker run --rm \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   -p 8091:8091 \
-  vllm/vllm-omni-rocm:v0.25.0 \
+  vllm/vllm-omni-rocm:v0.26.0 \
   --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
@@ -149,7 +149,7 @@ docker run --rm -it \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
   --entrypoint bash \
-  vllm/vllm-omni-rocm:v0.25.0
+  vllm/vllm-omni-rocm:v0.26.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.25.0 --torch-backend=auto
+uv pip install vllm==0.26.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723
+uv pip install vllm==0.26.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.26.0/rocm723
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -34,7 +34,7 @@ For additional installation methods — please see the [installation guide](inst
 !!! note
     It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
 
-    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.25.0` and vLLM Omni `0.25.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
+    If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.26.0` and vLLM Omni `0.26.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
 
 ## Offline Inference
 

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/attention/test_fish_kvcache_attn.py

@@ -100,8 +100,9 @@ def _decode_kv_cache(
     head_size: int = 128,
     dtype: torch.dtype = torch.float16,
 ) -> torch.Tensor:
-    # vLLM >=0.23.0 KV cache layout: key/value live on dim 1.
-    return torch.zeros((num_blocks, 2, block_size, num_kv_heads, head_size), dtype=dtype)
+    # vLLM >=0.23.0 KV cache layout: (num_blocks, num_kv_heads, block_size,
+    # 2*head_size) with K and V packed in the last dimension.
+    return torch.zeros((num_blocks, num_kv_heads, block_size, 2 * head_size), dtype=dtype)
 
 
 def test_fish_kvcache_enabled_by_default(monkeypatch):
@@ -221,10 +222,11 @@ def _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size):
 
 
 def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch):
-    # vLLM >=0.23.0 lays the KV cache out as (num_blocks, 2, block_size,
-    # num_kv_heads, head_size). The backend must unbind on dim 1 so each of
-    # key/value comes back as the 4-D (num_blocks, block_size, num_kv_heads,
-    # head_size) tensor the decode kernel expects.
+    # vLLM >=0.23.0 lays out the KV cache as (num_blocks, num_kv_heads,
+    # block_size, 2*head_size) with K and V packed in the last dimension.
+    # The backend must split on the content dim so each of key/value comes
+    # back as the 4-D (num_blocks, block_size, num_kv_heads, head_size)
+    # tensor the decode kernel expects.
     monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
     monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
     monkeypatch.setattr(fish_kvcache_backend, "load_error", lambda: None)
@@ -232,7 +234,8 @@ def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch):
 
     num_blocks, block_size, num_kv_heads, head_size = 8, 16, 8, 128
     cache_shape = _vllm_kv_cache_shape(num_blocks, block_size, num_kv_heads, head_size)
-    assert cache_shape == (num_blocks, 2, block_size, num_kv_heads, head_size)
+    # Upstream now returns a 4-D packed shape: (B, H, N, 2*D).
+    assert cache_shape == (num_blocks, num_kv_heads, block_size, 2 * head_size)
 
     captured = {}
 
@@ -262,20 +265,21 @@ def test_fish_kvcache_backend_unbinds_kv_on_vllm_cache_layout(monkeypatch):
 def test_fish_kvcache_vllm_cache_unbinds_then_falls_back_until_contiguous(monkeypatch):
     # End-to-end with the real vLLM cache layout and the real (unmocked) guard.
     #
-    # The point of the unbind(1) fix: unbind(0) raised "too many values to
-    # unpack" on the 5-D vLLM 0.23.0 layout, so the request crashed outright.
-    # unbind(1) makes the request *work* again -- it unpacks cleanly and the
-    # backend gracefully falls back to the original attention with correct
-    # output. This test guards that fix and must NOT depend on the fast path
-    # firing.
+    # The point of the transpose+split fix: unbind(1) on the old 5-D layout
+    # raised "too many values to unpack" on the vLLM 0.23.0 layout, and the
+    # upstream now packs K/V into the last dimension (4-D shape). The
+    # transpose(1, 2) + split correctly unpacks into separate key/value tensors
+    # and the backend gracefully falls back to the original attention with
+    # correct output. This test guards that fix and must NOT depend on the fast
+    # path firing.
     #
-    # It also documents a known limitation (tracked separately): unbind(1) on
-    # the interleaved (num_blocks, 2, ...) layout yields *non-contiguous*
-    # key/value views, which the guard rejects on the is_contiguous() check.
-    # So the dimensions are compatible (4-D, block_size=16, head_size=128) --
-    # the only thing keeping the fast path from engaging is contiguity, not
-    # shape. Making the fast path actually fire would require a .contiguous()
-    # copy and is out of scope here.
+    # It also documents a known limitation (tracked separately):
+    # transpose(1, 2) on the interleaved (B, H, N, 2*D) layout yields
+    # *non-contiguous* key/value views, which the guard rejects on the
+    # is_contiguous() check. So the dimensions are compatible (4-D,
+    # block_size=16, head_size=128) -- the only thing keeping the fast path
+    # from engaging is contiguity, not shape. Making the fast path actually
+    # fire would require a .contiguous() copy and is out of scope here.
     fish_kvcache_backend.reset_fish_kvcache_attn_stats()
     monkeypatch.setenv("VLLM_OMNI_FISH_KVCACHE_ATTN", "1")
     monkeypatch.setattr(fish_kvcache_backend, "is_available", lambda: True)
@@ -300,13 +304,14 @@ def test_fish_kvcache_vllm_cache_unbinds_then_falls_back_until_contiguous(monkey
     output = torch.zeros_like(query)
     kv_cache = torch.zeros(cache_shape, dtype=torch.float16)
 
-    # Root cause record: unbind(1) gives the right 4-D shape but a non-contiguous view.
-    key_cache, value_cache = kv_cache.unbind(1)
+    # Root cause record: transpose(1,2) + split gives the right 4-D shape but
+    # a non-contiguous view.
+    key_cache, value_cache = kv_cache.transpose(1, 2).split(head_size, dim=-1)
     assert tuple(key_cache.shape) == (num_blocks, block_size, num_kv_heads, head_size)
     assert not key_cache.is_contiguous()
 
-    # unbind(1) no longer crashes (unbind(0) would have raised here) and the
-    # request completes via the correct fallback path.
+    # transpose+split no longer crashes (unbind(1) would have raised here) and
+    # the request completes via the correct fallback path.
     result = model.layers[0].self_attn.attn.impl.forward(None, query, None, None, kv_cache, _metadata(), output)
 
     assert result is output

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,182 @@
+"""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_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)
+
+
+@pytest.mark.skip(reason="#issue 5484")
+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)
+        return 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_ar_scheduler_streaming.py

@@ -73,6 +73,9 @@ def _run_resumable_segment_stop(
 
     sched._update_request_with_output.side_effect = stop_request
     sched._handle_stopped_request.return_value = session_finished
+    # vLLM 0.26 returns (kv_xfer_params, ec_xfer_params); an unconfigured
+    # MagicMock iterates empty and fails to unpack at the call site.
+    sched._free_request.return_value = (None, None)
     sched.chunk_transfer_adapter = None
     sched.running = [session]
     sched.waiting_for_transfer_free = set()
@@ -137,6 +140,23 @@ def test_resumable_session_terminal_is_not_marked_as_segment_boundary() -> None:
     assert output.is_segment_finished is False
 
 
+def test_update_from_output_settles_in_flight_tokens() -> None:
+    """vLLM 0.26: schedule() increments num_in_flight_tokens per scheduled
+    token; update_from_output must decrement it symmetrically. If the
+    decrement is dropped the counter grows monotonically and both readers
+    (allocate_slots, _connector_finished) clamp
+    max(0, num_computed_tokens - num_in_flight_tokens) to zero forever,
+    silently freezing sliding-window block freeing.
+    """
+    session = _make_request()
+    session.status = RequestStatus.RUNNING
+    session.num_in_flight_tokens = 1  # as left by schedule() for this step
+
+    _run_resumable_segment_stop(session)
+
+    assert session.num_in_flight_tokens == 0
+
+
 def test_running_decode_step_without_inter_stage_payload_does_not_raise() -> None:
     """A decode step that neither stops nor carries an inter-stage payload.
 

tests/core/sched/test_omni_scheduler_finish_requests_purge.py

@@ -76,8 +76,9 @@ def test_finish_requests_purges_finished_request_from_running(
     )
 
     def fake_finish_requests(self, request_ids, finished_status):
-        # super() returning the abort tuple but leaving self.running untouched
-        return [(finished.request_id, 0)]
+        # vLLM 0.26 super() returns the finished Request objects but here
+        # leaves self.running untouched to simulate the missed removal.
+        return [finished]
 
     monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
 
@@ -88,7 +89,7 @@ def test_finish_requests_purges_finished_request_from_running(
     )
 
     assert scheduler.running == []
-    assert result == [(finished.request_id, 0)]
+    assert result == [finished]
 
 
 @pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
@@ -111,7 +112,7 @@ def test_finish_requests_purges_untracked_request_from_running(
     )
 
     def fake_finish_requests(self, request_ids, finished_status):
-        return [(untracked.request_id, 0)]
+        return [untracked]
 
     monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
 
@@ -122,7 +123,7 @@ def test_finish_requests_purges_untracked_request_from_running(
     )
 
     assert scheduler.running == []
-    assert result == [(untracked.request_id, 0)]
+    assert result == [untracked]
 
 
 @pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
@@ -136,7 +137,7 @@ def test_finish_requests_leaves_healthy_running_intact(
     is conservative and won't aggressively drop healthy entries.
     """
     alive = _StubRequest("req-alive", RequestStatus.RUNNING)
-    leaving_id = "req-leaving"
+    leaving = _StubRequest("req-leaving", RequestStatus.FINISHED_STOPPED)
     scheduler = _make_scheduler(
         scheduler_cls,
         requests={alive.request_id: alive},
@@ -145,13 +146,13 @@ def test_finish_requests_leaves_healthy_running_intact(
     )
 
     def fake_finish_requests(self, request_ids, finished_status):
-        return [(leaving_id, 0)]
+        return [leaving]
 
     monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
 
     scheduler_cls.finish_requests(
         scheduler,
-        [leaving_id],
+        [leaving.request_id],
         RequestStatus.FINISHED_STOPPED,
     )
 

tests/core/sched/test_omni_scheduler_input_coordinator_cleanup.py

@@ -53,7 +53,11 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
     def fake_finish_requests(self, request_ids, finished_status):
         assert request_ids == ["req-a", "req-b"]
         assert finished_status == RequestStatus.FINISHED_STOPPED
-        return [("req-a", 0), ("req-b", 1)]
+        # vLLM 0.26 returns the finished Request objects, not (req_id, client_index) tuples.
+        return [
+            SimpleNamespace(request_id="req-a", client_index=0),
+            SimpleNamespace(request_id="req-b", client_index=1),
+        ]
 
     monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
 
@@ -63,7 +67,7 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
         RequestStatus.FINISHED_STOPPED,
     )
 
-    assert result == [("req-a", 0), ("req-b", 1)]
+    assert [(r.request_id, r.client_index) for r in result] == [("req-a", 0), ("req-b", 1)]
     assert coordinator.freed == ["req-a", "req-b"]
 
 
@@ -82,7 +86,7 @@ def test_ar_free_request_cleans_input_coordinator_on_normal_free() -> None:
 
     result = OmniARScheduler._free_request(scheduler, DummyRequest())
 
-    assert result is None
+    assert result == (None, None)
     assert coordinator.freed == ["req-free"]
     assert scheduler._omits_kv_transfer_cache == {}
     assert scheduler._new_prompt_len_snapshot == {}

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/boogu_image/test_boogu_image_transformer.py

@@ -1,6 +1,7 @@
 # SPDX-License-Identifier: Apache-2.0
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
+import logging
 import os
 from types import SimpleNamespace
 
@@ -338,27 +339,43 @@ def test_transformer_load_weights_round_trip():
         assert torch.allclose(reloaded[native_name], expected)
 
 
-def test_transformer_load_weights_warns_for_unexpected_and_unloaded(caplog):
+def test_transformer_load_weights_warns_for_unexpected_and_unloaded():
     from vllm_omni.diffusion.models.boogu_image.boogu_image_transformer import (
         BooguImageTransformer2DModel,
     )
 
-    model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
-    native_params = dict(model.named_parameters())
-    loaded_name = next(iter(native_params))
-    checkpoint_name = _native_to_checkpoint_name(loaded_name)
+    # vllm_omni's logger hierarchy hangs off vLLM's `vllm` logger, which sets
+    # propagate=False, so records never reach the root logger that pytest's
+    # caplog listens on. Capture at the emitting logger instead.
+    messages: list[str] = []
 
-    loaded = model.load_weights(
-        [
-            (checkpoint_name, torch.randn_like(native_params[loaded_name])),
-            ("unexpected.weight", torch.ones(1)),
-        ]
-    )
+    class _Capture(logging.Handler):
+        def emit(self, record: logging.LogRecord) -> None:
+            messages.append(record.getMessage())
 
+    log = logging.getLogger("vllm_omni.diffusion.models.boogu_image.boogu_image_transformer")
+    handler = _Capture(level=logging.WARNING)
+    log.addHandler(handler)
+    try:
+        model = BooguImageTransformer2DModel(od_config=_tiny_od_config())
+        native_params = dict(model.named_parameters())
+        loaded_name = next(iter(native_params))
+        checkpoint_name = _native_to_checkpoint_name(loaded_name)
+
+        loaded = model.load_weights(
+            [
+                (checkpoint_name, torch.randn_like(native_params[loaded_name])),
+                ("unexpected.weight", torch.ones(1)),
+            ]
+        )
+    finally:
+        log.removeHandler(handler)
+
+    text = "\n".join(messages)
     assert loaded == {loaded_name}
-    assert "Skipping unexpected checkpoint weight unexpected.weight" in caplog.text
-    assert "Model parameters not loaded from checkpoint" in caplog.text
-    assert next(name for name in native_params if name != loaded_name) in caplog.text
+    assert "Skipping unexpected checkpoint weight unexpected.weight" in text
+    assert "Model parameters not loaded from checkpoint" in text
+    assert next(name for name in native_params if name != loaded_name) in text
 
 
 def test_transformer_forward_t2i_shape():

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/distributed/omni_connectors/test_chunk_transfer_adapter.py

@@ -252,6 +252,37 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
     assert len(adapter._pending_save_reqs) == 1
 
 
+def test_send_single_request_terminal_chunk_still_flushes_processor(build_adapter, monkeypatch):
+    """A terminal stop is not a segment boundary (#5383), but the producer-side
+    processor must still receive the flush signal on the terminal chunk.
+    Passing only ``is_segment_finished`` starved processors of their final
+    accumulated payload once terminal stops stopped setting it (#5413: the
+    downstream stage got ``meta.finished`` with the tail data missing).
+    """
+    adapter, connector = build_adapter(stage_id=0)
+    request = _req("req-terminal", RequestStatus.FINISHED_STOPPED, external_req_id="ext-terminal")
+
+    seen_flush_flags = []
+
+    def recording_processor(**kwargs):
+        seen_flush_flags.append(kwargs["is_finished"])
+        return OmniPayloadStruct(
+            codes=CodesStruct(audio=torch.tensor([1, 2, 3], dtype=torch.long)),
+        )
+
+    adapter.custom_process_next_stage_input_func = recording_processor
+    monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
+
+    adapter._send_single_request(
+        {"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": False}
+    )
+
+    assert seen_flush_flags == [True]
+    sent_payload = connector.put.call_args.kwargs["data"]
+    assert bool(sent_payload.meta.finished.item()) is True
+    assert bool(sent_payload.meta.is_segment_finished.item()) is False
+
+
 def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, monkeypatch):
     """Producer may return a struct with ``meta=None`` (e.g. payload that
     carries only ``embed`` or ``codes``). The sender's ``meta is not None``
@@ -780,6 +811,10 @@ def test_restore_queues_skips_requests_missing_from_scheduler_requests(build_ada
 class _HashableRequest(SimpleNamespace):
     """SimpleNamespace that can be added to a set (needed by scheduler internals)."""
 
+    # vLLM 0.26: update_from_output settles this counter for every scheduled
+    # request; real Requests initialise it to 0 (vllm/v1/request.py).
+    num_in_flight_tokens = 0
+
     def __hash__(self):
         return hash(self.request_id)
 
@@ -833,7 +868,7 @@ def test_generation_scheduler_calls_cleanup_on_finished(monkeypatch, mocker: Moc
     scheduler.requests = {"req-s1": request}
 
     scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
-    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
     scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
     scheduler.running = [request]
     scheduler.waiting = mocker.MagicMock()
@@ -918,7 +953,7 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF
     scheduler._update_request_with_output = mocker.MagicMock(return_value=([], True))
     scheduler._process_kv_transfer_trigger = mocker.MagicMock(return_value=False)
     scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
-    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
     scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
     scheduler.running = [request]
     scheduler.waiting = mocker.MagicMock()
@@ -1047,7 +1082,7 @@ def _build_deferred_finish_scheduler(mocker, *, running, pending_finish_reqs):
     scheduler._pending_finish_reqs = list(pending_finish_reqs)
 
     scheduler._handle_stopped_request = mocker.MagicMock(return_value=True)
-    scheduler._free_request = mocker.MagicMock(return_value=None)
+    scheduler._free_request = mocker.MagicMock(return_value=(None, None))
     scheduler._get_routed_experts = mocker.MagicMock(return_value=None)
     scheduler.running = list(running)
     scheduler.waiting = mocker.MagicMock()
@@ -1103,7 +1138,7 @@ def test_deferred_finish_emits_finished_output(mocker: MockerFixture):
         running=[request],
         pending_finish_reqs=[request],
     )
-    scheduler._free_request.return_value = {"mock": "kv_params"}
+    scheduler._free_request.return_value = ({"mock": "kv_params"}, None)
 
     result = OmniGenerationScheduler.update_from_output(scheduler, sched_out, model_out)
 

tests/distributed/omni_connectors/test_kv_flow.py

@@ -330,6 +330,31 @@ def test_normalize_layer_kv_rejects_invalid_inputs(kv_config, common_constants,
     assert normalized is None
 
 
+def test_normalize_layer_kv_unpacks_packed_layout(common_constants):
+    """The vLLM >= 0.23 packed layout (num_blocks, n_heads, block_size,
+    2*head_dim) is unpacked into (num_blocks, block_size, n_heads, head_dim)
+    key/value blocks when block_size is provided."""
+    block_size = common_constants["block_size"]
+    num_heads = common_constants["num_heads"]
+    head_dim = common_constants["head_dim"]
+    req_id = common_constants["req_id"]
+    num_blocks = 10
+
+    layer_kv = torch.randn(num_blocks, num_heads, block_size, 2 * head_dim)
+
+    # Without block_size the 4-D packed layout is not recognized.
+    assert normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0) is None
+
+    normalized = normalize_layer_kv(layer_kv, req_id=req_id, layer_idx=0, block_size=block_size)
+    assert normalized is not None
+    key_blocks, value_blocks = normalized
+    assert key_blocks.shape == (num_blocks, block_size, num_heads, head_dim)
+    assert value_blocks.shape == (num_blocks, block_size, num_heads, head_dim)
+    reference = layer_kv.transpose(1, 2)
+    assert torch.equal(key_blocks, reference[..., :head_dim])
+    assert torch.equal(value_blocks, reference[..., head_dim:])
+
+
 def test_manager_reception(kv_config, mock_connector, common_constants):
     """Test reception and injection logic in OmniKVTransferManager."""
     num_layers = common_constants["num_layers"]

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/offline_inference/test_voxtral_tts_expansion.py

@@ -58,7 +58,7 @@ def _compose_request(model_name: str, text: str, voice: str) -> dict:
 
 
 @pytest.mark.parametrize("omni_runner_function", [_OMNI_RUNNER_PARAM], indirect=True)
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
 def test_voxtral_tts_offline_basic(omni_runner_function: OmniRunner) -> None:
     """Offline sync path; function-scoped runner so AsyncOmni test does not overlap a live Omni."""
     omni = omni_runner_function.omni

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/e2e/online_serving/test_voxtral_tts_expansion.py

@@ -30,7 +30,7 @@ pytestmark = [
 ]
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
 def test_speech_english_basic(omni_server, openai_client) -> None:
     """Test basic English TTS generation."""
     openai_client.send_audio_speech_request(
@@ -46,7 +46,7 @@ def test_speech_english_basic(omni_server, openai_client) -> None:
     )
 
 
-@hardware_test(res={"cuda": "L4"}, num_cards=1)
+@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=1)
 def test_speech_english_streaming(omni_server, openai_client) -> None:
     """Test basic streaming English TTS generation (PCM via streaming API)."""
     openai_client.send_audio_speech_request(

tests/engine/test_orchestrator.py

@@ -1813,6 +1813,56 @@ async def test_stage_pool_submit_update_refreshes_output_processor_state() -> No
 
 
 @pytest.mark.asyncio
+async def test_handle_streaming_update_unknown_request_is_dropped() -> None:
+    """A streaming update for an untracked request must not resurrect it.
+
+    Updates always follow the first-chunk add_request through the same
+    ordered submission queue, so an unknown id can only mean the request
+    already finished or was aborted (e.g. the client disconnected).
+    Falling back to add_request created headless sessions that kept
+    cycling through the stages (issue #4271).
+    """
+
+    class RecordingPool:
+        def __init__(self) -> None:
+            self.calls: list[tuple[str, Any]] = []
+
+        async def submit_update(self, request_id, req_state, request, *, prompt_text=None) -> int:
+            self.calls.append((request_id, prompt_text))
+            return 0
+
+    pool = RecordingPool()
+    orchestrator = object.__new__(Orchestrator)
+    orchestrator.async_chunk = False
+    orchestrator.request_states = {}
+    orchestrator.stage_pools = [pool]
+    add_request_calls: list[str] = []
+
+    async def record_add_request(msg) -> None:
+        add_request_calls.append(msg.request_id)
+
+    orchestrator._handle_add_request = record_add_request
+
+    await orchestrator._handle_streaming_update(
+        StageSubmissionMessage(
+            type="streaming_update",
+            request_id="req-unknown",
+            prompt=SimpleNamespace(request_id="req-unknown", prompt_token_ids=[1], resumable=True),
+            original_prompt={"prompt": "segment-2"},
+            output_prompt_text="segment-2",
+            sampling_params_list=[_sampling_params()],
+            final_stage_id=0,
+            preprocess_ms=0.0,
+            request_timestamp=time.time(),
+            enqueue_ts=time.perf_counter(),
+        )
+    )
+
+    assert pool.calls == []
+    assert add_request_calls == []
+    assert "req-unknown" not in orchestrator.request_states
+
+
 async def test_handle_streaming_update_passes_prompt_text_to_stage_pool() -> None:
     class RecordingPool:
         def __init__(self) -> 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/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py

@@ -201,6 +201,7 @@ def _assert_realtime_smoke(result: dict) -> None:
 def _assert_realtime_accuracy(
     result: dict,
     whisper_model_size: str = "large-v3",
+    threshold: float = 0.8,
 ) -> None:
     """Assert that whisper transcription of audio output matches model text.
 
@@ -214,6 +215,12 @@ def _assert_realtime_accuracy(
                    variability even though audio generation was correct. large-v3
                    transcribes these clips reliably, so a failure here now points
                    at the model, not the ASR grader.
+        threshold: Minimum cosine similarity (with length penalty) required to
+                   pass. Default 0.8. Do not lower per-callsite without data:
+                   at 0.35 the assertion no longer detects real audio
+                   regressions. If a variant genuinely needs a different gate
+                   (e.g. whisper partial transcripts under async_chunk), propose
+                   it in its own PR with measurements.
     """
     final_text = (result["transcription_text"] or "").strip()
     assert final_text, "Expected non-empty transcription (model text stream)"
@@ -223,8 +230,9 @@ def _assert_realtime_accuracy(
     assert whisper_text, "Whisper returned empty string for synthesized output audio"
 
     sim = cosine_similarity_text(whisper_text.lower(), final_text.lower())
-    assert sim > 0.8, (
-        f"Output audio transcript should match model text (sim={sim:.3f}): "
+    assert sim > threshold, (
+        f"Output audio transcript should match model text (sim={sim:.3f}, "
+        f"threshold={threshold}): "
         f"whisper={whisper_text!r}, model_text={final_text!r}"
     )
 

tests/entrypoints/test_pd_disaggregation.py

@@ -305,21 +305,11 @@ def mock_get_config(monkeypatch):
     """Auto-mock get_config and related model loading functions."""
     import sys
 
-    fake_tokenizer = _ns()
-    fake_tokenizer.encode = lambda *args, **kwargs: [1, 2, 3]
-    fake_tokenizer.decode = lambda *args, **kwargs: "test"
-
-    def _mock_init_tokenizer_from_configs(model_config=None, **kwargs):
-        return fake_tokenizer
-
-    monkeypatch.setattr(
-        "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs",
-        _mock_init_tokenizer_from_configs,
-        raising=False,
-    )
-    tokenizer_module_path = "vllm.transformers_utils.tokenizer"
-    if tokenizer_module_path in sys.modules:
-        setattr(sys.modules[tokenizer_module_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
+    # TODO(pd-unskip): the tokenizer mock that lived here targeted
+    # vllm.transformers_utils.tokenizer.init_tokenizer_from_configs, which was
+    # deleted in vLLM 0.26 (the raising=False setattr made it a silent no-op).
+    # When the module-level skip is lifted, mock whatever tokenizer seam the
+    # revived PD entrypoint code actually uses.
 
     def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None):
         if prompt_token_ids is not None:
@@ -344,13 +334,6 @@ def mock_get_config(monkeypatch):
             _mock_length_from_prompt_token_ids_or_embeds,
         )
 
-    monkeypatch.setattr(
-        "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs", _mock_init_tokenizer_from_configs, raising=False
-    )
-    async_omni_path = "vllm_omni.entrypoints.async_omni"
-    if async_omni_path in sys.modules:
-        setattr(sys.modules[async_omni_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs)
-
     fake_hf_config = _ns()
     fake_hf_config.model_type = "qwen2_5_omni"
 

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_executor/models/minicpmo_4_5/test_code2wav_model_dir.py

@@ -0,0 +1,88 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Unit tests for MiniCPM-o 4.5 code2wav model-dir resolution (#5442).
+
+In hub/CI deployments ``model_config.model`` is a repo id rather than a local
+directory, so asset lookups must resolve to the downloaded snapshot. The
+resolution must stay lazy: constructing the model with a fake path (as the
+CPU unit tests do) must not touch the hub.
+"""
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import huggingface_hub
+import pytest
+
+from vllm_omni.model_executor.models.minicpmo_4_5.minicpmo_4_5_code2wav import (
+    MiniCPMO45Code2Wav,
+    _resolve_model_dir,
+)
+
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
+
+
+def _no_hub(monkeypatch):
+    def _fail(*args, **kwargs):  # pragma: no cover - must not be reached
+        raise AssertionError("snapshot_download must not be called here")
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", _fail)
+
+
+def test_local_directory_is_returned_unchanged(tmp_path, monkeypatch):
+    _no_hub(monkeypatch)
+    assert _resolve_model_dir(str(tmp_path)) == str(tmp_path)
+
+
+def test_repo_id_resolves_via_snapshot_download(tmp_path, monkeypatch):
+    calls = {}
+
+    def _fake_snapshot_download(model_ref, revision=None, allow_patterns=None):
+        calls["model_ref"] = model_ref
+        calls["revision"] = revision
+        calls["allow_patterns"] = allow_patterns
+        return str(tmp_path / "snapshot")
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", _fake_snapshot_download)
+    resolved = _resolve_model_dir("openbmb/MiniCPM-o-4_5", revision="abc123")
+    assert resolved == str(tmp_path / "snapshot")
+    assert calls["model_ref"] == "openbmb/MiniCPM-o-4_5"
+    assert calls["revision"] == "abc123"
+    assert calls["allow_patterns"] == ["assets/*"]
+
+
+def test_snapshot_download_failure_propagates(monkeypatch):
+    def _raise(*args, **kwargs):
+        raise FileNotFoundError("offline and not cached")
+
+    monkeypatch.setattr(huggingface_hub, "snapshot_download", _raise)
+    with pytest.raises(FileNotFoundError):
+        _resolve_model_dir("openbmb/MiniCPM-o-4_5")
+
+
+def test_init_with_fake_path_does_not_resolve(monkeypatch):
+    """Mirrors the CPU-test construction: fake model path, no ``revision``."""
+    _no_hub(monkeypatch)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            model="/fake/model",
+            stage_connector_config=None,
+        )
+    )
+    model = MiniCPMO45Code2Wav(vllm_config=config)
+    assert model.model_path == "/fake/model"
+    assert model._default_prompt_wav == "/fake/model/assets/HT_ref_audio.wav"
+
+
+def test_default_prompt_wav_follows_resolved_model_path(monkeypatch):
+    _no_hub(monkeypatch)
+    config = SimpleNamespace(
+        model_config=SimpleNamespace(
+            model="openbmb/MiniCPM-o-4_5",
+            stage_connector_config=None,
+        )
+    )
+    model = MiniCPMO45Code2Wav(vllm_config=config)
+    model.model_path = "/resolved/snapshot"
+    assert model._default_prompt_wav == "/resolved/snapshot/assets/HT_ref_audio.wav"

tests/model_executor/models/qwen2_5_omni/test_qwen2_5_omni_embed.py

@@ -15,6 +15,7 @@ import functools
 import pytest
 import torch
 from pytest_mock import MockerFixture
+from vllm.multimodal.utils import set_mm_embedding_modality
 
 pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
@@ -117,13 +118,30 @@ class TestCheckInterleavedAudioVideo:
 # ---------------------------------------------------------------------------
 
 
-def make_mock_model(mocker: MockerFixture, hidden: int = 8):
+@functools.lru_cache(maxsize=1)
+def _omni_qwen2_5_thinker_cls():
+    """vllm_omni's thinker override, loaded lazily like the upstream one."""
+    from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import (
+        Qwen2_5OmniThinkerForConditionalGeneration,
+    )
+
+    return Qwen2_5OmniThinkerForConditionalGeneration
+
+
+def make_mock_model(mocker: MockerFixture, hidden: int = 8, thinker_cls=None):
     """
     Return a minimal mock of Qwen2_5OmniThinkerForConditionalGeneration
     that has enough structure to run embed_input_ids.
+
+    ``thinker_cls`` selects whose ``embed_input_ids`` the mock dispatches to.
+    It defaults to upstream's. Only the interleaved case passes omni's override:
+    omni's non-interleaved path falls through to a zero-arg ``super()`` call,
+    which raises TypeError when ``self`` is a Mock rather than an instance.
     """
     m = _qwen2_5_omni_thinker_mod()
     Qwen2_5OmniThinkerForConditionalGeneration = m.Qwen2_5OmniThinkerForConditionalGeneration
+    if thinker_cls is None:
+        thinker_cls = Qwen2_5OmniThinkerForConditionalGeneration
 
     model = mocker.Mock(spec=Qwen2_5OmniThinkerForConditionalGeneration)
 
@@ -155,9 +173,7 @@ def make_mock_model(mocker: MockerFixture, hidden: int = 8):
             is_multimodal=is_multimodal,
         )
 
-    model.embed_input_ids = lambda *a, **kw: m.Qwen2_5OmniThinkerForConditionalGeneration.embed_input_ids(  # noqa: E501
-        model, *a, **kw
-    )
+    model.embed_input_ids = lambda *a, **kw: thinker_cls.embed_input_ids(model, *a, **kw)
 
     model._super_embed_input_ids = fake_super_embed
 
@@ -171,11 +187,11 @@ def build_mm_embeds(audio_n, image_n, video_n, hidden, audio_val=10.0, image_val
     """
     embs = []
     if audio_n:
-        embs.append(torch.full((audio_n, hidden), audio_val))
+        embs.append(set_mm_embedding_modality(torch.full((audio_n, hidden), audio_val), "audio"))
     if image_n:
-        embs.append(torch.full((image_n, hidden), image_val))
+        embs.append(set_mm_embedding_modality(torch.full((image_n, hidden), image_val), "image"))
     if video_n:
-        embs.append(torch.full((video_n, hidden), video_val))
+        embs.append(set_mm_embedding_modality(torch.full((video_n, hidden), video_val), "video"))
     return embs
 
 
@@ -259,6 +275,11 @@ class TestEmbedInputIds:
         """
         Interleaved (use_audio_in_video): video chunks interleaved with audio.
         Video embeddings must go to video positions, audio to audio positions.
+
+        Runs against **omni's** thinker override, not upstream's: this is the
+        only unit-level coverage of omni's call into
+        ``merge_interleaved_embeddings``, whose upstream signature dropped
+        ``num_video``/``num_audio`` in vLLM #46213.
         """
         hidden = 8
         audio_val, video_val = 10.0, 30.0
@@ -270,11 +291,11 @@ class TestEmbedInputIds:
         audio_n = sum(audio_chunks)  # 6
 
         mm_embeds = [
-            torch.full((video_n, hidden), video_val),
-            torch.full((audio_n, hidden), audio_val),
+            set_mm_embedding_modality(torch.full((video_n, hidden), video_val), "video"),
+            set_mm_embedding_modality(torch.full((audio_n, hidden), audio_val), "audio"),
         ]
 
-        model, _ = make_mock_model(mocker, hidden)
+        model, _ = make_mock_model(mocker, hidden, thinker_cls=_omni_qwen2_5_thinker_cls())
         result = model.embed_input_ids(input_ids, mm_embeds, is_multimodal=is_multimodal)
 
         video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
@@ -308,8 +329,8 @@ class TestMergeInterleavedEmbeddings:
 
         inputs_embeds = torch.zeros(len(input_ids), hidden)
         mm_embeds = [
-            torch.full((num_video, hidden), 30.0),
-            torch.full((num_audio, hidden), 10.0),
+            set_mm_embedding_modality(torch.full((num_video, hidden), 30.0), "video"),
+            set_mm_embedding_modality(torch.full((num_audio, hidden), 10.0), "audio"),
         ]
 
         result = m.merge_interleaved_embeddings(
@@ -318,8 +339,6 @@ class TestMergeInterleavedEmbeddings:
             is_video,
             is_audio,
             is_multimodal,
-            num_video,
-            num_audio,
         )
 
         video_pos = is_video.nonzero(as_tuple=True)[0]

tests/model_executor/models/qwen3_omni/test_qwen3_omni_moe_lora.py

@@ -84,6 +84,7 @@ class _CpuFusedMoE3DWithLoRA(FusedMoE3DWithLoRA):
         self.tp_size = 1
         self.tp_rank = 0
         self._w13_slices = 1
+        self.enable_moe_shared_loras = False
 
 
 def _save_peft_adapter(tmp_path: Path) -> Path:

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

tests/worker/test_gpu_ar_model_runner.py

@@ -69,6 +69,9 @@ def test_speech_extra_params_reach_model_sampler_as_sampling_metadata(monkeypatc
         vocab_size=1024,
         block_sizes=[1],
         kernel_block_sizes=[1],
+        # vLLM 0.26: one entry per KV cache group; max_model_len=8 with
+        # block_sizes=[1] means one group of 8 blocks.
+        max_num_blocks_per_req=[8],
     )
     input_batch.add_request(
         CachedRequestState(

tests/worker/test_gpu_generation_model_runner.py

@@ -4,7 +4,7 @@ from types import SimpleNamespace
 import pytest
 import torch
 from vllm.config import CacheConfig
-from vllm.v1.worker.gpu_model_runner import EMPTY_MODEL_RUNNER_OUTPUT
+from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 
 import vllm_omni.worker.gpu_generation_model_runner as gen_runner_module
 from vllm_omni.worker.gpu_generation_model_runner import (

tests/worker/test_omni_connector_mixin.py

@@ -1245,7 +1245,7 @@ class TestRankAwareKVRouting(unittest.TestCase):
 
 class TestAttachOmniConnectorOutput(unittest.TestCase):
     def test_wraps_empty_model_runner_output_when_signals_exist(self):
-        from vllm.v1.worker.gpu_model_runner import EMPTY_MODEL_RUNNER_OUTPUT
+        from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
 
         host = MixinHost()
         host.get_omni_connector_output = lambda: OmniConnectorOutput(chunk_ready_req_ids={"req-1"})

vllm_omni/__init__.py

@@ -27,8 +27,14 @@ except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
     patch = None  # type: ignore
 
 # Register custom configs (AutoConfig, AutoTokenizer) as early as possible.
-from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
-from vllm_omni.transformers_utils import parsers as _parsers  # noqa: F401, E402
+try:
+    from vllm_omni.transformers_utils import configs as _configs  # noqa: F401, E402
+    from vllm_omni.transformers_utils import parsers as _parsers  # noqa: F401, E402
+except ModuleNotFoundError as exc:  # pragma: no cover - optional dependency
+    if exc.name != "vllm":
+        raise
+    # Allow importing vllm_omni without vllm (e.g., documentation builds)
+    pass
 
 from .config import OmniModelConfig
 

vllm_omni/attention/fish_kvcache_backend.py

@@ -234,11 +234,11 @@ def _forward_with_fish_kvcache(
     if attn_metadata is not None and not attn_metadata.use_cascade:
         num_actual_tokens = attn_metadata.num_actual_tokens
         # vLLM >=0.23.0 lays out the KV cache as
-        # (num_blocks, 2, block_size, num_kv_heads, head_size), so key/value
-        # live on dim 1 (matches FlashAttentionImpl's kv_cache.unbind(1)).
-        # The pre-0.23 layout had them on dim 0; unbind(0) here would split
-        # along num_blocks and raise "too many values to unpack".
-        key_cache, value_cache = kv_cache.unbind(1)
+        # (num_blocks, num_kv_heads, block_size, 2*head_size) with K and V
+        # packed in the last dimension.  Match FlashAttentionImpl's own
+        # unpacking: transpose(1, 2) then split on the content dim.
+        head_size = kv_cache.shape[-1] // 2
+        key_cache, value_cache = kv_cache.transpose(1, 2).split(head_size, dim=-1)
         q = query[:num_actual_tokens]
         out = output[:num_actual_tokens]
         active_batch = int(q.shape[0])

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

@@ -146,7 +146,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,