Update docs/community/governance.md

Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Signed-off-by: Roger Wang <hey@rogerw.io>

.buildkite/bootstrap-amd-omni.sh

@@ -90,18 +90,10 @@ upload_pipeline() {
     FAIL_FAST=$(fail_fast)
 
     cd .buildkite
-
-    # Select test definition file: merge suite for main, ready suite for PRs
-    if [[ $BUILDKITE_BRANCH == "main" ]]; then
-        TEST_YAML="test-amd-merge.yml"
-    else
-        TEST_YAML="test-amd-ready.yaml"
-    fi
-
     (
         set -x
         # Output pipeline.yaml with all blank lines removed
-        minijinja-cli test-template.j2 "$TEST_YAML" \
+        minijinja-cli test-template.j2 test-amd.yaml \
             -D branch="$BUILDKITE_BRANCH" \
             -D list_file_diff="$LIST_FILE_DIFF" \
             -D run_all="$RUN_ALL" \

.buildkite/bootstrap-intel-omni.sh

@@ -1,51 +0,0 @@
-#!/bin/bash
-# vllm-omni Intel bootstrap
-# Uses static pipeline-intel.yaml for Intel XPU tests
-# Based on: bootstrap-amd-omni.sh
-
-set -euo pipefail
-
-if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
-    DOCS_ONLY_DISABLE=0
-fi
-
-upload_pipeline() {
-    echo "--- 🛠 Preparing Intel pipeline"
-    buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
-}
-
-get_diff() {
-    $(git add .)
-    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
-}
-
-get_diff_main() {
-    $(git add .)
-    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
-}
-
-file_diff=$(get_diff)
-if [[ $BUILDKITE_BRANCH == "main" ]]; then
-    file_diff=$(get_diff_main)
-fi
-
-# ----------------------------------------------------------------------
-# Early exit: skip pipeline if all changed files are under docs/
-# ----------------------------------------------------------------------
-if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
-    docs_only=1
-    while IFS= read -r f; do
-        [[ -z "$f" ]] && continue
-        if [[ "$f" != docs/* ]]; then
-            docs_only=0
-            break
-        fi
-    done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
-
-    if [[ "$docs_only" -eq 1 ]]; then
-        buildkite-agent annotate ":memo: CI skipped — docs only" --style "info"
-        exit 0
-    fi
-fi
-
-upload_pipeline

.buildkite/pipeline-intel.yaml

@@ -1,16 +0,0 @@
-# .buildkite/pipeline-intel.yaml
-steps:
-  - group: "INTEL Tests"
-    steps:
-      - label: "Intel: XPU Test"
-        command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
-        agents:
-          queue: "intel-gpu-omni"
-        env:
-          DOCKER_BUILDKIT: "1"
-          # Buildkite will automatically replace this with the actual commit hash
-          VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.18.0"
-        priority: 100
-        timeout_in_minutes: 60
-        soft_fail: true

.buildkite/scripts/hardware_ci/run-amd-test.sh

@@ -87,11 +87,7 @@ HF_CACHE="$(realpath ~)/huggingface"
 mkdir -p "${HF_CACHE}"
 HF_MOUNT="/root/.cache/huggingface"
 
-if [[ -n "${TEST_COMMAND:-}" ]]; then
-    commands="$TEST_COMMAND"
-else
-    commands="$@"
-fi
+commands=$@
 echo "Commands:$commands"
 
 PARALLEL_JOB_COUNT=8
@@ -106,7 +102,6 @@ if [[ -z "$render_gid" ]]; then
 fi
 
 # check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
-# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
 if [[ $commands == *"--shard-id="* ]]; then
   # assign job count as the number of shards used
   commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
@@ -123,7 +118,7 @@ if [[ $commands == *"--shard-id="* ]]; then
         --rm \
         -e MIOPEN_DEBUG_CONV_DIRECT=0 \
         -e MIOPEN_DEBUG_CONV_GEMM=0 \
-        -e VLLM_ROCM_USE_AITER=0 \
+        -e VLLM_ROCM_USE_AITER=1 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
         -e AWS_ACCESS_KEY_ID \
@@ -158,7 +153,7 @@ else
           --rm \
           -e MIOPEN_DEBUG_CONV_DIRECT=0 \
           -e MIOPEN_DEBUG_CONV_GEMM=0 \
-          -e VLLM_ROCM_USE_AITER=0 \
+          -e VLLM_ROCM_USE_AITER=1 \
           -e HF_TOKEN \
           -e AWS_ACCESS_KEY_ID \
           -e AWS_SECRET_ACCESS_KEY \

.buildkite/test-amd-merge.yml

@@ -1,197 +0,0 @@
-steps:
-
-- label: "Simple Unit Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
-
-- label: "Diffusion Model Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
-
-- label: "Diffusion Images API LoRA E2E"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
-
-- label: "Diffusion Model CPU offloading Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-        EXIT1=\$?
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
-
-## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
-# - label: "Audio Generation Model Test"
-#   agent_pool: mi325_1
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export GPU_ARCHS=gfx942
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
-
-- label: "Diffusion Cache Backend Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
-
-- label: "Diffusion Sequence Parallelism Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
-
-# merge-only tests
-- label: "Diffusion Tensor Parallelism Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
-
-- label: "Diffusion GPU Worker Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-
-- label: "Benchmark & Engine Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/benchmarks/test_serve_cli.py
-        EXIT1=\$?
-        pytest -s -v tests/engine/test_async_omni_engine_abort.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
-
-- label: "Omni Model Test Qwen2-5-Omni"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-
-- label: "Omni Model Test Qwen3-Omni"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
-
-- label: "Qwen3-TTS E2E Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
-
-- label: "Diffusion Image Edit Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
-
-# split Bagel Model Test with H100 (Real Weights) into three tests
-- label: "Bagel Text2Img Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
-
-- label: "Bagel Img2Img Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
-
-- label: "Bagel Online Serving Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - export VLLM_IMAGE_FETCH_TIMEOUT=60
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"

.buildkite/test-amd-ready.yaml

@@ -1,178 +0,0 @@
-steps:
-
-- label: "Simple Unit Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_ROCM_USE_AITER=0
-    - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
-
-- label: "Diffusion Model Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
-
-- label: "Diffusion Model CPU offloading Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 20m bash -c '
-        set +e
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-        EXIT1=\$?
-        pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
-
-## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
-# - label: "Audio Generation Model Test"
-#   agent_pool: mi325_1
-#   depends_on: amd-build
-#   mirror_hardwares: [amdproduction]
-#   grade: Blocking
-#   commands:
-#     - export GPU_ARCHS=gfx942
-#     - export VLLM_LOGGING_LEVEL=DEBUG
-#     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-#     - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
-
-- label: "Diffusion Cache Backend Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
-
-- label: "Diffusion Sequence Parallelism Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
-
-- label: "Diffusion GPU Worker Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
-
-- label: "Benchmark & Engine Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - |
-      timeout 30m bash -c '
-        set +e
-        pytest -s -v tests/benchmarks/test_serve_cli.py
-        EXIT1=\$?
-        pytest -s -v tests/engine/test_async_omni_engine_abort.py
-        EXIT2=\$?
-        exit \$((EXIT1 | EXIT2))
-      '
-
-- label: "Omni Model Test Qwen2-5-Omni"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-
-- label: "Omni Model Test Qwen3-Omni"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
-    - timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
-
-- label: "Qwen3-TTS E2E Test"
-  agent_pool: mi325_2
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
-
-- label: "Diffusion Image Edit Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
-
-- label: "Bagel Text2Img Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
-
-- label: "Bagel Img2Img Model Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
-
-- label: "Bagel Online Serving Test"
-  agent_pool: mi325_1
-  depends_on: amd-build
-  mirror_hardwares: [amdproduction]
-  grade: Blocking
-  commands:
-    - export GPU_ARCHS=gfx942
-    - export VLLM_TEST_CLEAN_GPU_MEMORY=1
-    - export VLLM_IMAGE_FETCH_TIMEOUT=60
-    - export VLLM_LOGGING_LEVEL=DEBUG
-    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"

.buildkite/test-amd.yaml

@@ -1,7 +1,7 @@
 steps:
 
 - label: "Diffusion Model Test"
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -11,7 +11,7 @@ steps:
     - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
 - label: "Diffusion Images API LoRA E2E"
-  timeout_in_minutes: 30
+  timeout_in_minutes: 20
   agent_pool: mi325_1
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
@@ -101,7 +101,6 @@ steps:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
-    - export VLLM_ROCM_USE_AITER=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
 

.buildkite/test-merge.yml

@@ -12,12 +12,11 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
-    timeout_in_minutes: 30
+    timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
@@ -30,12 +29,11 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Images API LoRA E2E"
-    timeout_in_minutes: 30
+    timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
       - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
@@ -48,7 +46,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -67,7 +64,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -85,7 +81,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -103,7 +98,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -122,7 +116,6 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -141,7 +134,6 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -160,17 +152,21 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Engine Test"
+  - label: "Benchmark & Engine Test"
     depends_on: upload-merge-pipeline
     commands:
       - |
         timeout 15m bash -c '
                 export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
         '
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -182,7 +178,6 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -202,7 +197,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -232,11 +226,6 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -272,11 +261,6 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -288,64 +272,12 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: "Bagel Model Test with H100 (Real Weights)"
-    timeout_in_minutes: 60
-    depends_on: upload-merge-pipeline
-    commands:
-      - |
-        timeout 55m bash -c '
-          set -e
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          export VLLM_IMAGE_FETCH_TIMEOUT=60
-          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
-          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
-          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  # - label: "Voxtral TTS E2E Test"
-  #   timeout_in_minutes: 20
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   timeout_in_minutes: 30
   #   depends_on: upload-merge-pipeline
   #   commands:
-  #     - |
-  #       timeout 20m bash -c '
-  #         export VLLM_LOGGING_LEVEL=DEBUG
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-  #       '
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
   #   agents:
   #     queue: "mithril-h100-pool"
   #   plugins:

.buildkite/test-nightly.yml

@@ -11,9 +11,6 @@ steps:
         EXIT1=$$?
         pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
         EXIT2=$$?
-        # TODO: uncomment when mistralai/tts-model is public
-        # pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-        # EXIT3=$$?
         exit $$((EXIT1 | EXIT2))
     agents:
       queue: "mithril-h100-pool"
@@ -33,11 +30,6 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -66,14 +58,13 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: ":full_moon: Diffusion Model Test with H100"
     timeout_in_minutes: 60
     depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1"
+    # if: build.env("NIGHTLY") == "1"
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
       - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model"
@@ -95,11 +86,6 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -130,7 +116,6 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -164,13 +149,6 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
-                    value: "120"
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:
@@ -185,7 +163,7 @@ steps:
   - label: ":email: Nightly Collection & Email"
     key: nightly-perf-distribution
     depends_on:
-      - nightly-performance
+    - nightly-performance
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl

.buildkite/test-ready.yml

@@ -5,24 +5,6 @@ steps:
       - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
     agents:
       queue: "gpu_1_queue"
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Voxtral TTS CUDA Unit Test"
-    timeout_in_minutes: 10
-    depends_on: upload-ready-pipeline
-    commands:
-      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
-    agents:
-      queue: "gpu_1_queue"
     plugins:
       - docker#v5.2.0:
           image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
@@ -36,7 +18,7 @@ steps:
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -46,7 +28,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -63,7 +44,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -80,7 +60,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -97,7 +76,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -115,7 +93,6 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -133,18 +110,22 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
 
-  - label: "Engine Test"
+  - label: "Benchmark & Engine Test"
     depends_on: upload-ready-pipeline
     commands:
       - |
         timeout 15m bash -c '
                 export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                set +e
+                pytest -s -v tests/benchmarks/test_serve_cli.py
+                EXIT1=$$?
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
+                EXIT2=$$?
+                exit $$((EXIT1 | EXIT2))
         '
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -156,7 +137,6 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -179,7 +159,6 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -242,49 +221,9 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
-            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  # TODO: uncomment when mistralai/tts-model is public
-  # - label: "Voxtral-TTS E2E Test"
-  #   depends_on: upload-ready-pipeline
-  #   commands:
-  #     - |
-  #       timeout 20m bash -c '
-  #         export VLLM_LOGGING_LEVEL=DEBUG
-  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
-  #         pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
-  #       '
-  #   agents:
-  #     queue: "mithril-h100-pool"
-  #   plugins:
-  #     - kubernetes:
-  #         podSpec:
-  #           containers:
-  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-  #               resources:
-  #                 limits:
-  #                   nvidia.com/gpu: 1
-  #               volumeMounts:
-  #                 - name: devshm
-  #                   mountPath: /dev/shm
-  #                 - name: hf-cache
-  #                   mountPath: /root/.cache/huggingface
-  #               env:
-  #                 - name: HF_HOME
-  #                   value: /root/.cache/huggingface
-  #           nodeSelector:
-  #             node.kubernetes.io/instance-type: gpu-h100-sxm
-  #           volumes:
-  #             - name: devshm
-  #               emptyDir:
-  #                 medium: Memory
-  #             - name: hf-cache
-  #               hostPath:
-  #                 path: /mnt/hf-cache
-  #                 type: DirectoryOrCreate
-
   # - label: "Diffusion Image Edit Test with H100 (1 GPU)"
   #   depends_on: upload-ready-pipeline
   #   commands:
@@ -322,132 +261,39 @@ steps:
   #                 path: /mnt/hf-cache
   #                 type: DirectoryOrCreate
 
-  - label: "Bagel Text2Img Model Test with H100"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 30m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: "Bagel Img2Img Model Test with H100"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 30m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
-
-  - label: "Bagel Online Serving Test with H100"
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 40m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          export VLLM_TEST_CLEAN_GPU_MEMORY=1
-          export VLLM_IMAGE_FETCH_TIMEOUT=60
-          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
-        '
-    agents:
-      queue: "mithril-h100-pool"
-    plugins:
-      - kubernetes:
-          podSpec:
-            containers:
-              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-                resources:
-                  limits:
-                    nvidia.com/gpu: 1
-                volumeMounts:
-                  - name: devshm
-                    mountPath: /dev/shm
-                  - name: hf-cache
-                    mountPath: /root/.cache/huggingface
-                env:
-                  - name: HF_HOME
-                    value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
-            nodeSelector:
-              node.kubernetes.io/instance-type: gpu-h100-sxm
-            volumes:
-              - name: devshm
-                emptyDir:
-                  medium: Memory
-              - name: hf-cache
-                hostPath:
-                  path: /mnt/hf-cache
-                  type: DirectoryOrCreate
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 30m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
+  #       '
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 1
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate

.buildkite/test-template-amd-omni.j2

@@ -40,14 +40,9 @@
           {% else %}
           queue: amd_mi325_1
           {% endif %}
-{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
-{% set indented_cmd = cmd_body | replace("\n", "\n            ") %}
-        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
+        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe  }} ; {{ step.command  or (step.commands | join(" && ")) | safe }}"
         env:
           DOCKER_BUILDKIT: "1"
-          TEST_COMMAND: |-
-            (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
-            {{ indented_cmd | safe }}
         priority: 100
         {% if step.grade and step.grade == "Blocking" %}
         soft_fail: false

.github/ISSUE_TEMPLATE/400-bug-report.yml

@@ -74,21 +74,28 @@ body:
       If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
 
       ```python
-      from vllm_omni.entrypoints.omni import Omni
-      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-      from vllm import SamplingParams
+      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
 
-      omni = Omni(
-          model="Qwen/Qwen-Image",
-          stage_configs_path="/path/to/stage_configs.yaml",
+      # Create stage configurations
+      ar_config = create_ar_stage_config(
+          stage_id=0,
+          model_path="Qwen/Qwen3-0.6B",
+          input_modalities=["text"],
+          output_modalities=["text"]
       )
 
-      prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
-      sampling_params_list = [
-          SamplingParams(max_tokens=1),
-          OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
-      ]
-      outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
+      dit_config = create_dit_stage_config(
+          stage_id=1,
+          model_path="stabilityai/stable-diffusion-2-1",
+          input_modalities=["text"],
+          output_modalities=["image"]
+      )
+
+      # Initialize OmniLLM
+      omni = OmniLLM([ar_config, dit_config])
+
+      # Generate
+      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
       ```
 
       If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.

README.md

@@ -17,7 +17,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 
 *Latest News* 🔥
 
-- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
+- [2026/03] Joined vLLM HK First Meetup! Video on [youtube](https://youtu.be/sgwNfsNnR9I)
 - [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
 - [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
 - [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.

benchmarks/diffusion/backends.py

@@ -1,4 +1,3 @@
-import asyncio
 import base64
 import mimetypes
 import os
@@ -36,7 +35,6 @@ class RequestFuncOutput:
     error: str = ""
     start_time: float = 0.0
     response_body: dict[str, Any] = field(default_factory=dict)
-    stage_durations: dict[str, float] = field(default_factory=dict)
     peak_memory_mb: float = 0.0
     slo_achieved: bool | None = None
 
@@ -57,7 +55,6 @@ async def async_request_chat_completions(
     input: RequestFuncInput,
     session: aiohttp.ClientSession,
     pbar: tqdm | None = None,
-    enable_diffusion_pipeline_profiler: bool = False,
 ) -> RequestFuncOutput:
     output = RequestFuncOutput()
     output.start_time = time.perf_counter()
@@ -111,18 +108,6 @@ async def async_request_chat_completions(
                 output.success = True
                 if "peak_memory_mb" in resp_json:
                     output.peak_memory_mb = resp_json["peak_memory_mb"]
-                try:
-                    choices = resp_json.get("choices", [])
-                    if choices and isinstance(choices, list):
-                        msg = choices[0].get("message", {})
-                        if isinstance(msg, dict):
-                            content = msg.get("content", [])
-                            if content and isinstance(content, list) and len(content) > 0:
-                                first_item = content[0]
-                                if isinstance(first_item, dict):
-                                    output.stage_durations = first_item.get("stage_durations")
-                except (IndexError, TypeError, AttributeError):
-                    pass
             else:
                 output.error = f"HTTP {response.status}: {await response.text()}"
                 output.success = False
@@ -207,140 +192,7 @@ async def async_request_openai_images(
     return output
 
 
-async def async_request_v1_videos(
-    input: RequestFuncInput,
-    session: aiohttp.ClientSession,
-    pbar: tqdm | None = None,
-) -> RequestFuncOutput:
-    output = RequestFuncOutput()
-    output.start_time = time.perf_counter()
-
-    files = dict(input.extra_body)
-    if input.prompt:
-        files.setdefault("prompt", input.prompt)
-    if input.width and input.height:
-        files.setdefault("height", input.height)
-        files.setdefault("width", input.width)
-    if input.num_frames:
-        files.setdefault("num_frames", input.num_frames)
-    if input.num_inference_steps:
-        files.setdefault("num_inference_steps", input.num_inference_steps)
-    if input.seed is not None:
-        files.setdefault("seed", input.seed)
-    if input.fps:
-        files.setdefault("fps", input.fps)
-
-    form = aiohttp.FormData()
-    for k, v in files.items():
-        form.add_field(k, str(v))
-
-    image_file = None
-    if input.image_paths and len(input.image_paths) > 0:
-        image_path = input.image_paths[0]
-        image_file = open(image_path, "rb")
-        form.add_field(
-            "input_reference",
-            image_file,
-            filename=os.path.basename(image_path),
-            content_type="application/octet-stream",
-        )
-
-    job_id = None
-    job_status = None
-    poll_json = {}
-    resp_json = {}
-
-    try:
-        # invoke a post request (POST /v1/videos)
-        async with session.post(input.api_url, data=form) as response:
-            if response.status == 200:
-                resp_json = await response.json()
-                job_id = resp_json.get("id")
-                job_status = resp_json.get("status")
-                if not job_id or not job_status:
-                    output.error = "API response missing job 'id' or 'status' field."
-                    output.success = False
-                    return output
-            else:
-                output.error = f"HTTP {response.status}: {await response.text()}"
-                output.success = False
-                return output
-
-        # invoke a poll request (GET /v1/videos/{video_id})
-        poll_interval = 2.0  # Unit(s)
-        timeout_seconds = 600.0
-        deadline = time.perf_counter() + timeout_seconds
-        job_url = f"{input.api_url}/{job_id}"
-
-        while job_status not in {"completed", "failed"}:
-            await asyncio.sleep(poll_interval)
-
-            async with session.get(job_url) as poll_response:
-                if poll_response.status != 200:
-                    output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
-                    output.success = False
-                    return output
-
-                poll_json = await poll_response.json()
-                job_status = poll_json.get("status")
-
-                if time.perf_counter() >= deadline:
-                    output.error = f"Timed out waiting for video job {job_id} to complete."
-                    output.success = False
-                    return output
-
-        if job_status == "failed":
-            output.error = f"Video job failed: {poll_json}"
-            output.success = False
-            return output
-
-        # invoke a get request (GET /v1/videos/{video_id}/content)
-        content_url = f"{job_url}/content"
-        async with session.get(content_url) as content_response:
-            if content_response.status != 200:
-                output.error = (
-                    f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
-                )
-                output.success = False
-                return output
-
-            video_bytes = await content_response.read()
-            output.response_body = video_bytes
-            output.success = True
-            if "peak_memory_mb" in poll_json:
-                output.peak_memory_mb = poll_json["peak_memory_mb"]
-            elif "peak_memory_mb" in resp_json:
-                output.peak_memory_mb = resp_json["peak_memory_mb"]
-    except Exception as e:
-        output.error = str(e)
-        output.success = False
-    finally:
-        if image_file is not None:
-            image_file.close()
-
-        if job_id is not None:
-            try:
-                async with session.delete(f"{input.api_url}/{job_id}") as _:
-                    pass
-            except Exception as e:
-                print(f"Failed to clean up video job {job_id}: {e}")
-
-    output.latency = time.perf_counter() - output.start_time
-
-    if output.success and input.slo_ms is not None:
-        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
-
-    if pbar:
-        pbar.update(1)
-    return output
-
-
 backends_function_mapping = {
-    "2i": {
-        "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
-        "openai": (async_request_openai_images, "/v1/images/generations"),
-    },
-    "2v": {
-        "v1/videos": (async_request_v1_videos, "/v1/videos"),
-    },
+    "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+    "openai": (async_request_openai_images, "/v1/images/generations"),
 }

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -9,7 +9,6 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
 Supports multiple backends:
     - vllm-omni: Uses /v1/chat/completions endpoint (default)
     - openai: Uses /v1/images/generations endpoint
-    - v1/videos: Use /v1/videos endpoint
 
 Usage:
     # Video (vllm-omni backend)
@@ -29,16 +28,6 @@ Usage:
         --backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024
 
-    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend vllm-omni --dataset random --task t2i --num-prompts 1 \
-        --max-concurrency 1 --enable-negative-prompt \
-        --random-request-config '[
-            {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
-            {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
-            {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
-            {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
-        ]'
-
     i2i:
     python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
         --backend vllm-omni --dataset vbench --task i2i --num-prompts 10
@@ -49,16 +38,6 @@ Usage:
         --backend openai --dataset vbench --task t2i --num-prompts 10 \
         --height 1024 --width 1024 --port 3000
 
-    # Video (v1/vedeos)
-    t2v:
-    python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
-        --backend v1/videos --dataset random --task t2v --num-prompts 1 \
-        --max-concurrency 1 --enable-negative-prompt \
-        --random-request-config '[
-            {"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
-        ]'
-
-
 """
 
 import argparse
@@ -66,10 +45,7 @@ import ast
 import asyncio
 import glob
 import json
-import logging
 import os
-import random
-import tempfile
 import time
 import uuid
 from abc import ABC, abstractmethod
@@ -81,11 +57,8 @@ import aiohttp
 import numpy as np
 import requests
 from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
-from PIL import Image
 from tqdm.asyncio import tqdm
 
-logger = logging.getLogger(__name__)
-
 
 class BaseDataset(ABC):
     def __init__(self, args, api_url: str, model: str):
@@ -543,66 +516,26 @@ class TraceDataset(BaseDataset):
 
 
 class RandomDataset(BaseDataset):
-    def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
-        super().__init__(args, api_url, model)
+    def __init__(self, args, api_url: str, model: str):
+        self.args = args
+        self.api_url = api_url
+        self.model = model
         self.num_prompts = args.num_prompts
-        self.enable_negative_prompt = enable_negative_prompt
-        self.random_request_config = getattr(args, "random_request_config", None)
-        if self.random_request_config:
-            self.random_request_config = json.loads(self.random_request_config)
-            self._weights = [p["weight"] for p in self.random_request_config]
-
-            self.random_request_config = [
-                {k: v for k, v in p.items() if k != "weight"} for p in self.random_request_config
-            ]
-
-            seed = getattr(args, "random_request_seed", 42)
-            self._rng = random.Random(seed)
-
-            self._sampled_requests = self._rng.choices(
-                self.random_request_config,
-                weights=self._weights,
-                k=self.num_prompts,
-            )
-        else:
-            self._sampled_requests = None
-
-        # Random image generate
-        if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
-            img = Image.new("RGB", (512, 512), (255, 255, 255))
-
-            image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
-            self._random_image_path = [image_path]
-            img.save(image_path)
-        else:
-            self._random_image_path = None
 
     def __len__(self) -> int:
         return self.num_prompts
 
     def __getitem__(self, idx: int) -> RequestFuncInput:
-        extra_body = {}
-        if self.enable_negative_prompt:
-            extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
-
-        params = {
-            "width": self.args.width,
-            "height": self.args.height,
-            "num_frames": self.args.num_frames,
-            "num_inference_steps": self.args.num_inference_steps,
-            "fps": self.args.fps,
-        }
-        if self._sampled_requests:
-            profile = self._sampled_requests[idx]
-            params.update(profile)
         return RequestFuncInput(
             prompt=f"Random prompt {idx} for benchmarking diffusion models",
             api_url=self.api_url,
             model=self.model,
+            width=self.args.width,
+            height=self.args.height,
+            num_frames=self.args.num_frames,
+            num_inference_steps=self.args.num_inference_steps,
             seed=self.args.seed,
-            extra_body=extra_body,
-            image_paths=self._random_image_path,
-            **params,
+            fps=self.args.fps,
         )
 
     def get_requests(self) -> list[RequestFuncInput]:
@@ -753,7 +686,6 @@ def calculate_metrics(
         "latency_mean": np.mean(latencies) if latencies else 0,
         "latency_median": np.median(latencies) if latencies else 0,
         "latency_p99": np.percentile(latencies, 99) if latencies else 0,
-        "latency_p95": np.percentile(latencies, 95) if latencies else 0,
         "latency_p50": np.percentile(latencies, 50) if latencies else 0,
         "peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
         "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
@@ -810,32 +742,8 @@ async def benchmark(args):
     if args.base_url is None:
         args.base_url = f"http://{args.host}:{args.port}"
 
-    VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
-    IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
-
-    if args.task in VIDEO_TASKS:
-        task_type = "2v"
-    elif args.task in IMAGE_TASKS:
-        task_type = "2i"
-    else:
-        raise ValueError(
-            f"Unsupported task: '{args.task}'. "
-            f"Valid video tasks: {sorted(VIDEO_TASKS)}, "
-            f"Valid image tasks: {sorted(IMAGE_TASKS)}"
-        )
-
-    valid_backends = sorted(backends_function_mapping[task_type].keys())
-
-    if args.backend not in valid_backends:
-        logger.error(
-            f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
-            f"Valid backends for this task type: {valid_backends}\n"
-            f"Example usage: --task {args.task} --backend {valid_backends[0]}"
-        )
-        raise ValueError("Backend validation failed. See log above for valid options.")
-
     # Setup API URL and request function based on backend
-    request_func, api_url = backends_function_mapping[task_type][args.backend]
+    request_func, api_url = backends_function_mapping[args.backend]
     api_url = f"{args.base_url}{api_url}"
 
     if args.dataset == "vbench":
@@ -843,7 +751,7 @@ async def benchmark(args):
     elif args.dataset == "trace":
         dataset = TraceDataset(args, api_url, args.model)
     elif args.dataset == "random":
-        dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
+        dataset = RandomDataset(args, api_url, args.model)
     else:
         raise ValueError(f"Unknown dataset: {args.dataset}")
 
@@ -939,7 +847,6 @@ async def benchmark(args):
     print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"]))
     print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
     print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
-    print("{:<40} {:<15.4f}".format("Latency P95 (s):", metrics["latency_p95"]))
 
     if args.slo:
         print(f"{'-' * 50}")
@@ -976,7 +883,7 @@ if __name__ == "__main__":
         "--backend",
         type=str,
         default="vllm-omni",
-        choices=["vllm-omni", "openai", "v1/videos"],
+        choices=["vllm-omni", "openai"],
         help="Backend to target the benchmark to.",
     )
     parser.add_argument(
@@ -1065,25 +972,6 @@ if __name__ == "__main__":
         help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
     )
     parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
-    parser.add_argument(
-        "--enable-negative-prompt",
-        action="store_true",
-        default=False,
-        help="Generate negative prompts when using the random dataset.",
-    )
-    parser.add_argument(
-        "--random-request-config",
-        type=str,
-        default=None,
-        help=(
-            "JSON string defining random request profiles. "
-            "Each profile may contain: width, height, num_inference_steps, etc. "
-            "The 'weight' field controls sampling probability (relative weight). "
-            "Example: "
-            '[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
-            '{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
-        ),
-    )
 
     args = parser.parse_args()
 

benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md

@@ -1,169 +0,0 @@
-# Qwen-Image Serving Performance Dashboard
-
-This document describes how to deploy and benchmark **Qwen-Image** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
-
----
-
-# 1. Overview
-
-Qwen-Image is a multimodal text-to-image generation model served through the vLLM-Omni infrastructure.
-
-This document covers:
-
-* Service launch configuration (including acceleration options)
-* Benchmark scripts and usage
-* Dataset and workload settings
-* Performance measurement results
-* Reproducibility guidelines
-
----
-
-# 2. Test Environment
-| Component | Specification |
-|------------|----------------|
-| GPU | NVIDIA A100-SXM4-80GB |
-| Diffusion Attention Backend | FlashAttention |
-
-# 3. Service Launch Configuration
-
-## 3.1 Basic Serving Command
-
-```bash
-vllm serve Qwen/Qwen-Image --omni \
-    --port 8091
-```
-
-## 3.2 Key Parameters
-
-| Parameter             | Description              |
-| --------------------- | ------------------------ |
-| `--cfg-parallel-size` | CFG parallelism degree   |
-| `--ulysses-degree`    | Ulysses parallel degree  |
-| `--vae-patch-parallel-size`    | VAE parallel degree  |
-| `--tensor-parallel-size` | Tensor parallelism degree |
-
-Record these parameters when reporting performance results.
-
----
-
-# 4. Benchmark Script
-
-## 4.1 Benchmark Entry
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend vllm-omni \
-    --dataset <DATASET_NAME> \
-    --task t2i \
-    --num-prompts <N> \
-    --max-concurrency <C> \
-    --enable-negative-prompt \
-    --random-request-config <CFG>
-```
-
-## 4.2 Key Benchmark Arguments
-
-| Parameter              | Description                       |
-| ---------------------- | --------------------------------- |
-| `--backend`            | Serving backend (use `vllm-omni`) |
-| `--dataset`            | Dataset name (`random` or custom) |
-| `--task`               | Task type (e.g., `t2i`)           |
-| `--num-prompts`        | Total number of requests          |
-| `--max-concurrency`    | Client-side concurrency           |
-| `--random-request-config`| JSON string defining random request |
-
----
-
-# 5. Dataset & Workload Settings
-
-## 5.1 Recommended Evaluation Configurations
-
-### Dataset A ( 512 Resolution)
-
-* Dataset: `random`
-* Task: t2i
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":512,"height":512,"num_inference_steps":20,"weight":1}
-]
-```
-
-### Dataset B (1536 Resolution)
-
-* Dataset: `random`
-* Task: t2i
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":1536,"height":1536,"num_inference_steps":35,"weight":1}
-]
-```
-
-### Dataset C (Mix Resolution)
-
-* Dataset: `random`
-* Task: t2i
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
- {"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
- {"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
- {"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
-]
-```
----
-
-## 5.2 Example Benchmark Command
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend vllm-omni \
-    --dataset random \
-    --task t2i \
-    --num-prompts 1 \
-    --max-concurrency 1 \
-    --enable-negative-prompt \
-    --random-request-config '[
-        {"width":512,"height":512,"num_inference_steps":20,"weight":1}
-    ]'
-```
-
----
-
-# 6. Performance Metrics
-
-The following metrics are collected during benchmarking:
-
-| Metric             | Description                   | Unit    |
-| ------------------ | ----------------------------- | ------- |
-| Mean Latency        | Mean of latency       | seconds |
-| P99 Latency        | P99 of latency             | seconds |
-
----
-
-# 7. Performance Results
-
-| Dataset Configuration | Max Concur. | CFG | Usp | Tp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
-|-----------------------|-----|-----|-----|----|--------------|------------------|------------------|
-| Dataset A | 1 | 2 | 2 | Off | Off          | 2.2087          | 2.2087          |
-| Dataset B | 1 | 2 | 2 | Off | Off          | 19.6739          | 19.6739          |
-| Dataset C | 1 | 2 | 2 | Off | Off          | 5.67259        | 18.6234 |
----
-
-# 8. Reproducibility Checklist
-
-To ensure consistent and comparable benchmark results:
-
-* Record GPU type
-* Record parallel configuration
-* Record benchmark parameters (resolution, concurrency, number of prompts)
-* Ensure no background workload on GPUs during testing
-
----
-
-This document serves as the official Qwen-Image serving performance reference under vLLM-Omni.

benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md

@@ -1,170 +0,0 @@
-# Wan2.2 Serving Performance Dashboard
-
-This document describes how to deploy and benchmark **Wan-AI/Wan2.2-T2V-A14B-Diffusers** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
-
----
-
-# 1. Overview
-
-Wan-AI/Wan2.2-T2V-A14B-Diffusers is a multimodal text-to-video generation model served through the vLLM-Omni infrastructure.
-
-This document covers:
-
-* Service launch configuration (including acceleration options)
-* Benchmark scripts and usage
-* Dataset and workload settings
-* Performance measurement results
-* Reproducibility guidelines
-
----
-
-# 2. Test Environment
-| Component | Specification |
-|------------|----------------|
-| GPU | NVIDIA A100-SXM4-80GB |
-| Diffusion Attention Backend | FlashAttention |
-
-# 3. Service Launch Configuration
-
-## 3.1 Basic Serving Command
-
-```bash
-vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
-    --port 8091
-```
-
-## 3.2 Key Parameters
-
-| Parameter             | Description              |
-| --------------------- | ------------------------ |
-| `--cfg-parallel-size` | CFG parallelism degree   |
-| `--ulysses-degree`    | Ulysses parallel degree  |
-| `--vae-patch-parallel-size`    | VAE parallel degree  |
-| `--tensor-parallel-size` | Tensor parallelism degree |
-| `--use-hsdp` | Enable Hybrid Sharded Data Parallel to shard model weights across GPUs |
-
-Record these parameters when reporting performance results.
-
----
-
-# 4. Benchmark Script
-
-## 4.1 Benchmark Entry
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend v1/videos \
-    --dataset <DATASET_NAME> \
-    --task t2v \
-    --num-prompts <N> \
-    --max-concurrency <C> \
-    --enable-negative-prompt \
-    --random-request-config <CFG>
-```
-
-## 4.2 Key Benchmark Arguments
-
-| Parameter              | Description                       |
-| ---------------------- | --------------------------------- |
-| `--backend`            | Serving backend (use `v1/videos`) |
-| `--dataset`            | Dataset name (`random` or custom) |
-| `--task`               | Task type (e.g., `t2v`)           |
-| `--num-prompts`        | Total number of requests          |
-| `--max-concurrency`    | Client-side concurrency           |
-| `--random-request-config`| JSON string defining random request |
-
----
-
-# 5. Dataset & Workload Settings
-
-## 5.1 Recommended Evaluation Configurations
-
-### Dataset A (480p)
-
-* Dataset: `random`
-* Task: t2v
-* Concurrency: 1
-* Mix Resolution
-```
-[
-    {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":1}
-]
-```
-### Dataset B (720p)
-
-* Dataset: `random`
-* Task: t2v
-* Concurrency: 1
-* Mix Resolution
-```
-[
-    {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":1}
-]
-```
-### Dataset C (Mix Resolution)
-
-* Dataset: `random`
-* Task: t2v
-* Concurrency: 1
-* Mix Resolution
-```
-[
- {"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":0.15},
- {"width":854,"height":480,"num_inference_steps":4,"num_frames":120,"fps":24,"weight":0.25},
- {"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":0.6}
-]
-```
----
-
-## 5.2 Example Benchmark Command
-
-```bash
-python benchmarks/diffusion/diffusion_benchmark_serving.py \
-    --backend v1/videos \
-    --dataset random \
-    --task t2v \
-    --num-prompts 1 \
-    --max-concurrency 1 \
-    --enable-negative-prompt \
-    --random-request-config '[
-        {"width":854,"height":480,"num_inference_steps":18,"num_frames": 33,"fps":16",weight":1}
-    ]'
-```
-
----
-
-# 6. Performance Metrics
-
-The following metrics are collected during benchmarking:
-
-| Metric             | Description                   | Unit    |
-| ------------------ | ----------------------------- | ------- |
-| Mean Latency        | Mean of latency       | seconds |
-| P99 Latency        | P99 of latency             | seconds |
-
----
-
-# 7. Performance Results
-
-| Dataset Configuration | Max Concur. | CFG | Usp | Tp | Hsdp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
-|-----------------------|-----|-----|-----|-----|----|--------------|------------------|------------------|
-| Dataset A | 1 | 2 | 2 | 1 | On | 1          | 24.6766          | 24.6766          |
-| Dataset A | 1 | 2 | 2 | 1 | On | 4          | 21.6810          | 21.6810          |
-| Dataset B | 1 | 2 | 2 | 1 | On | 1          | 124.6639         | 124.6639          |
-| Dataset B | 1 | 2 | 2 | 1 | On | 4          | 117.44          | 117.44          |
-| Dataset C | 1 | 2 | 2 | 1 | On | 1          | 79.2175        | 124.2565 |
-| Dataset C | 1 | 2 | 2 | 1 | On | 4          | 74.4977        | 117.710 |
----
-
-# 8. Reproducibility Checklist
-
-To ensure consistent and comparable benchmark results:
-
-* Record GPU type
-* Record parallel configuration
-* Record benchmark parameters (resolution, concurrency, number of prompts)
-* Ensure no background workload on GPUs during testing
-
----
-
-This document serves as the official Wan2.2 serving performance reference under vLLM-Omni.

benchmarks/qwen3-omni/README.md

@@ -56,9 +56,9 @@ What it does:
 - Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
 - Uses `benchmarks/build_dataset/top100.txt` and writes to:
   - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
-    - `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
-    - `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
-    - `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
+    - `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
+    - `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
+    - `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
   - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
 
 Key checks:

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -26,12 +26,12 @@ else
                       --log-stats \
                       --log-dir $log_dir
     echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
-    echo "  - omni_pipeline_text                       run dir/base name"
-    echo "  - omni_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
-    echo "  - omni_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
-    echo "  - omni_pipeline_text.stage0.log                per-stage detailed logs"
-    echo "  - omni_pipeline_text.stage1.log"
-    echo "  - omni_pipeline_text.stage2.log"
+    echo "  - omni_llm_pipeline_text                       run dir/base name"
+    echo "  - omni_llm_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
+    echo "  - omni_llm_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
+    echo "  - omni_llm_pipeline_text.stage0.log                per-stage detailed logs"
+    echo "  - omni_llm_pipeline_text.stage1.log"
+    echo "  - omni_llm_pipeline_text.stage2.log"
     echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
     echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
 fi

benchmarks/qwen3-tts/README.md

@@ -32,8 +32,8 @@ bash run_benchmark.sh --hf-only
 # Use a different model (e.g. 1.7B)
 MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
 
-# Use bs16 config for higher throughput
-STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
+# Use batch_size=4 config for higher throughput
+STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
 
 # Custom GPU, prompt count, concurrency levels
 GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
@@ -54,7 +54,7 @@ CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
 ### 2) Run online serving benchmark
 
 ```bash
-python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
+python vllm_omni/bench_tts_serve.py \
     --port 8000 \
     --num-prompts 50 \
     --max-concurrency 1 4 10 \
@@ -65,7 +65,7 @@ python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
 ### 3) Run HuggingFace baseline
 
 ```bash
-python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
+python transformers/bench_tts_hf.py \
     --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
     --num-prompts 50 \
     --gpu-device 0 \
@@ -75,7 +75,7 @@ python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
 ### 4) Generate comparison plots
 
 ```bash
-python benchmarks/qwen3-tts/plot_results.py \
+python plot_results.py \
     --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
     --labels "vllm-omni" "hf_transformers" \
     --output results/comparison.png
@@ -83,14 +83,14 @@ python benchmarks/qwen3-tts/plot_results.py \
 
 ## Stage Configs
 
-| Config | max_num_seqs | Description |
-|--------|:------------:|-------------|
-| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
-| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
+| Config | Batch Size | Description |
+|--------|:----------:|-------------|
+| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lower latency) |
+| `vllm_omni/configs/qwen3_tts_bs4.yaml` | 4 | Concurrent request processing (higher throughput) |
 
-All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
+Both configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
 
-The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
+The model is specified via the CLI `--model` flag (or `MODEL` env var in `run_benchmark.sh`), so the same configs work for both the 0.6B and 1.7B model variants.
 
 ## Metrics
 

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml

@@ -4,18 +4,19 @@ async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
-    is_comprehension: true
     runtime:
       devices: "0"
+      max_batch_size: 1
     engine_args:
-      max_num_seqs: 1
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
       trust_remote_code: true
-      async_scheduling: true
+      async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: latent
       gpu_memory_utilization: 0.3
@@ -38,18 +39,20 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
+      max_batch_size: 1
     engine_args:
-      max_num_seqs: 1
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
-      async_scheduling: true
+      async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: audio
-      gpu_memory_utilization: 0.3
+      gpu_memory_utilization: 0.2
       distributed_executor_backend: "mp"
       max_num_batched_tokens: 8192
       max_model_len: 32768

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs16.yaml

@@ -1,94 +0,0 @@
-# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
-# High-throughput concurrent request processing
-# 2-stage pipeline: Talker -> Code2Wav
-async_chunk: true
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_comprehension: true
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 16
-      model_stage: qwen3_tts
-      model_arch: Qwen3TTSTalkerForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      enforce_eager: false
-      trust_remote_code: true
-      async_scheduling: true
-      enable_prefix_caching: false
-      engine_output_type: latent
-      gpu_memory_utilization: 0.3
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 4096
-      max_model_len: 4096
-      custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
-    output_connectors:
-      to_stage_1: connector_of_shared_memory
-    default_sampling_params:
-      temperature: 0.9
-      top_k: 50
-      max_tokens: 4096
-      seed: 42
-      detokenize: false
-      repetition_penalty: 1.05
-      stop_token_ids: [2150]
-
-  - stage_id: 1
-    stage_type: llm
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 16
-      model_stage: code2wav
-      model_arch: Qwen3TTSCode2Wav
-      worker_type: generation
-      scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
-      enforce_eager: true
-      trust_remote_code: true
-      async_scheduling: true
-      enable_prefix_caching: false
-      engine_output_type: audio
-      gpu_memory_utilization: 0.2
-      distributed_executor_backend: "mp"
-      max_num_batched_tokens: 16384
-      max_model_len: 32768
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: audio
-    input_connectors:
-      from_stage_0: connector_of_shared_memory
-    tts_args:
-      max_instructions_length: 500
-    default_sampling_params:
-      temperature: 0.0
-      top_p: 1.0
-      top_k: -1
-      max_tokens: 65536
-      seed: 42
-      detokenize: true
-      repetition_penalty: 1.0
-
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 16
-
-  connectors:
-    connector_of_shared_memory:
-      name: SharedMemoryConnector
-      extra:
-        shm_threshold_bytes: 65536
-        codec_streaming: true
-        connector_get_sleep_s: 0.01
-        connector_get_max_wait_first_chunk: 3000
-        connector_get_max_wait: 300
-        codec_chunk_frames: 25
-        codec_left_context_frames: 25
-
-  edges:
-    - from: 0
-      to: 1
-      window_size: -1

benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs4.yaml

@@ -1,22 +1,23 @@
 # Qwen3-TTS batch_size=4 config (streaming with async_chunk)
-# Enables concurrent request processing
+# Enables concurrent request processing with max_inflight=4
 # 2-stage pipeline: Talker -> Code2Wav
 async_chunk: true
 stage_args:
   - stage_id: 0
     stage_type: llm
-    is_comprehension: true
     runtime:
       devices: "0"
+      max_batch_size: 4
     engine_args:
-      max_num_seqs: 4
       model_stage: qwen3_tts
       model_arch: Qwen3TTSTalkerForConditionalGeneration
+      hf_overrides:
+        architectures: [Qwen3TTSTalkerForConditionalGeneration]
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       enforce_eager: false
       trust_remote_code: true
-      async_scheduling: true
+      async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: latent
       gpu_memory_utilization: 0.3
@@ -39,15 +40,17 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
+      max_batch_size: 4
     engine_args:
-      max_num_seqs: 4
       model_stage: code2wav
       model_arch: Qwen3TTSCode2Wav
+      hf_overrides:
+        architectures: [Qwen3TTSCode2Wav]
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       enforce_eager: true
       trust_remote_code: true
-      async_scheduling: true
+      async_scheduling: false
       enable_prefix_caching: false
       engine_output_type: audio
       gpu_memory_utilization: 0.2

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.18.0
+ARG VLLM_BASE_TAG=v0.17.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.npu

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.17.0rc1
+ARG VLLM_ASCEND_TAG=v0.14.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
+RUN pip install -v -e .
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+RUN pip install -v -e .
+
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}
 
 COPY . .
 
-RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
-    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
-    source /usr/local/Ascend/nnal/atb/set_env.sh && \
-    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
-    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
+RUN pip install -v -e .
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.17.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=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
 ARG ARG_PYTORCH_ROCM_ARCH
 ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
 

docker/Dockerfile.xpu

@@ -76,7 +76,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.18.0
+ARG VLLM_VERSION=v0.17.0
 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
 WORKDIR /workspace/vllm
 

docs/.nav.yml

@@ -58,9 +58,7 @@ nav:
       - Quantization:
         - Overview: user_guide/diffusion/quantization/overview.md
         - FP8: user_guide/diffusion/quantization/fp8.md
-        - Int8: user_guide/diffusion/quantization/int8.md
         - GGUF: user_guide/diffusion/quantization/gguf.md
-      - Step Execution: user_guide/diffusion/step_execution.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
       - LoRA: user_guide/diffusion/lora.md
@@ -92,7 +90,6 @@ nav:
       - design/feature/cache_dit.md
       - design/feature/teacache.md
       - design/feature/async_chunk_design.md
-      - design/feature/diffusion_step_execution.md
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/api/README.md

@@ -6,13 +6,19 @@ Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
 - [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
+- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
 - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
 - [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
 - [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
 - [vllm_omni.entrypoints.omni.Omni][]
-- [vllm_omni.entrypoints.omni_base.OmniBase][]
+- [vllm_omni.entrypoints.omni.OmniBase][]
+- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
+- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
+- [vllm_omni.entrypoints.omni_stage.OmniStage][]
+- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
+- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
 
 ## Inputs
 
@@ -42,7 +48,9 @@ Engine classes for offline and online inference.
 - [vllm_omni.engine.OmniEngineCoreOutputs][]
 - [vllm_omni.engine.OmniEngineCoreRequest][]
 - [vllm_omni.engine.PromptEmbedsPayload][]
+- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
 - [vllm_omni.engine.arg_utils.OmniEngineArgs][]
+- [vllm_omni.engine.input_processor.OmniInputProcessor][]
 - [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
 - [vllm_omni.engine.output_processor.OmniRequestState][]
 

docs/configuration/README.md

@@ -14,10 +14,6 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
 
 - **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
 
-## Multi-Stage Recipes
-
-- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
-
 ## Optimization Features
 
 - **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss

docs/configuration/gpu_memory_utilization.md

@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
 
 **Solutions:**
 1. Reduce `max_num_batched_tokens`
-2. Reduce `max_num_seqs` in engine_args
+2. Reduce `max_batch_size` in runtime config
 3. Lower `gpu_memory_utilization` slightly
 4. Enable quantization if supported
 

docs/configuration/pd_disaggregation.md

@@ -1,177 +0,0 @@
-# Prefill-Decode (PD) Disaggregation
-
-PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
-stages so prompt processing and token generation can run on different workers.
-
-This is documented as a stage-config recipe instead of a bundled YAML because the
-deployment-specific values usually change per environment:
-
-- GPU placement
-- `tensor_parallel_size`
-- connector backend and connector ports
-- connector IPs or bootstrap addresses
-
-Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
-and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
-changes below.
-
-## Requirements
-
-- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
-- A KV connector supported by vLLM, such as `MooncakeConnector`
-- Matching `tensor_parallel_size` on the prefill and decode thinker stages
-
-## 1. Split the thinker into prefill and decode stages
-
-Replace the original thinker stage with two stages:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    stage_type: llm
-    is_prefill_only: true
-    runtime:
-      devices: "0"
-    engine_args:
-      max_num_seqs: 16
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-      kv_transfer_config:
-        kv_connector: "MooncakeConnector"
-        kv_role: "kv_producer"
-        kv_rank: 0
-        kv_parallel_size: 2
-        kv_connector_extra_config:
-          mooncake_bootstrap_port: 25201
-    final_output: false
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-
-  - stage_id: 1
-    stage_type: llm
-    is_decode_only: true
-    runtime:
-      devices: "1"
-    engine_args:
-      max_num_seqs: 64
-      model_stage: thinker
-      model_arch: Qwen3OmniMoeForConditionalGeneration
-      worker_type: ar
-      scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
-      gpu_memory_utilization: 0.9
-      enforce_eager: true
-      trust_remote_code: true
-      engine_output_type: latent
-      distributed_executor_backend: "mp"
-      enable_prefix_caching: false
-      max_num_batched_tokens: 32768
-      hf_config_name: thinker_config
-      tensor_parallel_size: 1
-      kv_transfer_config:
-        kv_connector: "MooncakeConnector"
-        kv_role: "kv_consumer"
-        kv_rank: 1
-        kv_parallel_size: 2
-        kv_connector_extra_config:
-          mooncake_bootstrap_port: 25202
-    engine_input_source: [0]
-    final_output: true
-    final_output_type: text
-    is_comprehension: true
-    default_sampling_params:
-      temperature: 0.4
-      top_p: 0.9
-      top_k: 1
-      max_tokens: 2048
-      seed: 42
-      detokenize: True
-      repetition_penalty: 1.05
-```
-
-Notes:
-
-- `is_prefill_only: true` marks the thinker stage that only saves KV.
-- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
-- `kv_transfer_config` is required on both stages.
-- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
-  prefill side only processes the prompt and exports KV.
-
-## 2. Shift the downstream stages by one index
-
-After inserting the extra thinker stage, renumber the remaining stages:
-
-```yaml
-  - stage_id: 2
-    runtime:
-      devices: "2"
-    engine_input_source: [1]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
-
-  - stage_id: 3
-    runtime:
-      devices: "2"
-    engine_args:
-      max_num_seqs: 1
-    engine_input_source: [2]
-    custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
-```
-
-Compared with the default Qwen3-Omni config:
-
-- the talker becomes stage `2` instead of stage `1`
-- the code2wav stage becomes stage `3` instead of stage `2`
-- the talker now reads from decode stage `1`
-
-## 3. Add runtime edges for the four-stage pipeline
-
-```yaml
-runtime:
-  enabled: true
-  defaults:
-    window_size: -1
-    max_inflight: 1
-  edges:
-    - from: 0
-      to: 1
-      window_size: -1
-    - from: 1
-      to: 2
-      window_size: -1
-    - from: 2
-      to: 3
-      window_size: -1
-```
-
-## 4. Launch with your custom config
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
-    --stage-configs-path /path/to/qwen3_omni_pd.yaml
-```
-
-## Operational Notes
-
-- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
-  pair. Keep prefill and decode at the same `tensor_parallel_size`.
-- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
-  separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
-- Choose connector ports and addresses that match your deployment. The values
-  shown above are examples only.

docs/configuration/stage_configs.md

@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
 For offline (Assume necessary dependencies have ben imported):
 ```python
 model_name = "Qwen/Qwen2.5-Omni-7B"
-omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni_llm = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
 ```
 
 For online serving:
@@ -30,15 +30,15 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
 
 Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
 ```python
-# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration
       process: true  # Run this stage in a separate process
       devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      max_batch_size: 1 # the batch_size for offline inference
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
-      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
       worker_type: ar # The specific worker used
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
@@ -62,9 +62,9 @@ stage_args:
     runtime:
       process: true
       devices: "1"
+      max_batch_size: 3
     engine_args:
       model_stage: talker
-      max_num_seqs: 3
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -88,9 +88,9 @@ stage_args:
     runtime:
       process: true
       devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      max_batch_size: 1
     engine_args:
       model_stage: code2wav
-      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -159,9 +159,9 @@ Visible devices for this stage, specified as a string. This controls which GPU d
 
 Default: `"0"`
 
-#### `engine_args.max_num_seqs`
+#### `runtime.max_batch_size`
 
-The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
+The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
 
 Default: `1`
 

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -1,12 +1,12 @@
-# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
+# stage config for running qwen2.5-omni with architecture of OmniLLM.
 stage_args:
   - stage_id: 0
     runtime:
       process: true            # Run this stage in a separate process
       devices: "0"            # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
+      max_batch_size: 1
     engine_args:
       model_stage: thinker
-      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -30,9 +30,9 @@ stage_args:
     runtime:
       process: true
       devices: "1"
+      max_batch_size: 1
     engine_args:
       model_stage: talker
-      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -56,9 +56,9 @@ stage_args:
     runtime:
       process: true
       devices: "0"            # Example: use a different GPU than the previous stage; use "0" if single GPU
+      max_batch_size: 1
     engine_args:
       model_stage: code2wav
-      max_num_seqs: 1
       model_arch: Qwen2_5OmniForConditionalGeneration
       worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler

docs/contributing/README.md

@@ -107,7 +107,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
 - `[CI/Build]` for build or continuous integration improvements.
 - `[Doc]` for documentation fixes and improvements.
 - `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
-- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
+- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
 - `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
 - `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
 - `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.

docs/contributing/ci/CI_5levels.md

@@ -168,6 +168,11 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
+│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
+│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
+│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
+│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)

docs/contributing/ci/tests_style.md

@@ -24,6 +24,8 @@ End-to-end tests verify the complete functionality of a system or component. For
 
 - **`tests/e2e/online_serving/`**: Tests for online serving scenarios (e.g., API server tests)
 
+**Example:** The test file for `vllm_omni/entrypoints/omni_llm.py` should be located at `tests/entrypoints/test_omni_llm.py`.
+
 ## Test Directory Structure
 
 The ideal directory structure mirrors the source code organization. Legend: `✅` = test exists, `⬜` = suggested to add.
@@ -73,6 +75,11 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── omni.py                                 │   ├── test_omni.py                    ⬜  (E2E covered by e2e/offline, e2e/online)
+│   ├── omni_llm.py                             │   ├── test_omni_llm.py                ✅
+│   ├── omni_stage.py                            │   ├── test_omni_stage.py              ⬜  (partial in test_omni_stage_diffusion_config.py)
+│   ├── omni_diffusion.py                       │   ├── test_omni_diffusion.py          ✅
+│   ├── async_omni.py                            │   ├── test_async_omni.py              ✅ actually in e2e/online_serving/test_async_omni.py
 │   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
@@ -163,7 +170,7 @@ vllm_omni/                                    tests/
 
 ### Naming Conventions
 
-- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.py`
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
 
 - **E2E Tests**: Place in `tests/e2e/offline_inference/` or `tests/e2e/online_serving/` with descriptive names. Example: `tests/e2e/offline_inference/test_qwen3_omni.py`, `tests/e2e/offline_inference/test_diffusion_model.py`
 

docs/contributing/model/adding_diffusion_model.md

@@ -739,22 +739,6 @@ See detailed guide: [How to add Sequence Parallel support](../../design/feature/
 omni = Omni(model="your-model", ulysses_degree=2, ring_degree=2)
 ```
 
-### Step Execution
-
-See detailed design guide: [How to add step execution support](../../design/feature/diffusion_step_execution.md)
-
-Use this only when your pipeline can be split into stable request-scoped and
-step-scoped phases. The reference implementation is
-`QwenImagePipeline`, which maps its request-level `forward()` into:
-
-1. `prepare_encode()` for prompt encoding, latent init, timestep prep, and per-request scheduler setup.
-2. `denoise_step()` for one transformer/noise prediction.
-3. `step_scheduler()` for one scheduler update and `step_index` advance.
-4. `post_decode()` for the final VAE decode.
-
-Do not enable `step_execution=True` until those four methods are implemented
-and validated against the request-level path.
-
 ### Cache Acceleration
 
 #### TeaCache
@@ -831,107 +815,6 @@ class WanTransformer3DModel(nn.Module):
 
 ---
 
-### Diffusion Timing (Performance Profiling)
-When adapting a new diffusion model, it is often useful to analyze the latency of key components such as text encoding, diffusion denoising, and VAE decoding.
-vLLM-Omni provides a timing utility via `DiffusionPipelineProfilerMixin` to help developers quickly identify performance bottlenecks.
-
-!!! info
-      `DiffusionPipelineProfilerMixin` is different from using `torch.profiler` for diffusion models, as introduced in this [tutorial](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/profiling.md#3-profiling-diffusion-models). `DiffusionPipelineProfilerMixin` only prints the timing information of multiple functions (such as `vae.decode`), while `torch.profiler` saves detailed GPU/CPU computation time, call/execution steps.
-
-This tool automatically measures the execution time of selected pipeline modules and prints the results in the logs.
-
-**Enabling Diffusion Timing**
-
-
-Enable timing by setting:
-```
-vllm serve Qwen/Qwen-Image --omni --port 8091 --enable-diffusion-pipeline-profiler
-```
-You can optionally specify which modules to profile:
-```
-class YourPipeline(xxx, DiffusionPipelineProfilerMixin):
-    def __init__(self, xxx):
-        ...
-        self.setup_diffusion_pipeline_profiler(profiler_targets=["diffuse"], enable_diffusion_pipeline_profiler)
-```
-If not specified, the default targets are used:
-```
-["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
-```
-**Adding DiffusionPipelineProfilerMixin to a Pipeline**
-To enable timing support in your pipeline, inherit from DiffusionPipelineProfilerMixin.
-```python
-from vllm_omni.diffusion.utils.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
-
-class YourModelPipeline(nn.Module, DiffusionPipelineProfilerMixin):
-    # Optional: Specify custom timing targets
-    _PROFILER_TARGETS = ["vae.encode", "vae.decode", "diffuse", "text_encoder.forward", "tokenizer.forward"]
-
-    def __init__(
-        self,
-        *,
-        od_config: OmniDiffusionConfig,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.od_config = od_config
-        self.parallel_config = od_config.parallel_config
-        # initialize pipeline components
-        ...
-
-        # initialize timing profiler
-        self.setup_diffusion_pipeline_profiler(enable_diffusion_pipeline_profiler)
-```
-The mixin dynamically wraps selected methods and records their execution time during inference.
-
-If you need to fetch the execution time of different modules, you will need to pass `self.stage_durations` to `DiffusionOutput`, as shown below:
-
-```diff
--   return DiffusionOutput(output=img)
-+   return DiffusionOutput(
-            output=image, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None
-        )
-```
-
-**Pipeline Design for Timing**
-The current diffusion timing utility is function-based, meaning it measures the execution time of individual methods.
-
-When implementing a new pipeline, avoid putting all logic inside a single function (e.g., forward). Instead, structure the pipeline in a modular way by separating key stages into independent methods, such as the diffusion loop.
-
-For example:
-```
-def forward(self, req: OmniDiffusionRequest):
-    prompt_embeds = self.encode_prompt(req)
-    latents = self.diffuse(prompt_embeds, req)
-    images = self.vae.decode(latents)
-    return DiffusionOutput(output=images)
-```
-This allows the timing utility to measure each stage (e.g., encode_prompt, diffuse, vae.decode) separately and helps identify performance bottlenecks more easily.
-
-
-**Default Profiled Modules**
-
-By default, the following pipeline modules are timed:
-```
-vae.encode
-vae.decode
-diffuse
-text_encoder.forward
-tokenizer.forward
-```
-
-**Example Output**
-
-When enabled, timing logs appear like this:
-```
-[DiffusionTiming] text_encoder.forward took 0.018s
-[DiffusionTiming] diffuse took 2.412s
-[DiffusionTiming] vae.decode took 0.063s
-```
-These measurements help identify bottlenecks during model adaptation and optimization
-
-
-
 ## Troubleshooting
 
 
@@ -983,7 +866,7 @@ hidden_states = hidden_states.reshape(batch_size, seq_len, -1)
 
 1. **Reduce batch size:**
    ```python
-   omni.generate(prompts=[...], max_num_seqs=2)
+   omni.generate(prompts=[...], max_batch_size=2)
    ```
 
 2. **Use smaller image size:**

docs/contributing/model/adding_omni_model.md

@@ -330,46 +330,54 @@ Stage transitions are the mechanism by which outputs from one stage are converte
 
 ### Where Stage Transitions Are Called
 
-Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
+Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
 
-1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
+1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
 2. **Trigger**: When a stage completes processing and produces outputs
 3. **Execution Flow**:
    ```python
-   # In orchestrator.py
-   next_stage_id = stage_id + 1
-   next_client = self.stage_clients[next_stage_id]
-   params = req_state.sampling_params_list[next_stage_id]
+   # In omni_llm.py, _run_generation() method (around line 345-460)
 
-   # Save current stage outputs so stage_input_processors can consume them.
-   self.stage_clients[stage_id].set_engine_outputs([output])
+   # Main orchestrator loop polls each stage for completed requests
+   for stage_id, stage in enumerate(self.stage_list):
+       result = stage.try_collect()  # Get completed request
+       if result is None:
+           continue
 
-   # THIS IS WHERE STAGE TRANSITION HAPPENS
-   next_inputs = next_client.process_engine_inputs(
-       stage_list=self.stage_clients,
-       prompt=req_state.prompt,
-   )
+       # Store outputs from this stage
+       engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
+       stage.set_engine_outputs(engine_outputs)
 
-   # Build and submit request(s) to the next stage.
-   for next_input in next_inputs:
-       request = build_engine_core_request_from_tokens(
-           request_id=req_id,
-           prompt=next_input,
-           params=params,
-           model_config=self.stage_vllm_configs[next_stage_id].model_config,
-       )
-       await next_client.add_request_async(request)
+       # Check if there's a next stage to forward to
+       next_stage_id = stage_id + 1
+       if next_stage_id < len(self.stage_list):
+           next_stage: OmniStage = self.stage_list[next_stage_id]
+
+           # THIS IS WHERE STAGE TRANSITION HAPPENS
+           next_inputs = next_stage.process_engine_inputs(
+               self.stage_list,
+               [request_id_to_prompt[req_id]]
+           )
+
+           # Submit to next stage
+           task = {
+               "type": OmniStageTaskType.GENERATE,
+               "request_id": req_id,
+               "engine_inputs": next_inputs[0],
+               "sampling_params": sampling_params_list[next_stage_id],
+           }
+           next_stage.submit(task)
    ```
 
 ### How Stage Transitions Work
 
 The stage transition process follows these steps:
 
-1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_engine_outputs(...)`
+1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
 
 2. **Transition Detection**: The orchestrator checks if there's a next stage and calls `process_engine_inputs()` on it
 
-3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
+3. **Input Processing**: The `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
    ```python
    def process_engine_inputs(
        self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None

docs/contributing/model/adding_tts_model.md

@@ -361,9 +361,9 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
+      max_batch_size: 64
     engine_args:
       model_stage: ar_stage
-      max_num_seqs: 64
       model_arch: YourTTSModelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
@@ -403,9 +403,9 @@ stage_args:
     stage_type: llm
     runtime:
       devices: "0"
+      max_batch_size: 64
     engine_args:
       model_stage: ar_stage
-      max_num_seqs: 64
       model_arch: YourTTSModelForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler

docs/contributing/profiling.md

@@ -23,32 +23,30 @@ export VLLM_PROFILER_MAX_ITERS=1
 The profiler is default to function across all stages. But It is highly recommended to profile specific stages by passing the stages list, preventing from producing too large trace files:
 ```python
 # Profile all stages
-omni.start_profile()
+omni_llm.start_profile()
 
 # Only profile Stage 1
-omni.start_profile(stages=[1])
+omni_llm.start_profile(stages=[1])
 ```
 
 ```python
 # Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
-omni.start_profile(stages=[0, 2])
+omni_llm.start_profile(stages=[0, 2])
 ```
 
 **Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni.entrypoints.omni import Omni
-
-omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+from vllm_omni import omni_llm
 
 profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
 # 1. Start profiling if enabled
 if profiler_enabled:
-    omni.start_profile(stages=[0])
+    omni_llm.start_profile(stages=[0])
 
 # Initialize generator
-omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
 total_requests = len(prompts)
 processed_count = 0
@@ -59,21 +57,21 @@ for stage_outputs in omni_generator:
     # ... [Output processing logic for text/audio would go here] ...
 
     # Update count to track when to stop profiling
-    processed_count += 1
+    processed_count += len(stage_outputs.request_output)
 
     # 2. Check if all requests are done to stop the profiler safely
     if profiler_enabled and processed_count >= total_requests:
         print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
 
         # Stop the profiler while workers are still active
-        omni.stop_profile()
+        omni_llm.stop_profile()
 
         # Wait for traces to flush to disk
         print("[Info] Waiting 30s for workers to write trace files to disk...")
         time.sleep(30)
         print("[Info] Trace export wait time finished.")
 
-omni.close()
+omni_llm.close()
 ```
 
 

docs/design/architecture_overview.md

@@ -67,12 +67,12 @@ According to analysis for current popular open-source models, most of them have
 | Component         | Description                                                                                                                              |
 | ----------------- | ---------------------------------------------------------------------------------------------------------------------------------------- |
 | **OmniRouter**    | provide an intelligent router for Omni-modality requests dispatch                                                                        |
-| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni) and provide the OmniStage abstraction for different AR/DiT stages |
 | **AR**            | adapted for omni-modality models while inheriting efficient features from vLLM, such as cache management                                 |
 | **Diffusion**     | natively implemented and optimized using acceleration components                                                                         |
 | **OmniConnector** | supports fully disaggregation based on E/P/D/G (Encoding/Processing/Decoding/Generation) disaggregation across stages                    |
 
-Disaggregated stages are managed through stage configuration. In Qwen3-Omni, Thinker/Talker/Code2wav are declared as separate configured stages, and runtime routing is handled by `Orchestrator` over `StageEngineCoreClient` / `StageDiffusionClient`.
+Disaggregated stages are managed through configuration, such as in the Qwen3-Omni example, where stages like Thinker, Talker, and Code2wav are defined as separate OmniStage instances with specific resources and input/output type.
 
 ## Main features
 
@@ -127,10 +127,10 @@ Taking **Qwen3-Omni** as an example:
 The **Omni** class provides a Python interface for offline batched inference. Users initialize the Omni class with a Hugging Face model name and use the generate method, passing inputs that include both text prompts and multi-modal data:
 
 ```
-# Create an omni runtime with HF model name.
+# Create an omni_lm with HF model name.
 from vllm_omni.entrypoints.omni import Omni
 
-omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 # Example prompts.
 om_inputs = {"prompt": prompt,
@@ -140,7 +140,7 @@ om_inputs = {"prompt": prompt,
              }}
 
 # Generate texts and audio from the multi-modality inputs.
-outputs = omni.generate(om_inputs, sampling_params_list)
+outputs = omni_lm.generate(om_inputs, sampling_params_list)
 ```
 
 ## Online Serving

docs/design/feature/async_chunk_design.md

@@ -171,7 +171,7 @@ stage_args:
 - `custom_process_next_stage_input_func: str`: Path to custom chunk processing function; receives `(transfer_manager, pooling_output, request)`. For qwen3-omni: `thinker2talker_async_chunk`, `talker2code2wav_async_chunk`
 - `stage_connector_config: dict`: Connector configuration
 - `worker_type: str`: Model type, e.g. `"ar"` or `"generation"` (used by OmniChunkTransferAdapter for mode-specific payload handling)
-- `max_num_seqs: int`: Maximum number of sequences for concurrent processing in the stage
+- `max_batch_size: int`: Maximum batch size for the stage
 
 
 ### Connector Configuration
@@ -195,9 +195,9 @@ stage_args:
   - stage_id: 2  # code2wav stage
     runtime:
       devices: "1"
+      max_batch_size: 64  # Enables batched audio generation
     engine_args:
       model_stage: code2wav
-      max_num_seqs: 64  # Enables batched audio generation
 ```
 
 ## Related Files

docs/design/feature/diffusion_step_execution.md

@@ -1,121 +0,0 @@
-# Diffusion Step Execution
-
-This guide documents vLLM-Omni's stepwise diffusion contract for model authors
-and contributors implementing `step_execution=True` support for a diffusion
-pipeline.
-
-For end-user enablement, supported models, and current limitations, see
-[Step Execution](../../user_guide/diffusion/step_execution.md).
-
-## Current Support Scope
-
-`step_execution` is **not** a generic diffusion toggle. It only works for
-pipelines that implement the segmented stateful contract in
-[`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py).
-
-Current in-tree support:
-
-| Pipeline | Example models | Step execution |
-|----------|----------------|----------------|
-| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
-| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
-
-Current engine/runtime limitations:
-
-- `StepScheduler` only schedules `batch_size=1`.
-- `cache_backend` is not supported in step mode.
-- Request-mode extras such as KV transfer are not wired into step mode yet.
-- Unsupported pipelines now fail early during model loading instead of failing on the first request.
-
-## Execution Contract
-
-Step mode is driven by four pipeline methods plus the shared mutable request
-state object:
-
-- `prepare_encode(state)`: one-time request preparation.
-- `denoise_step(state)`: compute the noise prediction for the current step.
-- `step_scheduler(state, noise_pred)`: mutate latents and advance step state.
-- `post_decode(state)`: decode the final output after denoising is complete.
-
-The state lives in
-[`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
-as `DiffusionRequestState`. Store request-scoped tensors there, or use
-`state.extra` for model-specific fields that do not justify extending the core
-dataclass.
-
-The worker-side step loop lives in
-[`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py):
-
-1. `prepare_encode()` runs once for a new request.
-2. `denoise_step()` runs every scheduler tick.
-3. `step_scheduler()` mutates `state.latents` and advances `state.step_index`.
-4. `post_decode()` runs exactly once after `state.denoise_completed` becomes true.
-
-## Recommended Split
-
-When converting an existing request-level `forward()` pipeline, keep the split
-strict and mechanical:
-
-| Request-level phase | Stepwise method | What belongs there |
-|---------------------|-----------------|--------------------|
-| Input validation, prompt encoding, latent init, timestep prep, per-request scheduler creation | `prepare_encode()` | Anything that should happen once per request |
-| Transformer forward / noise prediction | `denoise_step()` | Pure denoise computation for the current timestep |
-| `scheduler.step(...)` and `step_index += 1` | `step_scheduler()` | Only latent/state mutation for one step |
-| VAE decode / postprocess | `post_decode()` | Final decode only |
-
-Keep the stepwise path reusing the same helpers as the request-level path
-whenever possible. Reimplementing the denoise loop from scratch is the easiest
-way to introduce behavioral drift.
-
-## Qwen-Image Reference
-
-[`pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)
-is the reference implementation and is split correctly for the current
-contract:
-
-- `prepare_encode()` reuses `_prepare_generation_context()` so prompt encoding,
-  latent init, timestep creation, CFG setup, and shape bookkeeping stay aligned
-  with `forward()`.
-- `prepare_encode()` deep-copies `self.scheduler` **after**
-  `prepare_timesteps()` so request-specific scheduler state is isolated.
-- `denoise_step()` reuses `_build_denoise_kwargs()` plus
-  `predict_noise_maybe_with_cfg()`, so sequential CFG, CFG-parallel, and
-  non-CFG behavior stay identical to the request-level path.
-- `step_scheduler()` only calls
-  `scheduler_step_maybe_with_cfg(..., per_request_scheduler=state.scheduler)`
-  and increments `state.step_index`.
-- `post_decode()` reuses `_decode_latents()`, so the final image decode matches
-  the normal `forward()` path.
-
-That decomposition is the target pattern for future models.
-
-## Rules For New Pipelines
-
-- Do not keep request-scoped scheduler state on `self.scheduler`. Copy it into
-  `state.scheduler` during `prepare_encode()`.
-- Do not mutate `state.step_index` inside `denoise_step()`. Only
-  `step_scheduler()` should advance the step.
-- Do not decode partial outputs in `denoise_step()` or `step_scheduler()`.
-- If the request-level pipeline has condition latents, masks, or edit-specific
-  tensors, store them in `state` or `state.extra`, not in global pipeline
-  attributes.
-- Preserve CFG behavior by sharing the same helper path used by `forward()`.
-- Keep `post_decode()` equivalent to the tail of `forward()`.
-
-## Validation Checklist
-
-Before marking a pipeline as `supports_step_execution = True`, verify:
-
-- Stepwise output matches request-level output for the same seed and sampling params.
-- Per-request scheduler state is isolated across concurrent requests.
-- Abort during denoise does not leak cached state.
-- `step_index` reported by `RunnerOutput` matches the scheduler progress.
-- CFG-parallel and non-CFG paths both work if the request-level pipeline supports them.
-
-## Related Files
-
-- Contract: [`vllm_omni/diffusion/models/interface.py`](gh-file:vllm_omni/diffusion/models/interface.py)
-- State: [`vllm_omni/diffusion/worker/utils.py`](gh-file:vllm_omni/diffusion/worker/utils.py)
-- Runner loop: [`vllm_omni/diffusion/worker/diffusion_model_runner.py`](gh-file:vllm_omni/diffusion/worker/diffusion_model_runner.py)
-- Scheduler transport: [`vllm_omni/diffusion/sched/interface.py`](gh-file:vllm_omni/diffusion/sched/interface.py)
-- Reference pipeline: [`vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py`](gh-file:vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py)

docs/design/module/ar_module.md

@@ -85,6 +85,9 @@ classDiagram
     class InputProcessor {
         +process_inputs() EngineCoreRequest
     }
+    class OmniInputProcessor {
+        +process_inputs() OmniEngineCoreRequest
+    }
 
     class VLLMOutputProcessor {
         +process_outputs() OutputProcessorOutput
@@ -93,6 +96,7 @@ classDiagram
         +process_outputs() OutputProcessorOutput
         +_route_and_normalize()
     }
+    InputProcessor <|-- OmniInputProcessor
     VLLMOutputProcessor <|-- MultimodalOutputProcessor
 ```
 
@@ -101,7 +105,7 @@ classDiagram
 - **Scheduler**: `OmniARScheduler` extends `vllm.v1.core.sched.scheduler.Scheduler` to enrich scheduled requests with omni-specific payloads
 - **Worker**: `GPUARWorker` extends `vllm.v1.worker.gpu_worker.Worker` to initialize AR-specific model runners
 - **ModelRunner**: `GPUARModelRunner` extends `OmniGPUModelRunner` → `vllm.v1.worker.gpu_model_runner.GPUModelRunner` to expose hidden states and handle multimodal outputs
-- **InputProcessor**: Stage-0 uses upstream `vllm.v1.engine.input_processor.InputProcessor`; `AsyncOmniEngine` then restores omni-specific payloads (for example `additional_information` and `prompt_embeds`) when building `OmniEngineCoreRequest`
+- **InputProcessor**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
 - **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
 
 ## 3. Scheduler Design
@@ -114,7 +118,7 @@ The following diagram illustrates the request flow through the AR module compone
 
 ```mermaid
 flowchart TD
-    A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
+    A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
     B -->|schedule: OmniNewRequestData| C[GPUARWorker]
     C -->|SchedulerOutput| D[GPUARModelRunner]
     D -->|execute_model: None| E[Model Forward Pass]
@@ -293,17 +297,15 @@ The input/output processing pipeline handles serialization, routing, and accumul
 ```mermaid
 sequenceDiagram
     participant Client
-    participant AsyncOmniEngine
-    participant InputProcessor
+    participant OmniInputProcessor
     participant Scheduler
     participant ModelRunner
     participant MultimodalOutputProcessor
     participant Client
 
-    Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
-    AsyncOmniEngine->>InputProcessor: process_inputs()
-    InputProcessor->>Scheduler: EngineCoreRequest
-    AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
+    Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
+    OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
+    OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
     Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
     ModelRunner->>ModelRunner: Decode payloads → CPU tensors
     ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
@@ -316,18 +318,43 @@ sequenceDiagram
     MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
 ```
 
-### Stage-0 Input Processing
+### OmniInputProcessor
 
-Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
+`OmniInputProcessor` extends the base input processor to serialize prompt embeddings and additional information for inter-stage transfer.
+
+#### Payload Serialization
+
+Converts PyTorch tensors to serialized payloads:
 
 ```python
-request = self.input_processor.process_inputs(
-    request_id=request_id,
-    prompt=prompt,
-    params=params,
-    supported_tasks=self.supported_tasks,
-)
-request = _upgrade_to_omni_request(request, prompt)
+def process_inputs(self, ...) -> OmniEngineCoreRequest:
+    # Serialize prompt_embeds
+    if "prompt_embeds" in decoder_inputs:
+        pe_cpu = decoder_inputs["prompt_embeds"].detach().to("cpu").contiguous()
+        prompt_embeds_payload = PromptEmbedsPayload(
+            data=pe_cpu.numpy().tobytes(),
+            shape=[seq_len, hidden_size],
+            dtype=dtype_str,
+        )
+
+    # Serialize additional_information
+    if "additional_information" in decoder_inputs:
+        entries = {}
+        for key, value in raw_info.items():
+            if isinstance(value, torch.Tensor):
+                entry = AdditionalInformationEntry(
+                    tensor_data=value.numpy().tobytes(),
+                    tensor_shape=list(value.shape),
+                    tensor_dtype=dtype_str,
+                )
+            entries[key] = entry
+        additional_information_payload = AdditionalInformationPayload(entries=entries)
+
+    return OmniEngineCoreRequest(
+        # ... standard fields ...
+        prompt_embeds=prompt_embeds_payload,
+        additional_information=additional_information_payload,
+    )
 ```
 
 ### MultimodalOutputProcessor
@@ -373,13 +400,13 @@ The AR module of vLLM-Omni extends vLLM through strategic inheritance and minima
 ### Key Design Patterns
 
 1. **Inheritance over composition**: Extends vLLM classes to preserve compatibility with existing scheduling, batching, and execution mechanisms
-2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage transfer
+2. **Payload serialization**: Uses `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
 3. **Two-phase execution**: Maintains vLLM's execute/sample separation for AR models while supporting single-phase execution for generation models
 4. **Multimodal routing**: Routes outputs by `output_type` and accumulates tensors incrementally to support streaming
 
 ### Differences from vLLM
 
-- **Payload support**: Serialized additional information and prompt embeddings enable direct transfer between pipeline stages
+- **Payload support**: Serialized prompt embeddings and additional information enable direct transfer between pipeline stages
 - **Multimodal handling**: Extended input/output processors support images, audio, and other modalities alongside text
 - **Hidden state exposure**: AR model runners expose per-request hidden states via `pooler_output` for downstream consumption
 - **Generation scheduler**: Fast-path scheduling for basic heterogeneous architectures that complete in one step

docs/design/module/async_omni_architecture.md

@@ -1,203 +0,0 @@
-# AsyncOmni Architecture (Qwen3-Omni Example)
-
-## 1. System Architecture
-
-```text
-• ┌─────────────────────────────────────────────────────────────────────────────────┐
-  │                                    API Layer                                    │
-  │  ┌─────────────────────────────────────┐  ┌──────────────────────────────────┐  │
-  │  │ AsyncOmni (EngineClient)            │  │ Omni                             │  │
-  │  │ • generate() / abort() / shutdown() │  │ • generate()                     │  │
-  │  │ • _final_output_handler()           │  │                                  |  │
-  │  └─────────────────────────────────────┘  └──────────────────────────────────┘  │
-  ├─────────────────────────────────────────────────────────────────────────────────┤
-  │                              Engine Layer (Proxy)                               │
-  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
-  │  │ AsyncOmniEngine                                                           │  │
-  │  │ • _bootstrap_orchestrator() & _initialize_stages()                        │  │
-  │  │ • add_request() / add_request_async() -> input_processor.process_inputs() │  │
-  │  │ • try_get_output() / try_get_output_async()                               │  │
-  │  └───────────────────┬─────────────────────────────────▲─────────────────────┘  │
-  │         request_queue (janus.Queue)        output_queue (janus.Queue)           │
-  ├──────────────────────┼─────────────────────────────────┼────────────────────────┤
-  │                      ▼        Orchestration Layer      │                        │
-  │  ┌───────────────────────────────────────────────────────────────────────────┐  │
-  │  │ Orchestrator [background thread]                                          │  │
-  │  │ • _request_handler()                                                      │  │
-  │  │     -  stage_client.add_request_async() & _prewarm_async_chunk_stages()   │  │
-  │  │ • _orchestration_output_handler()                                         │  │
-  │  │     -  _process_stage_outputs() -> output_processors[i].process_outputs() │  │
-  │  │     -  _route_output() & _forward_to_next_stage()                         │  │
-  │  └──────────┬─────────────────────────┬────────────────────────┬─────────────┘  │
-  ├─────────────┼─────────────────────────┼────────────────────────┼────────────────┤
-  │             │                 Communication Layer              │                │
-  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
-  │  │ StageEngineCoreClient │ │ StageEngineCoreClient │ │ StageDiffusionClient  │  │
-  │  │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │ │ • ZMQ ROUTER / PULL   │  │
-  │  │ • Msgpack codec       │ │ • Msgpack codec       │ │ • Msgpack codec       │  │
-  │  └──────────┬────────────┘ └──────────┬────────────┘ └──────────┬────────────┘  │
-  │             ▼ ZMQ IPC                 ▼ ZMQ IPC                 ▼ ZMQ IPC       │
-  ├─────────────────────────────────────────────────────────────────────────────────┤
-  │                                 Execution Layer                                 │
-  │  ┌───────────────────────┐ ┌───────────────────────┐ ┌───────────────────────┐  │
-  │  │ StageCoreProc         │ │ StageCoreProc         │ │ DiffusionEngine       │  │
-  │  │ [background process]  │ │ [background process]  │ │ [background process]  │  │
-  │  └───────────────────────┘ └───────────────────────┘ └───────────────────────┘  │
-  └─────────────────────────────────────────────────────────────────────────────────┘
-```
-
-## 2. Execution Flow (Arrow Steps, one generate request)
-
-```text
-[1] App
-    -> AsyncOmni.generate(prompt, request_id)
-
-[2] AsyncOmni
-    -> _final_output_handler()   (started on first request)
-    -> AsyncOmniEngine.add_request(stage_id=0, ...)
-
-[3] AsyncOmniEngine.add_request
-    -> (if stage-0 is llm and input is not EngineCoreRequest)
-       InputProcessor.process_inputs()
-       OutputProcessor[0].add_request()
-    -> request_queue.put(add_request_msg)
-
-[4] Orchestrator._request_handler
-    -> _handle_add_request(msg)
-    -> stage_clients[0].add_request_async(...)
-
-[5] Orchestrator._orchestration_loop (loop)
-    -> poll stage output
-       - llm stage: await get_output_async()
-       - diffusion stage: get_diffusion_output_async()
-    -> (llm stage) output_processors[i].process_outputs(...)
-    -> _route_output(...)
-    -> if finished and not final_stage and non-async-chunk:
-         _forward_to_next_stage(...)
-         -> next_stage.add_request_async(...)
-    -> output_queue.put(output)
-
-[6] AsyncOmni._final_output_loop (background coroutine)
-    -> AsyncOmniEngine.try_get_output_async()
-    -> route by request_id to ClientRequestState.queue
-
-[7] AsyncOmni._process_orchestrator_results
-    -> read from ClientRequestState.queue
-    -> _process_single_result(...)
-    -> yield OmniRequestOutput
-
-[8] Exit condition
-    -> receive result["finished"] == True
-    -> generate() ends
-```
-
-## 3. Runtime Sequence (one generate request)
-
-```mermaid
-sequenceDiagram
-    participant APP as App
-    participant AO as AsyncOmni
-    participant ENG as AsyncOmniEngine
-    participant ORCH as Orchestrator
-    participant S0 as Stage-0 Client
-    participant SN as Next Stage Client
-
-    APP->>AO: generate
-    AO->>AO: start output_handler once
-    AO->>ENG: add_request(stage_id=0, ...)
-    ENG->>ENG: input_processor.process_inputs()
-    ENG->>ORCH: request_queue.put(add_request)
-
-    ORCH->>ORCH: _handle_add_request
-    ORCH->>S0: add_request_async
-
-    loop poll route forward
-        ORCH->>S0: get_output_async / get_diffusion_output_async
-        ORCH->>ORCH: _route_output
-        alt need forward to next stage
-            ORCH->>SN: add_request_async
-        end
-        ORCH-->>ENG: output_queue.put
-    end
-
-    AO->>ENG: try_get_output_async
-    ENG-->>AO: message
-    AO-->>APP: yield OmniRequestOutput
-```
-
-## 4. Comparison
-
-Previous topology (reference):
-
-```text
-┌────────────────────────────────────────────────────────────────────────────┐
-│ Main Process                                                               │
-│  ┌──────────────────────┐   ┌────────────────────────────────────────────┐ │
-│  │ generate()           │   │ final_output_handler()                     │ │
-│  └──────────────────────┘   └────────────────────────────────────────────┘ │
-└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
-  mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)    mp.Queue (in_q/out_q)
-           ▼▲                        ▼▲                        ▼▲
-┌───────────────────────┐  ┌───────────────────────┐  ┌──────────────────────┐
-│ Worker Proc-0         │  │ Worker Proc-1         │  │ Worker Proc-2        │
-│ (Thinker LLM)         │  │ (Talker LLM)          │  │ (Vocoder)            │
-│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
-│  │_stage_worker   │   │  │  │_stage_worker   │   │  │  │_stage_worker   │  │
-│  │_async()        │   │  │  │_async()        │   │  │  │_async()        │  │
-│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
-│  ┌────────────────┐   │  │  ┌────────────────┐   │  │  ┌────────────────┐  │
-│  │output_handler()│   │  │  │output_handler()│   │  │  │output_handler()│  │
-│  └────────────────┘   │  │  └────────────────┘   │  │  └────────────────┘  │
-└──────────┬────────────┘  └──────────┬────────────┘  └──────────┬───────────┘
-       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ
-┌──────────────────────┐   ┌──────────────────────┐   ┌──────────────────────┐
-│ EngineCore Proc-0    │   │ EngineCore Proc-1    │   │ EngineCore Proc-2    │
-│ (Thinker)            │   │ (Talker)             │   │ (Vocoder)            │
-└──────────────────────┘   └──────────────────────┘   └──────────────────────┘
-```
-
-Current topology:
-
-```text
-┌────────────────────────────────────────────────────────────────────────────┐
-│ Main Process                                                               │
-│  ┌──────────────────────────────────────────────────────────────────────┐  │
-│  │ Main Thread                                                          │  │
-│  │  ┌──────────────────────┐   ┌─────────────────────────────────────┐  │  │
-│  │  │ generate()           │   │ final_output_handler()              │  │  │
-│  │  └──────────────────────┘   └─────────────────────────────────────┘  │  │
-│  └──────────────────────────────────────────────────────────────────────┘  │
-│         janus.Queue (request_queue) ▼  ▲ janus.Queue (output_queue)        │
-│  ┌──────────────────────────────────────────────────────────────────────┐  │
-│  │ Orchestrator Thread                                                  │  │
-│  │  ┌──────────────────────┐  ┌──────────────────────────────────────┐  │  │
-│  │  │ _request_handler()   │  │ _orchestration_output_handler()      │  │  │
-│  │  └──────────────────────┘  └──────────────────────────────────────┘  │  │
-│  │  ┌────────────────────────────────────────────────────────────────┐  │  │
-│  │  │ _orchestration_loop(): poll/process/route outputs for all stages│  │  │
-│  │  └────────────────────────────────────────────────────────────────┘  │  │
-│  └───────┬─────────────────────────┬─────────────────────────┬──────────┘  │
-└──────────┬─────────────────────────┬─────────────────────────┬─────────────┘
-       ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ               ZMQ ▼ ▲ ZMQ  
-  ┌──────────────────────┐  ┌──────────────────────┐  ┌──────────────────────┐
-  │ EngineCore Proc-0    │  │ EngineCore Proc-1    │  │ EngineCore Proc-2    │
-  │ (Thinker)            │  │ (Talker)             │  │ (Vocoder)            │
-  └──────────────────────┘  └──────────────────────┘  └──────────────────────┘
-```
-
-
-Test scripts:
-```bash
-# enter offline inference folder.
-cd examples/offline_inference/qwen2_5_omni
-python end2end.py --output-dir output_audio --query-type use_mixed_modalities
-
-cd ../qwen3_omni
-python end2end.py --output-dir output_audio --query-type text --async-chunk --enable-stats
-
-cd ../bagel
-python end2end.py --prompts "A cute cat"
-
-cd ../text_to_image
-python text_to_image.py --prompt "a cup of coffee on the table" --output output.png
-```

docs/features/custom_pipeline.md

@@ -97,7 +97,7 @@ outputs = omni.generate(
 )
 
 # Access custom trajectory data
-output = outputs[0].request_output
+output = outputs[0].request_output[0]
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```

docs/getting_started/installation/gpu.md

@@ -30,7 +30,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, and 0.18.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 only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0. For the latest version, 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

@@ -34,13 +34,13 @@ uv pip install vllm-omni
 # --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.18.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.18.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 ```
 
-The 0.18.0 release of vLLM ships CUDA 12.9-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 release 0.17.0 of vLLM is based on PyTorch 2.10.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -56,11 +56,11 @@ 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.18.0
+git checkout v0.17.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
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.17.0/vllm-0.17.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -91,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.18.0 \
+    vllm/vllm-omni:v0.16.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --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.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
 #### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
 # --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.18.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.17.0 release version of the library
 
 ```bash
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.17.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 ```
 
-The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The release 0.17.0 of vLLM requires ROCm 7.0 environment.
 
 #### 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.18.0
+git checkout v0.17.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.18.0 \
+  vllm/vllm-omni-rocm:v0.16.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.18.0
+  vllm/vllm-omni-rocm:v0.16.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/installation/npu/npu.inc.md

@@ -33,25 +33,18 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
-cd /vllm-workspace/vllm
-git pull origin main
-git fetch origin --tags
-git checkout v0.16.0
-
 # Because vllm-ascend will release v0.16.0rc1 after vllm-omni 0.16.0,
 # we have to pin vllm-ascend at the current commit.
 cd /vllm-workspace/vllm-ascend
-git pull origin main
 git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
 pip install -v -e .
 
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.16.0 https://github.com/vllm-project/vllm-omni.git
-cd vllm-omni
-pip install -v -e . --no-build-isolation
-# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
 
+cd vllm-omni
+pip install -v -e .
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
@@ -68,22 +61,22 @@ We are keeping [issue #886](https://github.com/vllm-project/vllm-omni/issues/886
 You can also build vLLM-Omni from the latest main branch if you want to use the latest features or bug fixes. (But sometimes it will break for a while. You can check [issue #886](https://github.com/vllm-project/vllm-omni/issues/886) for the status of the latest commit of vLLM-Omni main branch on NPU.)
 
 ```bash
-# Pin vLLM version to 0.18.0
+# Pin vLLM version to 0.16.0
 cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
-git checkout v0.18.0
+git checkout v0.16.0
 VLLM_TARGET_DEVICE=empty pip install -v -e .
 
 # Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
-cd /vllm-workspace/vllm-ascend
+cd ../vllm-ascend
 git pull origin main
 git fetch origin --tags
-git checkout 1e05c4908f31737bc4eef865a9f351d030a77c9d
+git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
-cd /vllm-workspace/vllm-omni
+cd ../vllm-omni
 git clone https://github.com/vllm-project/vllm-omni.git
 pip install -v -e . --no-build-isolation
 # or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .

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.18.0 --torch-backend=auto
+uv pip install vllm==0.17.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
+uv pip install vllm==0.17.0 --extra-index-url https://wheels.vllm.ai/rocm/0.17.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
@@ -42,7 +42,7 @@ if __name__ == "__main__":
     omni = Omni(model="Tongyi-MAI/Z-Image-Turbo")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output.images
+    images = outputs[0].request_output[0].images
     images[0].save("coffee.png")
 ```
 
@@ -70,7 +70,7 @@ if __name__ == "__main__":
     ]
     omni_outputs = omni.generate(prompts)
     for i_prompt, prompt_output in enumerate(omni_outputs):
-        this_request_output = prompt_output.request_output
+        this_request_output = prompt_output.request_output[0]
         this_images = this_request_output.images
         for i_image, image in enumerate(this_images):
             image.save(f"p{i_prompt}-img{i_image}.jpg")
@@ -82,7 +82,7 @@ if __name__ == "__main__":
 
 !!! info
 
-    For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
     batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 

docs/mkdocs/hooks/generate_api_readme.py

@@ -141,7 +141,6 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
                 "vllm_omni.diffusion.models.qwen_image",
                 "vllm_omni.entrypoints.async_diffusion",
                 "vllm_omni.entrypoints.openai",
-                "vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
             ]
             if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
                 continue

docs/models/supported_models.md

@@ -28,7 +28,7 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
-| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
+| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` |
@@ -52,9 +52,6 @@ th {
 |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
 |`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
 |`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
-| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` |
-| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` |
-|`VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` |
 
 
 ## List of Supported Models for NPU
@@ -78,7 +75,7 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
-| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
+| `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
 |`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |

docs/serving/speech_api.md

@@ -4,7 +4,6 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
 
 - **Qwen3-TTS** (`Qwen/Qwen3-TTS-12Hz-*`) -- Qwen3-based TTS with CustomVoice, VoiceDesign, and Base (voice cloning) task types. Output: 24 kHz.
 - **Fish Speech S2 Pro** (`fishaudio/s2-pro`) -- Dual-AR TTS with DAC codec. Supports text-to-speech and voice cloning via reference audio. Output: 44.1 kHz.
-- **Voxtral TTS** (`mistralai/Voxtral-4B-TTS-2603`) -- AR + FlowMatching TTS with preset voices. Output: 24 kHz.
 
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
@@ -29,14 +28,6 @@ vllm-omni serve fishaudio/s2-pro \
     --trust-remote-code \
     --enforce-eager \
     --gpu-memory-utilization 0.9
-
-# Voxtral TTS
-vllm serve mistralai/Voxtral-4B-TTS-2603 \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
-    --omni \
-    --port 8091 \
-    --trust-remote-code \
-    --enforce-eager
 ```
 
 ### Generate Speech
@@ -356,12 +347,6 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 
 Fish Speech uses `ref_audio` and `ref_text` for voice cloning (no `task_type` needed). The `voice` field should be set to `"default"`. See the [Fish Speech online serving example](../user_guide/examples/online_serving/fish_speech.md) for details.
 
-### Voxtral TTS
-
-| Model | Description |
-|-------|-------------|
-| `mistralai/Voxtral-4B-TTS-2603` | 3B AR + FlowMatching TTS. Supports text-to-speech with preset voices. |
-
 ## Error Responses
 
 ### 400 Bad Request

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -24,22 +24,22 @@ The following table shows which models are currently supported by parallelism me
 
 ### ImageGen
 
-| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel | HSDP |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|:----:|
-| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |      N/A        | ❌   |
-| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |      N/A        | ❌   |
-| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |      N/A        | ❌   |
-| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |      N/A        | ❌   |
-| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |      N/A        | ✅   |
-| **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      N/A        | ✅   |
-| **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |      ✅        | ❌   |
-| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |      N/A        | ❌   |
+| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel | VAE-Patch-Parallel | Expert-Parallel |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
+| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |         ❌          |        N/A        |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |         ✅          |        N/A        |
+| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |         ✅          |        N/A        |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ✅        |         ✅          |        N/A        |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ✅       |        ✅        |         ❌          |        N/A        |
+| **FLUX.2-dev**           | `black-forest-labs/FLUX.2-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        N/A        |
+| **HunyuanImage3.0**      | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        ✅        |
+| **DreamID-Omni**           | `XuGuo699/DreamID-Omni`       |     ❌      |    ❌    |      ✅       |        ❌        |         ❌          |        N/A        |
 
 !!! note "TP Limitations for Diffusion Models"
     We currently implement Tensor Parallelism (TP) only for the DiT (Diffusion Transformer) blocks. This is because the `text_encoder` component in vLLM-Omni uses the original Transformers implementation, which does not yet support TP.
@@ -58,8 +58,6 @@ The following table shows which models are currently supported by parallelism me
 
 | Model | Model Identifier | Ulysses-SP | Ring-Attention | Tensor-Parallel | HSDP | VAE-Patch-Parallel |
 |-------|------------------|:----------:|:--------------:|:---------------:|:----:| :----:|
-| **Wan2.1** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.1** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **LTX-2** | `Lightricks/LTX-2` | ✅ | ✅ | ✅ | ❌ | ❌ |
 

docs/user_guide/diffusion/quantization/fp8.md

@@ -63,7 +63,6 @@ The available `ignored_layers` names depend on the model architecture (e.g., `to
 | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
 | Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Skip sensitive layers | `img_mlp` |
 | Flux | `black-forest-labs/FLUX.1-dev` | All layers | None |
-| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | All layers | None |
 
 ## Combining with Other Features
 

docs/user_guide/diffusion/quantization/gguf.md

@@ -22,16 +22,10 @@
 CLI (examples/offline_inference/text_to_image/text_to_image.py)
   |
   v
-Omni (vllm_omni/entrypoints/__init__.py)
+Omni (vllm_omni/entrypoints/omni.py)
   |
   v
-AsyncOmniEngine
-  |
-  v
-Orchestrator
-  |
-  v
-StageDiffusionClient
+OmniStage (diffusion)
   |
   v
 DiffusionWorker
@@ -70,13 +64,7 @@ _generate_with_async_omni
 AsyncOmni
   |
   v
-AsyncOmniEngine
-  |
-  v
-Orchestrator
-  |
-  v
-StageDiffusionClient
+DiffusionEngine
   |
   v
 OmniRequestOutput

docs/user_guide/diffusion/quantization/int8.md

@@ -1,75 +0,0 @@
-# Int8 Quantization
-
-## Overview
-
-Int8 quantization converts BF16/FP16 weights to Int8 at model load time. No calibration or pre-quantized checkpoint needed.
-
-Depending on the model, either all layers can be quantized, or some sensitive layers should stay in BF16/FP16. See the [per-model table](#supported-models) for which case applies.
-
-## Configuration
-
-1. **Python API**: set `quantization="int8"`. To skip sensitive layers, use `quantization_config` with `ignored_layers`.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-# All layers quantized
-omni = Omni(model="<your-model>", quantization="int8")
-
-# Skip sensitive layers
-omni = Omni(
-    model="<your-model>",
-    quantization_config={
-        "method": "int8",
-        "ignored_layers": ["<layer-name>"],
-    },
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-2. **CLI**: pass `--quantization int8` and optionally `--ignored-layers`.
-
-```bash
-# All layers
-python text_to_image.py --model <your-model> --quantization int8
-
-# Skip sensitive layers
-python text_to_image.py --model <your-model> --quantization int8 --ignored-layers "img_mlp"
-
-# Online serving
-vllm serve <your-model> --omni --quantization int8
-```
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `method` | str | — | Quantization method (`"int8"`) |
-| `ignored_layers` | list[str] | `[]` | Layer name patterns to keep in BF16/FP16 |
-| `activation_scheme` | str | `"dynamic"` | `"dynamic"` (no calibration) |
-
-
-The available `ignored_layers` names depend on the model architecture (e.g., `to_qkv`, `to_out`, `img_mlp`, `txt_mlp`). Consult the transformer source for your target model.
-
-## Supported Models
-
-| Model | HF Models | Recommendation | `ignored_layers` |
-|-------|-----------|---------------|------------------|
-| Z-Image | `Tongyi-MAI/Z-Image-Turbo` | All layers | None |
-| Qwen-Image | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | All layers | None |
-
-## Combining with Other Features
-
-Int8 quantization can be combined with cache acceleration:
-
-```python
-omni = Omni(
-    model="<your-model>",
-    quantization="int8",
-    cache_backend="tea_cache",
-    cache_config={"rel_l1_thresh": 0.2},
-)
-```

docs/user_guide/diffusion/quantization/overview.md

@@ -7,20 +7,12 @@ vLLM-Omni supports quantization of DiT linear layers to reduce memory usage and
 | Method | Guide |
 |--------|-------|
 | FP8 | [FP8](fp8.md) |
-| Int8 | [Int8](int8.md) |
 | GGUF | [GGUF](gguf.md) |
 
-## Device Compatibility for FP8
+## Device Compatibility
 
 | GPU Generation | Example GPUs | FP8 Mode |
 |---------------|-------------------|----------|
 | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
 
 Kernel selection is automatic.
-
-## Device Compatibility for Int8
-
-| Device Type | Generation | Example | Int8 Mode |
-|-------------|---------------|-------------------|----------|
-| NVIDIA GPU | Ada/Hopper (SM 89+) | RTX 4090, H100, H200 | Full W8A8 with native hardware |
-| Ascend NPU | Atlas A2/Atlas A3 | Atlas 800T A2/Atlas 900 A3 | Full W8A8 with native hardware |

docs/user_guide/diffusion/step_execution.md

@@ -1,61 +0,0 @@
-# Step Execution
-
-Step execution is an opt-in diffusion execution mode enabled with
-`step_execution=True` when constructing `Omni`.
-
-It is not a generic diffusion toggle for every pipeline. Only pipelines that
-implement the stepwise contract support it today.
-
-## Quick Start
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    step_execution=True,
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(
-        num_inference_steps=50,
-    ),
-)
-```
-
-## Supported Pipelines
-
-| Pipeline | Example models | Step execution |
-|----------|----------------|----------------|
-| `QwenImagePipeline` | `Qwen/Qwen-Image`, `Qwen/Qwen-Image-2512` | Yes |
-| All other diffusion pipelines | `QwenImageEditPipeline`, `QwenImageEditPlusPipeline`, `QwenImageLayeredPipeline`, GLM-Image, Wan, Flux, etc. | No |
-
-## Current Limitations
-
-- `step_execution` currently supports `batch_size=1` only.
-- `cache_backend` is not supported together with step execution.
-- Unsupported pipelines fail early during model loading.
-- Request-mode extras such as KV transfer are not wired into step mode yet.
-
-## When To Use It
-
-Use step execution only when you specifically need the pipeline to run through
-its stepwise request state machine. For normal diffusion inference, leave it
-disabled unless your workflow depends on this mode.
-
-If you are looking for general diffusion speedups, see
-[Diffusion Acceleration Overview](../diffusion_acceleration.md).
-
-## Troubleshooting
-
-If model loading fails with a message mentioning `prepare_encode()`,
-`denoise_step()`, `step_scheduler()`, and `post_decode()`, the selected
-pipeline does not support step execution.
-
-## For Model Authors
-
-If you want to add step execution support to a new diffusion pipeline, see the
-implementation guide:
-[Diffusion Step Execution Design](../../design/feature/diffusion_step_execution.md).

docs/user_guide/diffusion_acceleration.md

@@ -16,7 +16,7 @@ Both methods can provide significant speedups (typically **1.5x-2.0x**) while ma
 
 vLLM-Omni also supports quantization methods:
 
-3. **[Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8 or Int8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
+3. **[FP8 Quantization](diffusion/quantization/overview.md)** - Reduces DiT linear layers from BF16 to FP8, providing ~1.28x speedup with minimal quality loss. Supports per-layer skip for sensitive layers.
 
 vLLM-Omni also supports parallelism methods for diffusion models, including:
 
@@ -46,7 +46,6 @@ vLLM-Omni also supports parallelism methods for diffusion models, including:
 | Method | Configuration | Description | Best For |
 |--------|--------------|-------------|----------|
 | **FP8** | `quantization="fp8"` | FP8 W8A8 on Ada/Hopper, weight-only on older GPUs | Memory reduction, inference speedup |
-| **Int8** | `quantization="int8"` | Int8 W8A8 | Memory reduction, inference speedup |
 
 ## Supported Models
 
@@ -76,19 +75,17 @@ The following table shows which models are currently supported by each accelerat
 
 | Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel | HSDP | VAE-Patch-Parallel |
 |-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
-| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
-| **Wan2.1-T2V** | `Wan-AI/Wan2.1-T2V-14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
 | **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
 
 ### Quantization
 
-| Model | Model Identifier | FP8 | Int8 |
-|-------|------------------|:---:|:---:|
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ |
-| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ |
+| Model | Model Identifier | FP8 |
+|-------|------------------|:---:|
+| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ |
+| **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ |
 
 
 ## Performance Benchmarks
@@ -341,30 +338,13 @@ outputs = omni.generate(
 )
 ```
 
-### Using Int8 Quantization
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="<your-model>",
-    quantization="int8",
-)
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
 ## Documentation
 
 For detailed information on each acceleration method:
 
 - **[TeaCache Guide](diffusion/teacache.md)** - Complete TeaCache documentation, configuration options, and best practices
 - **[Cache-DiT Acceleration Guide](diffusion/cache_dit_acceleration.md)** - Comprehensive Cache-DiT guide covering DBCache, TaylorSeer, SCM, and configuration parameters
-- **[Quantization Guide](diffusion/quantization/overview.md)** - Quantization for DiT models with per-layer control
+- **[FP8 Quantization Guide](diffusion/quantization/overview.md)** - FP8 quantization for DiT models with per-layer control
 - **[Tensor Parallelism](diffusion/parallelism_acceleration.md#tensor-parallelism)** - Guidance on how to enable TP for diffusion models.
 - **[Sequence Parallelism](diffusion/parallelism_acceleration.md#sequence-parallelism)** - Guidance on how to set sequence parallelism with configuration.
 - **[CFG-Parallel](diffusion/parallelism_acceleration.md#cfg-parallel)** - Guidance on how to set CFG-Parallel to run positive/negative branches across ranks.

docs/user_guide/examples/offline_inference/bagel.md

@@ -130,7 +130,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :------------------------------ | :----------------------- |
 | `stage_type`                     | `llm`                           | Stage type               |
 | `devices`                        | `"0"`                           | GPU device ID            |
-| `max_num_seqs`                   | `1`                             | Maximum batch size       |
+| `max_batch_size`                 | `1`                             | Maximum batch size       |
 | `model_stage`                    | `thinker`                       | Model stage identifier   |
 | `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
 | `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
@@ -146,7 +146,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :---------- | :-------------------------- |
 | `stage_type`                     | `diffusion` | Stage type                  |
 | `devices`                        | `"0"`       | GPU device ID               |
-| `max_num_seqs`                   | `1`         | Maximum batch size          |
+| `max_batch_size`                 | `1`         | Maximum batch size          |
 | `model_stage`                    | `dit`       | Model stage identifier      |
 | `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
 | `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -107,7 +107,7 @@ completes. This demonstrates that audio data is available progressively rather t
 
 ## Batched Decoding
 
-The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
 
 ```
 python end2end.py --query-type CustomVoice \
@@ -116,7 +116,7 @@ python end2end.py --query-type CustomVoice \
     --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
 ```
 
-**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
 
 ## Notes
 

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -19,7 +19,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output.images
+    images = outputs[0].request_output[0].images
     images[0].save("coffee.png")
 ```
 
@@ -37,7 +37,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output.images[0].save(f"{i}.jpg")
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -49,7 +49,7 @@ if __name__ == "__main__":
 
 !!! info
 
-    For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
     batched inputs, you can [modify this configuration](https://github.com/vllm-project/vllm-omni/tree/main/configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 
@@ -72,7 +72,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output.images[0].save(f"{i}.jpg")
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -3,10 +3,7 @@
 Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_video>.
 
 
-The Wan T2V pipeline generates short videos from text prompts. Supported models include:
-
-- `Wan-AI/Wan2.2-T2V-A14B-Diffusers` (14B, requires ~48GB VRAM)
-- `Wan-AI/Wan2.1-T2V-1.3B-Diffusers` (1.3B, requires ~16GB VRAM, use `--flow-shift 3.0 --boundary-ratio 0.0`)
+The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts.
 
 ## Local CLI Usage
 

examples/offline_inference/bagel/README.md

@@ -127,7 +127,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :------------------------------ | :----------------------- |
 | `stage_type`                     | `llm`                           | Stage type               |
 | `devices`                        | `"0"`                           | GPU device ID            |
-| `max_num_seqs`                   | `1`                             | Maximum batch size       |
+| `max_batch_size`                 | `1`                             | Maximum batch size       |
 | `model_stage`                    | `thinker`                       | Model stage identifier   |
 | `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
 | `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
@@ -143,7 +143,7 @@ The default yaml configuration deploys Thinker and DiT on the same GPU. You can
 | :------------------------------- | :---------- | :-------------------------- |
 | `stage_type`                     | `diffusion` | Stage type                  |
 | `devices`                        | `"0"`       | GPU device ID               |
-| `max_num_seqs`                   | `1`         | Maximum batch size          |
+| `max_batch_size`                 | `1`         | Maximum batch size          |
 | `model_stage`                    | `dit`       | Model stage identifier      |
 | `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
 | `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |

examples/offline_inference/bagel/end2end.py

@@ -34,13 +34,6 @@ def parse_args():
         help="Path to input image for img2img.",
     )
 
-    parser.add_argument(
-        "--output",
-        type=str,
-        default=".",
-        help="Output directory to save images.",
-    )
-
     # OmniLLM init args
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-sleep-seconds", type=int, default=20)
@@ -65,11 +58,6 @@ def parse_args():
         help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
     parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
 
     args = parser.parse_args()
     return args
@@ -77,7 +65,6 @@ def parse_args():
 
 def main():
     args = parse_args()
-    os.makedirs(args.output, exist_ok=True)
     model_name = args.model
     prompts: list[OmniPromptType] = []
     try:
@@ -116,7 +103,6 @@ def main():
             "shm_threshold_bytes": args.shm_threshold_bytes,
             "worker_backend": args.worker_backend,
             "ray_address": args.ray_address,
-            "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         }
     )
 
@@ -177,17 +163,25 @@ def main():
 
     omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
-    img_idx = 0
-    for req_output in omni_outputs:
+    for i, req_output in enumerate(omni_outputs):
         images = getattr(req_output, "images", None)
+        if not images and hasattr(req_output, "output"):
+            if isinstance(req_output.output, list):
+                images = req_output.output
+            else:
+                images = [req_output.output]
 
-        if not images:
-            continue
+        if images:
+            for j, img in enumerate(images):
+                img.save(f"output_{i}_{j}.png")
 
-        for j, img in enumerate(images):
-            save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
-            img.save(save_path)
-        img_idx += 1
+        if hasattr(req_output, "request_output") and req_output.request_output:
+            for stage_out in req_output.request_output:
+                if hasattr(stage_out, "images") and stage_out.images:
+                    for k, img in enumerate(stage_out.images):
+                        save_path = f"output_{i}_stage_{getattr(stage_out, 'stage_id', '?')}_{k}.png"
+                        img.save(save_path)
+                        print(f"[Info] Saved stage output image to {save_path}")
 
     print(omni_outputs)
 

examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py

@@ -167,26 +167,27 @@ def run_e2e():
     print(f"Received {len(outputs)} outputs.")
     for i, output in enumerate(outputs):
         try:
-            ro = output.request_output
-            if ro is None:
+            ro_list = output.request_output or []
+            if not ro_list:
                 print("No request_output found.")
                 continue
 
-            # Multimodal output may be attached to RequestOutput or CompletionOutput.
-            mm = getattr(ro, "multimodal_output", None)
-            if not mm and ro.outputs:
-                mm = getattr(ro.outputs[0], "multimodal_output", None)
+            for ro in ro_list:
+                # Multimodal output may be attached to RequestOutput or CompletionOutput.
+                mm = getattr(ro, "multimodal_output", None)
+                if not mm and ro.outputs:
+                    mm = getattr(ro.outputs[0], "multimodal_output", None)
 
-            if mm:
-                print(f"Multimodal output keys: {mm.keys()}")
-                if "audio" in mm:
-                    audio_out = mm["audio"]
-                    print(f"Generated Audio Shape: {audio_out.shape}")
-                    out_path = f"output_{i}.wav"
-                    sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
-                    print(f"Saved audio to {out_path}")
-            else:
-                print("No multimodal output found.")
+                if mm:
+                    print(f"Multimodal output keys: {mm.keys()}")
+                    if "audio" in mm:
+                        audio_out = mm["audio"]
+                        print(f"Generated Audio Shape: {audio_out.shape}")
+                        out_path = f"output_{i}.wav"
+                        sf.write(out_path, audio_out.cpu().numpy().squeeze(), 22050)
+                        print(f"Saved audio to {out_path}")
+                else:
+                    print("No multimodal output found.")
         except Exception as e:
             print(f"Error inspecting output: {e}")
     omni.close()

examples/offline_inference/custom_pipeline/image_to_image/image_edit.py

@@ -209,7 +209,7 @@ async def main():
     if not outputs:
         raise ValueError("No output produced from omni.generate()")
 
-    first_out = outputs[0].request_output
+    first_out = outputs[0].request_output[0]
     req_out: OmniRequestOutput = first_out
 
     # Verify trajectory data (from custom pipeline)

examples/offline_inference/fish_speech/end2end.py

@@ -129,14 +129,8 @@ def main(args):
 
     t_start = time.perf_counter()
     for stage_outputs in omni.generate(inputs):
-        request_output = stage_outputs.request_output
-        if request_output is None or not request_output.outputs:
-            continue
-        _save_wav(
-            output_dir,
-            request_output.request_id,
-            request_output.outputs[0].multimodal_output,
-        )
+        for output in stage_outputs.request_output:
+            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
     t_end = time.perf_counter()
     logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
 

examples/offline_inference/glm_image/end2end.py

@@ -247,7 +247,6 @@ def main(args: argparse.Namespace) -> None:
         stage_configs_path=config_path,
         log_stats=args.enable_stats,
         stage_init_timeout=args.stage_init_timeout,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     init_time = time.time() - start_time
@@ -311,35 +310,36 @@ def main(args: argparse.Namespace) -> None:
 
     output_count = 0
     for stage_outputs in omni.generate(prompts, sampling_params_list, py_generator=True):
-        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "image":
-            request_id = output.request_id
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
 
-            # Get generated images
-            images = output.images if hasattr(output, "images") else []
-            if not images and hasattr(output, "multimodal_output"):
-                images = output.multimodal_output.get("images", [])
+                # Get generated images
+                images = output.images if hasattr(output, "images") else []
+                if not images and hasattr(output, "multimodal_output"):
+                    images = output.multimodal_output.get("images", [])
 
-            # Save each generated image
-            for idx, img in enumerate(images):
-                if args.num_prompts == 1 and len(images) == 1:
-                    output_path = args.output
-                else:
-                    base, ext = os.path.splitext(args.output)
-                    output_path = f"{base}_{request_id}_{idx}{ext}"
+                # Save each generated image
+                for idx, img in enumerate(images):
+                    if args.num_prompts == 1 and len(images) == 1:
+                        output_path = args.output
+                    else:
+                        base, ext = os.path.splitext(args.output)
+                        output_path = f"{base}_{request_id}_{idx}{ext}"
 
-                if isinstance(img, Image.Image):
-                    save_image(img, output_path)
-                else:
-                    print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
+                    if isinstance(img, Image.Image):
+                        save_image(img, output_path)
+                    else:
+                        print(f"Warning: Unexpected image type for request {request_id}: {type(img)}")
 
-                output_count += 1
+                    output_count += 1
 
         elif stage_outputs.final_output_type == "text":
             # AR stage output (intermediate, for debugging)
             if args.verbose:
-                print(f"AR output for request {output.request_id}:")
-                print(f"  Token count: {len(output.outputs[0].token_ids)}")
+                for output in stage_outputs.request_output:
+                    print(f"AR output for request {output.request_id}:")
+                    print(f"  Token count: {len(output.outputs[0].token_ids)}")
 
     gen_time = time.time() - gen_start_time
     print(f"\nGeneration completed in {gen_time:.2f}s")
@@ -464,11 +464,6 @@ def parse_args() -> argparse.Namespace:
         default=False,
         help="Enable verbose output",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
 
     return parser.parse_args()
 

examples/offline_inference/helios/end2end.py

@@ -173,11 +173,6 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of initial denoising steps using zero prediction (default: 1).",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
 
     # Memory & parallelism
     parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
@@ -212,7 +207,6 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     # Validate I2V / V2V arguments
@@ -299,11 +293,12 @@ def main():
                 )
 
             if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
-                inner_output = first_item.request_output
-                if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
-                    frames = inner_output.images[0] if inner_output.images else None
-                    if frames is None:
-                        raise ValueError("No video frames found in output.")
+                if isinstance(first_item.request_output, list) and len(first_item.request_output) > 0:
+                    inner_output = first_item.request_output[0]
+                    if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
+                        frames = inner_output.images[0] if inner_output.images else None
+                        if frames is None:
+                            raise ValueError("No video frames found in output.")
             elif hasattr(first_item, "images") and first_item.images:
                 frames = first_item.images
             else:

examples/offline_inference/hunyuan_image3/image_to_text.py

@@ -9,7 +9,7 @@ from PIL import Image
 from vllm_omni.entrypoints.omni import Omni
 
 """
-The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. It utilizes two tokenizer delimiter templates:
+The tencent/HunyuanImage-3.0-Instruct base model is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
 
 1) Pretrained template (default for gen_text mode), which concatenates system, image
    tokens, and user question WITHOUT role delimiters:
@@ -41,11 +41,6 @@ def parse_args() -> argparse.Namespace:
         required=True,
         help="Pretrain template prompt: <|startoftext|>{system}<img>{question}",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 
@@ -57,10 +52,7 @@ def load_image(image_path: str) -> Image.Image:
 
 
 def main(args: argparse.Namespace) -> None:
-    omni = Omni(
-        model=args.model,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
-    )
+    omni = Omni(model=args.model)
 
     prompt_dict = {
         "prompt": args.prompt,
@@ -78,8 +70,8 @@ def main(args: argparse.Namespace) -> None:
     prompts = [prompt_dict]
     omni_outputs = omni.generate(prompts=prompts)
 
-    prompt_text = omni_outputs[0].request_output.prompt
-    generated_text = omni_outputs[0].request_output.outputs[0].text
+    prompt_text = omni_outputs[0].request_output[0].prompt
+    generated_text = omni_outputs[0].request_output[0].outputs[0].text
     print(f"Prompt: {prompt_text}")
     print(f"Text: {generated_text}")
 

examples/offline_inference/image_to_image/image_edit.py

@@ -325,11 +325,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable layerwise (blockwise) offloading on DiT modules.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 
@@ -394,7 +389,6 @@ def main():
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         enable_cpu_offload=args.enable_cpu_offload,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     print("Pipeline loaded")
 
@@ -469,12 +463,12 @@ def main():
         raise ValueError("No output generated from omni.generate()")
 
     # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from request_output.images
+    # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output
+    req_out = first_output.request_output[0]
     if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
         raise ValueError("Invalid request_output structure or missing 'images' key")
 

examples/offline_inference/image_to_video/image_to_video.py

@@ -2,16 +2,15 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 """
-Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
+Image-to-Video generation example using Wan2.2 I2V/TI2V models or LTX2.
 
 Supports:
 - Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
 - Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
 - LTX2 image-to-video pipeline
-- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
 
 Usage:
-    # Wan I2V-A14B (MoE)
+    # I2V-A14B (MoE)
     python image_to_video.py --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
         --image input.jpg --prompt "A cat playing with yarn"
 
@@ -25,11 +24,6 @@ Usage:
         --image input.jpg --prompt "A cinematic dolly shot of a boat" \
         --num-frames 121 --num-inference-steps 40 --guidance-scale 4.0 \
         --frame-rate 24 --fps 24 --output ltx2_i2v.mp4
-
-    # HunyuanVideo-1.5 I2V (480p)
-    python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
-        --image input.jpg --prompt "A cat playing with yarn" \
-        --flow-shift 5.0 --guidance-scale 6.0
 """
 
 import argparse
@@ -49,11 +43,11 @@ from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
+    parser = argparse.ArgumentParser(description="Generate a video from an image with Wan2.2 or LTX2.")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-        help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
+        help="Diffusers Wan2.2 I2V model ID or local path.",
     )
     parser.add_argument(
         "--model-class-name",
@@ -182,11 +176,6 @@ def parse_args() -> argparse.Namespace:
             "Default 1 means pure sharding (no replication). "
         ),
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 
@@ -292,7 +281,6 @@ def main():
         model_class_name=model_class_name,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     if profiler_enabled:
@@ -350,6 +338,8 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
+            if isinstance(inner_output, list):
+                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_image_summarize.py

@@ -43,11 +43,6 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument("--top-p", type=float, default=0.9)
     parser.add_argument("--seed", type=int, default=42)
     parser.add_argument("--trust-remote-code", action="store_true")
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 
@@ -75,7 +70,6 @@ def main() -> None:
         model=args.model,
         stage_configs_path=args.stage_config,
         trust_remote_code=args.trust_remote_code,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     try:
         sp = SamplingParams(
@@ -105,12 +99,14 @@ def main() -> None:
 
     lines: list[str] = []
     for stage_outputs in outputs:
-        ro = getattr(stage_outputs, "request_output", stage_outputs)
-        text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
-        lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
-        lines.append("answer:\n")
-        lines.append(text.strip() + "\n")
-        lines.append("\n")
+        req_outputs = getattr(stage_outputs, "request_output", stage_outputs)
+        req_outputs = req_outputs if isinstance(req_outputs, list) else [req_outputs]
+        for ro in req_outputs:
+            text = ro.outputs[0].text if getattr(ro, "outputs", None) else str(ro)
+            lines.append(f"request_id: {getattr(ro, 'request_id', 'unknown')}\n")
+            lines.append("answer:\n")
+            lines.append(text.strip() + "\n")
+            lines.append("\n")
 
     print("\n".join(lines))
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -148,16 +148,19 @@ def _collect_images(outputs: list) -> list[torch.Tensor]:
     """Extract all image tensors produced by the final (DiT) stage."""
     images: list[torch.Tensor] = []
     for out in outputs:
-        ro_item = getattr(out, "request_output", out)
-        for completion in getattr(ro_item, "outputs", None) or []:
-            mm = getattr(completion, "multimodal_output", None)
-            if not isinstance(mm, dict) or "image" not in mm:
-                raise RuntimeError(f"Missing image in multimodal output: {mm}")
-            payload = mm["image"]
-            for tensor in payload if isinstance(payload, list) else [payload]:
-                if not isinstance(tensor, torch.Tensor):
-                    raise TypeError(f"Expected image tensor, got {type(tensor)}")
-                images.append(tensor)
+        ro_list = getattr(out, "request_output", out)
+        if not isinstance(ro_list, list):
+            ro_list = [ro_list]
+        for ro_item in ro_list:
+            for completion in getattr(ro_item, "outputs", None) or []:
+                mm = getattr(completion, "multimodal_output", None)
+                if not isinstance(mm, dict) or "image" not in mm:
+                    raise RuntimeError(f"Missing image in multimodal output: {mm}")
+                payload = mm["image"]
+                for tensor in payload if isinstance(payload, list) else [payload]:
+                    if not isinstance(tensor, torch.Tensor):
+                        raise TypeError(f"Expected image tensor, got {type(tensor)}")
+                    images.append(tensor)
     return images
 
 

examples/offline_inference/mimo_audio/end2end.py

@@ -180,16 +180,15 @@ def main(args):
     # Get the query function and call it with appropriate parameters
     query_func = query_map[args.query_type]
 
-    omni = Omni(
+    omni_llm = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.enable_stats,
-        log_file=("omni_pipeline.log" if args.enable_stats else None),
+        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
         init_sleep_seconds=args.init_sleep_seconds,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
         shm_threshold_bytes=args.shm_threshold_bytes,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -291,7 +290,7 @@ def main(args):
     prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
 
     print("prompts", prompts)
-    omni_outputs = omni.generate(prompts, sampling_params_list)
+    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
 
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     if args.query_type is not None:
@@ -299,44 +298,45 @@ def main(args):
     os.makedirs(output_dir, exist_ok=True)
 
     for stage_outputs in omni_outputs:
-        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            request_id = output.request_id
-            text_output = output.outputs[0].text
-            # Save aligned text file per request
-            prompt_text = output.prompt
-            out_txt = os.path.join(output_dir, f"{request_id}.txt")
-            lines = []
-            lines.append("Prompt:\n")
-            lines.append(str(prompt_text) + "\n")
-            lines.append("vllm_text_output:\n")
-            lines.append(str(text_output).strip() + "\n")
-            try:
-                with open(out_txt, "w", encoding="utf-8") as f:
-                    print("lines", lines)
-                    f.writelines(lines)
-            except Exception as e:
-                print(f"[Warn] Failed writing text file {out_txt}: {e}")
-            print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = output.prompt
+                out_txt = os.path.join(output_dir, f"{request_id}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        print("lines", lines)
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}\n")
         elif stage_outputs.final_output_type == "audio":
-            request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output.get("audio")
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                audio_tensor = output.outputs[0].multimodal_output.get("audio")
 
-            if audio_tensor is None:
-                continue
+                if audio_tensor is None:
+                    continue
 
-            output_wav = os.path.join(output_dir, f"{request_id}.wav")
+                output_wav = os.path.join(output_dir, f"{request_id}.wav")
 
-            # Convert to numpy array and ensure correct format
-            audio_numpy = audio_tensor.float().detach().cpu().numpy()
+                # Convert to numpy array and ensure correct format
+                audio_numpy = audio_tensor.float().detach().cpu().numpy()
 
-            # Ensure audio is 1D (flatten if needed)
-            if audio_numpy.ndim > 1:
-                audio_numpy = audio_numpy.flatten()
+                # Ensure audio is 1D (flatten if needed)
+                if audio_numpy.ndim > 1:
+                    audio_numpy = audio_numpy.flatten()
 
-            # Save audio file with explicit WAV format
-            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
-            print(f"Request ID: {request_id}, Audio saved to {output_wav}")
+                # Save audio file with explicit WAV format
+                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+                print(f"Request ID: {request_id}, Audio saved to {output_wav}")
 
 
 def parse_args():
@@ -434,11 +434,6 @@ def parse_args():
         default="../../../model_executor/stage_configs/mimo_audio.yaml",
         help="Path to a stage configs file.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -320,14 +320,13 @@ def main(args):
         query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
     else:
         query_result = query_func()
-    omni = Omni(
+    omni_llm = Omni(
         model=model_name,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
         batch_timeout=args.batch_timeout,
         init_timeout=args.init_timeout,
         shm_threshold_bytes=args.shm_threshold_bytes,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
     thinker_sampling_params = SamplingParams(
         temperature=0.0,  # Deterministic - no randomness
@@ -379,11 +378,9 @@ def main(args):
             prompt["modalities"] = output_modalities
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-    if profiler_enabled and hasattr(omni, "start_profile"):
-        omni.start_profile(stages=[0])
-    elif profiler_enabled:
-        print("[Warn] VLLM_TORCH_PROFILER_DIR is set, but current engine does not support profiler controls.")
-    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+    if profiler_enabled:
+        omni_llm.start_profile(stages=[0])
+    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
@@ -392,42 +389,43 @@ def main(args):
     total_requests = len(prompts)
     processed_count = 0
     for stage_outputs in omni_generator:
-        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            request_id = output.request_id
-            text_output = output.outputs[0].text
-            # Save aligned text file per request
-            prompt_text = output.prompt
-            out_txt = os.path.join(output_dir, f"{request_id}.txt")
-            lines = []
-            lines.append("Prompt:\n")
-            lines.append(str(prompt_text) + "\n")
-            lines.append("vllm_text_output:\n")
-            lines.append(str(text_output).strip() + "\n")
-            try:
-                with open(out_txt, "w", encoding="utf-8") as f:
-                    f.writelines(lines)
-            except Exception as e:
-                print(f"[Warn] Failed writing text file {out_txt}: {e}")
-            print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = output.prompt
+                out_txt = os.path.join(output_dir, f"{request_id}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output["audio"]
-            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
-            sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
-            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                audio_tensor = output.outputs[0].multimodal_output["audio"]
+                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+                sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += 1
-        if profiler_enabled and hasattr(omni, "stop_profile") and processed_count >= total_requests:
+        processed_count += len(stage_outputs.request_output)
+        if profiler_enabled and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni.stop_profile()
+            omni_llm.stop_profile()
 
             print("[Info] Waiting 30s for workers to write massive trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait finished.")
 
-    omni.close()
+    omni_llm.close()
 
 
 def parse_args():
@@ -541,11 +539,6 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_omni/end2end.py

@@ -294,12 +294,11 @@ def main(args):
     else:
         query_result = query_func()
 
-    omni = Omni(
+    omni_llm = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     thinker_sampling_params = SamplingParams(
@@ -355,8 +354,8 @@ def main(args):
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     if profiler_enabled:
-        omni.start_profile(stages=[0])
-    omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+        omni_llm.start_profile(stages=[0])
+    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
     # Determine output directory: prefer --output-dir; fallback to --output-wav
     output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
     os.makedirs(output_dir, exist_ok=True)
@@ -367,50 +366,51 @@ def main(args):
     print(f"query type: {args.query_type}")
 
     for stage_outputs in omni_generator:
-        output = stage_outputs.request_output
         if stage_outputs.final_output_type == "text":
-            request_id = output.request_id
-            text_output = output.outputs[0].text
-            # Save aligned text file per request
-            prompt_text = output.prompt
-            out_txt = os.path.join(output_dir, f"{request_id}.txt")
-            lines = []
-            lines.append("Prompt:\n")
-            lines.append(str(prompt_text) + "\n")
-            lines.append("vllm_text_output:\n")
-            lines.append(str(text_output).strip() + "\n")
-            try:
-                with open(out_txt, "w", encoding="utf-8") as f:
-                    f.writelines(lines)
-            except Exception as e:
-                print(f"[Warn] Failed writing text file {out_txt}: {e}")
-            print(f"Request ID: {request_id}, Text saved to {out_txt}")
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                text_output = output.outputs[0].text
+                # Save aligned text file per request
+                prompt_text = output.prompt
+                out_txt = os.path.join(output_dir, f"{request_id}.txt")
+                lines = []
+                lines.append("Prompt:\n")
+                lines.append(str(prompt_text) + "\n")
+                lines.append("vllm_text_output:\n")
+                lines.append(str(text_output).strip() + "\n")
+                try:
+                    with open(out_txt, "w", encoding="utf-8") as f:
+                        f.writelines(lines)
+                except Exception as e:
+                    print(f"[Warn] Failed writing text file {out_txt}: {e}")
+                print(f"Request ID: {request_id}, Text saved to {out_txt}")
         elif stage_outputs.final_output_type == "audio":
-            request_id = output.request_id
-            audio_tensor = output.outputs[0].multimodal_output["audio"]
-            output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+            for output in stage_outputs.request_output:
+                request_id = output.request_id
+                audio_tensor = output.outputs[0].multimodal_output["audio"]
+                output_wav = os.path.join(output_dir, f"output_{request_id}.wav")
 
-            # Convert to numpy array and ensure correct format
-            audio_numpy = audio_tensor.float().detach().cpu().numpy()
+                # Convert to numpy array and ensure correct format
+                audio_numpy = audio_tensor.float().detach().cpu().numpy()
 
-            # Ensure audio is 1D (flatten if needed)
-            if audio_numpy.ndim > 1:
-                audio_numpy = audio_numpy.flatten()
+                # Ensure audio is 1D (flatten if needed)
+                if audio_numpy.ndim > 1:
+                    audio_numpy = audio_numpy.flatten()
 
-            # Save audio file with explicit WAV format
-            sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
-            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+                # Save audio file with explicit WAV format
+                sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV")
+                print(f"Request ID: {request_id}, Saved audio to {output_wav}")
 
-        processed_count += 1
+        processed_count += len(stage_outputs.request_output)
         if profiler_enabled and processed_count >= total_requests:
             print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...")
             # Stop the profiler while workers are still alive
-            omni.stop_profile()
+            omni_llm.stop_profile()
 
             print("[Info] Waiting 30s for workers to write trace files to disk...")
             time.sleep(30)
             print("[Info] Trace export wait time finished.")
-    omni.close()
+    omni_llm.close()
 
 
 def parse_args():
@@ -527,11 +527,6 @@ def parse_args():
         default=False,
         help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
 
     return parser.parse_args()
 

examples/offline_inference/qwen3_omni/end2end_async_chunk.py

@@ -231,47 +231,52 @@ async def run_single_request(
             sampling_params_list=sampling_params_list,
             output_modalities=output_modalities,
         ):
-            output = omni_output.request_output
-            if omni_output.final_output_type == "text":
-                if stage_0_first_output_ts is None:
-                    stage_0_first_output_ts = time.perf_counter()
-                text_output = output.outputs[0].text
-                if output.finished:
-                    text_parts.append(text_output)
-            elif omni_output.final_output_type == "audio":
-                mm_out = output.outputs[0].multimodal_output
-                if mm_out and "audio" in mm_out:
-                    if first_audio_ts is None:
-                        first_audio_ts = time.perf_counter()
-                    if audio_sr is None and "sr" in mm_out:
-                        sr_val = mm_out["sr"]
-                        audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
-                        samplerate = audio_sr
-                    audio_data = mm_out["audio"]
-                    if isinstance(audio_data, list):
-                        new_chunks = audio_data[audio_list_consumed:]
-                        audio_list_consumed = len(audio_data)
-                    elif isinstance(audio_data, torch.Tensor):
-                        new_chunks = [audio_data]
-                        audio_last_tensor = audio_data
-                    else:
-                        new_chunks = []
+            if not isinstance(omni_output.request_output, list):
+                outputs_list = [omni_output.request_output]
+            else:
+                outputs_list = omni_output.request_output
 
-                    if stream_audio_to_disk and new_chunks:
-                        if sf_writer is None:
-                            sf_writer = sf.SoundFile(
-                                wav_file,
-                                mode="w",
-                                samplerate=samplerate,
-                                channels=1,
-                                subtype="FLOAT",
-                            )
-                        for chunk in new_chunks:
-                            chunk_np = chunk.float().detach().cpu().numpy().flatten()
-                            sf_writer.write(chunk_np)
-                            audio_samples_written += len(chunk_np)
-                    else:
-                        audio_chunks.extend(new_chunks)
+            for output in outputs_list:
+                if omni_output.final_output_type == "text":
+                    if stage_0_first_output_ts is None:
+                        stage_0_first_output_ts = time.perf_counter()
+                    text_output = output.outputs[0].text
+                    if output.finished:
+                        text_parts.append(text_output)
+                elif omni_output.final_output_type == "audio":
+                    mm_out = output.outputs[0].multimodal_output
+                    if mm_out and "audio" in mm_out:
+                        if first_audio_ts is None:
+                            first_audio_ts = time.perf_counter()
+                        if audio_sr is None and "sr" in mm_out:
+                            sr_val = mm_out["sr"]
+                            audio_sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+                            samplerate = audio_sr
+                        audio_data = mm_out["audio"]
+                        if isinstance(audio_data, list):
+                            new_chunks = audio_data[audio_list_consumed:]
+                            audio_list_consumed = len(audio_data)
+                        elif isinstance(audio_data, torch.Tensor):
+                            new_chunks = [audio_data]
+                            audio_last_tensor = audio_data
+                        else:
+                            new_chunks = []
+
+                        if stream_audio_to_disk and new_chunks:
+                            if sf_writer is None:
+                                sf_writer = sf.SoundFile(
+                                    wav_file,
+                                    mode="w",
+                                    samplerate=samplerate,
+                                    channels=1,
+                                    subtype="FLOAT",
+                                )
+                            for chunk in new_chunks:
+                                chunk_np = chunk.float().detach().cpu().numpy().flatten()
+                                sf_writer.write(chunk_np)
+                                audio_samples_written += len(chunk_np)
+                        else:
+                            audio_chunks.extend(new_chunks)
     finally:
         if sf_writer is not None:
             sf_writer.close()
@@ -382,7 +387,6 @@ async def run_all(args):
             stage_configs_path=args.stage_configs_path,
             log_stats=args.log_stats,
             stage_init_timeout=args.stage_init_timeout,
-            enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
         )
 
         # Use default sampling params from stage config (they are pre-configured
@@ -585,11 +589,6 @@ def parse_args():
         default=16000,
         help="Sampling rate for audio loading.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 

examples/offline_inference/qwen3_tts/README.md

@@ -104,7 +104,7 @@ completes. This demonstrates that audio data is available progressively rather t
 
 ## Batched Decoding
 
-The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_num_seqs > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
+The Code2Wav stage (stage 1) supports batched decoding, where multiple requests are decoded in a single forward pass through the SpeechTokenizer. To use it, provide a stage config with `max_batch_size > 1` and pass multiple prompts via `--txt-prompts` with a matching `--batch-size`.
 
 ```
 python end2end.py --query-type CustomVoice \
@@ -113,7 +113,7 @@ python end2end.py --query-type CustomVoice \
     --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts_batch.yaml
 ```
 
-**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_num_seqs >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_num_seqs`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
+**Important:** `--batch-size` must match a CUDA graph capture size (1, 2, 4, 8, 16...) because the Talker's code predictor KV cache is sized to `max_num_seqs`, and CUDA graphs pad the batch to the next capture size. Both stages need `max_batch_size >= batch_size` in the stage config for batching to take effect. If only stage 1 has a higher `max_batch_size`, it won't help — stage 1 can only batch chunks from requests that are in-flight simultaneously, which requires stage 0 to also process multiple requests concurrently.
 
 ## Notes
 

examples/offline_inference/qwen3_tts/end2end.py

@@ -51,68 +51,16 @@ def _estimate_prompt_len(
 
             tok = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True, padding_side="left")
             cfg = Qwen3TTSConfig.from_pretrained(model_name, trust_remote_code=True)
+            _cache[model_name] = (tok, getattr(cfg, "talker_config", None))
 
-            # Load speech tokenizer (codec encoder) for exact ref_code_len.
-            speech_tok = None
-            try:
-                import os
-
-                from transformers.utils import cached_file
-
-                from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_tokenizer import Qwen3TTSTokenizer
-
-                st_cfg_path = cached_file(model_name, "speech_tokenizer/config.json")
-                if st_cfg_path:
-                    speech_tok = Qwen3TTSTokenizer.from_pretrained(
-                        os.path.dirname(st_cfg_path), torch_dtype=torch.bfloat16
-                    )
-                    logger.info("Loaded speech tokenizer for exact ref_code_len estimation")
-            except Exception as e:
-                logger.debug("Could not load speech tokenizer: %s", e)
-
-            _cache[model_name] = (tok, getattr(cfg, "talker_config", None), speech_tok)
-
-        tok, tcfg, speech_tok = _cache[model_name]
+        tok, tcfg = _cache[model_name]
         task_type = (additional_information.get("task_type") or ["CustomVoice"])[0]
-
-        def _estimate_ref_code_len(ref_audio: object) -> int | None:
-            """Encode ref_audio with the actual codec to get exact frame count."""
-            if not isinstance(ref_audio, (str, list)):
-                return None
-            audio_path = ref_audio[0] if isinstance(ref_audio, list) else ref_audio
-            if not isinstance(audio_path, str) or not audio_path.strip():
-                return None
-            try:
-                from vllm.multimodal.media import MediaConnector
-
-                connector = MediaConnector(allowed_local_media_path="/")
-                audio, sr = connector.fetch_audio(audio_path)
-                import numpy as np
-
-                wav_np = np.asarray(audio, dtype=np.float32)
-
-                if speech_tok is not None:
-                    enc = speech_tok.encode(wav_np, sr=int(sr), return_dict=True)
-                    ref_code = getattr(enc, "audio_codes", None)
-                    if isinstance(ref_code, list):
-                        ref_code = ref_code[0] if ref_code else None
-                    if ref_code is not None and hasattr(ref_code, "shape"):
-                        shape = ref_code.shape
-                        return int(shape[0]) if len(shape) == 2 else int(shape[1]) if len(shape) == 3 else None
-
-                # Fallback: estimate from duration
-                codec_hz = getattr(tcfg, "codec_frame_rate", None) or 12
-                return int(len(audio) / sr * codec_hz)
-            except Exception:
-                return None
-
         return Qwen3TTSTalkerForConditionalGeneration.estimate_prompt_len_from_additional_information(
             additional_information=additional_information,
             task_type=task_type,
             tokenize_prompt=lambda t: tok(t, padding=False)["input_ids"],
             codec_language_id=getattr(tcfg, "codec_language_id", None),
             spk_is_dialect=getattr(tcfg, "spk_is_dialect", None),
-            estimate_ref_code_len=_estimate_ref_code_len,
         )
     except Exception as exc:
         logger.warning("Failed to estimate prompt length, using fallback 2048: %s", exc)
@@ -371,15 +319,14 @@ def main(args):
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     batch_size = args.batch_size
     for batch_start in range(0, len(inputs), batch_size):
         batch = inputs[batch_start : batch_start + batch_size]
         for stage_outputs in omni.generate(batch):
-            output = stage_outputs.request_output
-            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+            for output in stage_outputs.request_output:
+                _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
 
 
 async def main_streaming(args):
@@ -393,7 +340,6 @@ async def main_streaming(args):
         stage_configs_path=args.stage_configs_path,
         log_stats=args.log_stats,
         stage_init_timeout=args.stage_init_timeout,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     for i, prompt in enumerate(inputs):
@@ -535,11 +481,6 @@ def parse_args():
         default=1,
         help="Number of prompts per batch (default: 1, sequential).",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
 
     return parser.parse_args()
 

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -90,11 +90,6 @@ def parse_args() -> argparse.Namespace:
         default=44100,
         help="Sample rate for output audio (Stable Audio uses 44100 Hz).",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 
@@ -138,10 +133,7 @@ def main():
     print(f"{'=' * 60}\n")
 
     # Initialize Omni with Stable Audio model
-    omni = Omni(
-        model=args.model,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
-    )
+    omni = Omni(model=args.model)
 
     # Calculate audio end time
     audio_end_in_s = args.audio_start + args.audio_length
@@ -185,7 +177,7 @@ def main():
     output = outputs[0]
     if not hasattr(output, "request_output") or not output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
-    request_output = output.request_output
+    request_output = output.request_output[0]
     if not hasattr(request_output, "multimodal_output"):
         raise ValueError("No multimodal_output found in request_output")
 

examples/offline_inference/text_to_image/README.md

@@ -16,7 +16,7 @@ if __name__ == "__main__":
     omni = Omni(model="Qwen/Qwen-Image")
     prompt = "a cup of coffee on the table"
     outputs = omni.generate(prompt)
-    images = outputs[0].request_output.images
+    images = outputs[0].request_output[0].images
     images[0].save("coffee.png")
 ```
 
@@ -34,7 +34,7 @@ if __name__ == "__main__":
     ]
     outputs = omni.generate(prompts)
     for i, output in enumerate(outputs):
-      image = output.request_output.images[0].save(f"{i}.jpg")
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -46,7 +46,7 @@ if __name__ == "__main__":
 
 !!! info
 
-    For diffusion pipelines, the stage config field `stage_args.[].engine_args.max_num_seqs` is 1 by default, and the input
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
     list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
     batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
 
@@ -69,7 +69,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output.images[0].save(f"{i}.jpg")
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

examples/offline_inference/text_to_image/gradio_demo.py

@@ -112,7 +112,7 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
         )
         images_outputs = []
         for output in outputs:
-            req_out = output.request_output
+            req_out = output.request_output[0]
             if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
                 raise ValueError("Invalid request_output structure or missing 'images' key")
             images = req_out.images

examples/offline_inference/text_to_image/text_to_image.py

@@ -14,6 +14,7 @@ from vllm_omni.entrypoints.omni import Omni
 from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.lora.request import LoRARequest
 from vllm_omni.lora.utils import stable_lora_int_id
+from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
 
@@ -131,10 +132,12 @@ def parse_args() -> argparse.Namespace:
         "--quantization",
         type=str,
         default=None,
-        choices=["fp8", "int8", "gguf"],
-        help="Quantization method for the transformer. "
-        "Options: 'fp8' (FP8 W8A8 on Ada/Hopper, weight-only on older GPUs), 'int8' (Int8 W8A8), 'gguf' (GGUF quantized weights). "
-        "Default: None (no quantization, uses BF16).",
+        choices=["fp8", "gguf"],
+        help=(
+            "Quantization method for the transformer. "
+            "Options: 'fp8' (FP8 W8A8), 'gguf' (GGUF quantized weights). "
+            "Default: None (no quantization, uses BF16)."
+        ),
     )
     parser.add_argument(
         "--gguf-model",
@@ -215,11 +218,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
     return parser.parse_args()
 
 
@@ -307,7 +305,6 @@ def main():
         "parallel_config": parallel_config,
         "enforce_eager": args.enforce_eager,
         "enable_cpu_offload": args.enable_cpu_offload,
-        "enable_diffusion_pipeline_profiler": args.enable_diffusion_pipeline_profiler,
         **lora_args,
         **quant_kwargs,
     }
@@ -403,18 +400,20 @@ def main():
         else:
             print("[Profiler] No valid profiling data returned.")
 
-    # omni.generate() returns list[OmniRequestOutput]
+    # Extract images from OmniRequestOutput
+    # omni.generate() returns list[OmniRequestOutput], extract images from the first output
     if not outputs or len(outputs) == 0:
         raise ValueError("No output generated from omni.generate()")
     logger.info(f"Outputs: {outputs}")
 
+    # Extract images from request_output[0]['images']
     first_output = outputs[0]
     if not hasattr(first_output, "request_output") or not first_output.request_output:
         raise ValueError("No request_output found in OmniRequestOutput")
 
-    req_out = first_output.request_output
-    if not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images'.")
+    req_out = first_output.request_output[0]
+    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images' key")
 
     images = req_out.images
     if not images:

examples/offline_inference/text_to_video/text_to_video.md

@@ -1,19 +1,10 @@
 # Text-To-Video
 
-A unified script for text-to-video generation. Supports multiple models with model-aware defaults.
-
-## Supported Models
-
-| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
-|---|---|---|---|---|---|
-| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
-| `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 |
+The `Wan-AI/Wan2.2-T2V-A14B-Diffusers` pipeline generates short videos from text prompts. This script can also be used
+for `Lightricks/LTX-2` to generate video+audio.
 
 ## Local CLI Usage
 
-### Wan2.2 (default)
-
 ```bash
 python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
@@ -45,100 +36,22 @@ python text_to_video.py \
   --output ltx2_out.mp4
 ```
 
-### HunyuanVideo-1.5 (480p)
+Key arguments:
 
-```bash
-python text_to_video.py \
-  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
-  --prompt "A cat walks through a sunlit garden, flowers swaying gently in the breeze." \
-  --height 480 \
-  --width 832 \
-  --num-frames 121 \
-  --guidance-scale 6.0 \
-  --flow-shift 5.0 \
-  --num-inference-steps 50 \
-  --fps 24 \
-  --output hunyuan_video_15_output.mp4
-```
-
-### HunyuanVideo-1.5 (720p)
-
-```bash
-python text_to_video.py \
-  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v \
-  --prompt "A serene lakeside sunrise with mist over the water." \
-  --height 720 \
-  --width 1280 \
-  --num-frames 121 \
-  --guidance-scale 6.0 \
-  --flow-shift 9.0 \
-  --num-inference-steps 50 \
-  --fps 24 \
-  --output hunyuan_720p.mp4
-```
-
-### HunyuanVideo-1.5 with FP8 Quantization
-
-```bash
-python text_to_video.py \
-  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
-  --prompt "A dog running across a field of golden wheat." \
-  --quantization fp8 \
-  --height 480 --width 832 --num-frames 121 \
-  --guidance-scale 6.0 --flow-shift 5.0 \
-  --output hunyuan_fp8.mp4
-```
-
-Quick test (smaller resolution, fewer frames):
-
-```bash
-python text_to_video.py \
-  --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \
-  --prompt "A serene lakeside sunrise with mist over the water." \
-  --height 320 --width 576 --num-frames 17 --num-inference-steps 30 \
-  --flow-shift 5.0 \
-  --output quick_test.mp4
-```
-
-## Key Arguments
-
-### Common
-
-- `--model`: Diffusers model ID or local path.
 - `--prompt`: text description (string).
-- `--height/--width`: output resolution. Default depends on model.
-- `--num-frames`: number of frames. Default depends on model.
-- `--guidance-scale`: CFG scale. Default depends on model.
-- `--num-inference-steps`: sampling steps. Default depends on model.
-- `--fps`: frames per second for the saved MP4.
+- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--num-frames`: Number of frames (Wan default is 81).
+- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high).
+- `--negative-prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
+- `--boundary-ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
+- `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
 - `--vae-use-slicing`: enable VAE slicing for memory optimization.
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--tensor-parallel-size`: tensor parallel size (effective for models that support TP, e.g. LTX2).
 - `--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).
-- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
-- `--flow-shift`: scheduler flow_shift parameter.
 
-### Wan2.2-specific
-
-- `--negative-prompt`: artifacts to suppress.
-- `--guidance-scale-high`: separate CFG scale for high-noise stage.
-- `--boundary-ratio`: boundary split for low/high DiT (default 0.875).
-- `--flow-shift`: scheduler flow_shift (5.0 for 720p, 12.0 for 480p).
-- `--cache-backend`: `cache_dit` for acceleration.
-
-### HunyuanVideo-1.5 Optimal Configs
-
-| Variant | flow_shift | guidance_scale | steps |
-|---------|-----------|----------------|-------|
-| 480p T2V | 5.0 | 6.0 | 50 |
-| 720p T2V | 9.0 | 6.0 | 50 |
-| 480p I2V | 5.0 | 6.0 | 50 |
-| 720p I2V | 7.0 | 6.0 | 50 |
-| CFG-distilled | (same) | 1.0 | 50 |
-
-> If you encounter OOM errors, try `--vae-use-slicing`, `--vae-use-tiling`, `--enable-cpu-offload`, or `--quantization fp8`.
+> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/text_to_video/text_to_video.py

@@ -15,92 +15,56 @@ from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
 from vllm_omni.platforms import current_omni_platform
 
-_MODEL_PRESETS = {
-    "wan": {
-        "height": 720,
-        "width": 1280,
-        "num_frames": 81,
-        "num_inference_steps": 40,
-        "guidance_scale": 4.0,
-        "fps": 24,
-        "output": "wan22_output.mp4",
-    },
-    "hunyuan": {
-        "height": 480,
-        "width": 832,
-        "num_frames": 121,
-        "num_inference_steps": 50,
-        "guidance_scale": 6.0,
-        "fps": 24,
-        "output": "hunyuan_video_15_output.mp4",
-    },
-}
-
-
-def _detect_preset(model: str) -> dict:
-    model_lower = model.lower()
-    if "hunyuan" in model_lower:
-        return _MODEL_PRESETS["hunyuan"]
-    return _MODEL_PRESETS["wan"]
-
 
 def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(
-        description="Generate a video from a text prompt. "
-        "Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models."
-    )
+    parser = argparse.ArgumentParser(description="Generate a video with Wan2.2 T2V.")
     parser.add_argument(
         "--model",
         default="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
-        help="Diffusers model ID or local path. "
-        "Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
-        "hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
+        help="Diffusers Wan2.2 model ID or local path.",
     )
     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 (Wan2.2 only).")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
     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(
-        "--guidance-scale-high", type=float, default=None, help="Separate CFG for high-noise stage (Wan2.2 only)."
-    )
-    parser.add_argument("--height", type=int, default=None, help="Video height. Default: model-specific.")
-    parser.add_argument("--width", type=int, default=None, help="Video width. Default: model-specific.")
-    parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
+    parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (applied to low/high).")
+    parser.add_argument("--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise.")
+    parser.add_argument("--height", type=int, default=720, help="Video height.")
+    parser.add_argument("--width", type=int, default=1280, help="Video width.")
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames (Wan default is 81).")
     parser.add_argument(
         "--frame-rate",
         type=float,
         default=None,
         help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
     )
-    parser.add_argument(
-        "--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
-    )
+    parser.add_argument("--num-inference-steps", type=int, default=40, help="Sampling steps.")
     parser.add_argument(
         "--boundary-ratio",
         type=float,
-        default=None,
-        help="(Wan2.2) Boundary split ratio for low/high DiT. Default 0.875.",
+        default=0.875,
+        help="Boundary split ratio for low/high DiT. Default 0.875 uses both transformers for best quality. Set to 1.0 to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited).",
     )
     parser.add_argument(
-        "--flow-shift",
-        type=float,
-        default=None,
-        help="Scheduler flow_shift. Wan2.2: 5.0(720p)/12.0(480p). HunyuanVideo-1.5: 5.0(480p)/9.0(720p).",
+        "--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
     )
     parser.add_argument(
         "--cache-backend",
         type=str,
         default=None,
         choices=["cache_dit"],
-        help="Cache backend for acceleration (Wan2.2). Default: None.",
+        help=(
+            "Cache backend to use for acceleration. "
+            "Options: 'cache_dit' (DBCache + SCM + TaylorSeer). "
+            "Default: None (no cache acceleration)."
+        ),
     )
     parser.add_argument(
         "--enable-cache-dit-summary",
         action="store_true",
         help="Enable cache-dit summary logging after diffusion forward passes.",
     )
-    parser.add_argument("--output", type=str, default=None, help="Output path (mp4). Default: model-specific.")
-    parser.add_argument("--fps", type=int, default=None, help="Frames per second for the output video.")
+    parser.add_argument("--output", type=str, default="wan22_output.mp4", help="Path to save the video (mp4).")
+    parser.add_argument("--fps", type=int, default=24, help="Frames per second for the output video.")
     parser.add_argument(
         "--vae-use-slicing",
         action="store_true",
@@ -168,46 +132,32 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable expert parallelism for MoE layers.",
     )
-    parser.add_argument(
-        "--enable-diffusion-pipeline-profiler",
-        action="store_true",
-        help="Enable diffusion pipeline profiler to display stage durations.",
-    )
-    parser.add_argument(
-        "--quantization",
-        type=str,
-        default=None,
-        choices=["fp8", "gguf"],
-        help="Quantization method for the transformer (fp8 for online FP8 quantization).",
-    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
-
-    preset = _detect_preset(args.model)
-    for key, default_val in preset.items():
-        if getattr(args, key.replace("-", "_"), None) is None:
-            setattr(args, key.replace("-", "_"), default_val)
-
     generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-    # Cache-dit config (Wan2.2 only)
+    frame_rate = args.frame_rate if args.frame_rate is not None else float(args.fps)
+
+    # Wan2.2 cache-dit tuning (from cache-dit examples and cache_alignment).
     cache_config = None
     if args.cache_backend == "cache_dit":
         cache_config = {
-            "Fn_compute_blocks": 1,
-            "Bn_compute_blocks": 0,
-            "max_warmup_steps": 4,
+            # DBCache parameters [cache-dit only]
+            "Fn_compute_blocks": 1,  # Optimized for single-transformer models
+            "Bn_compute_blocks": 0,  # Number of backward compute blocks
+            "max_warmup_steps": 4,  # Maximum warmup steps (works for few-step models)
             "max_cached_steps": 20,
-            "residual_diff_threshold": 0.24,
-            "max_continuous_cached_steps": 3,
-            "enable_taylorseer": False,
-            "taylorseer_order": 1,
-            "scm_steps_mask_policy": None,
-            "scm_steps_policy": "dynamic",
+            "residual_diff_threshold": 0.24,  # Higher threshold for more aggressive caching
+            "max_continuous_cached_steps": 3,  # Limit to prevent precision degradation
+            # TaylorSeer parameters [cache-dit only]
+            "enable_taylorseer": False,  # Disabled by default (not suitable for few-step models)
+            "taylorseer_order": 1,  # TaylorSeer polynomial order
+            # SCM (Step Computation Masking) parameters [cache-dit only]
+            "scm_steps_mask_policy": None,  # SCM mask policy: None (disabled), "slow", "medium", "fast", "ultra"
+            "scm_steps_policy": "dynamic",  # SCM steps policy: "dynamic" or "static"
         }
-
     # Configure parallel settings
     parallel_config = DiffusionParallelConfig(
         ulysses_degree=args.ulysses_degree,
@@ -221,30 +171,20 @@ def main():
     # Check if profiling is requested via environment variable
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
-    omni_kwargs = dict(
+    omni = Omni(
         model=args.model,
         enable_layerwise_offload=args.enable_layerwise_offload,
         vae_use_slicing=args.vae_use_slicing,
         vae_use_tiling=args.vae_use_tiling,
+        boundary_ratio=args.boundary_ratio,
+        flow_shift=args.flow_shift,
+        enable_cache_dit_summary=args.enable_cache_dit_summary,
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
-        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
-    if args.boundary_ratio is not None:
-        omni_kwargs["boundary_ratio"] = args.boundary_ratio
-    if args.flow_shift is not None:
-        omni_kwargs["flow_shift"] = args.flow_shift
-    if args.quantization is not None:
-        omni_kwargs["quantization"] = args.quantization
-    if args.cache_backend is not None:
-        omni_kwargs["cache_backend"] = args.cache_backend
-        omni_kwargs["cache_config"] = cache_config
-        omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
-
-    omni = Omni(**omni_kwargs)
 
     if profiler_enabled:
         print("[Profiler] Starting profiling...")
@@ -264,25 +204,22 @@ def main():
     print(f"  Video size: {args.width}x{args.height}")
     print(f"{'=' * 60}\n")
 
-    prompt_dict = {"prompt": args.prompt}
-    if args.negative_prompt:
-        prompt_dict["negative_prompt"] = args.negative_prompt
-
-    sampling_kwargs = dict(
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
-    )
-    if args.guidance_scale_high is not None:
-        sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
-
     generation_start = time.perf_counter()
     frames = omni.generate(
-        prompt_dict,
-        OmniDiffusionSamplingParams(**sampling_kwargs),
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+            frame_rate=frame_rate,
+        ),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
@@ -303,6 +240,8 @@ def main():
             audio = frames.multimodal_output["audio"]
         if frames.is_pipeline_output and frames.request_output is not None:
             inner_output = frames.request_output
+            if isinstance(inner_output, list):
+                inner_output = inner_output[0] if inner_output else None
             if isinstance(inner_output, OmniRequestOutput):
                 if inner_output.multimodal_output and "audio" in inner_output.multimodal_output:
                     audio = inner_output.multimodal_output["audio"]

examples/offline_inference/voxtral_tts/README.md

@@ -1,64 +0,0 @@
-# Voxtral TTS Offline Inference
-
-`end2end.py` runs Voxtral TTS end-to-end offline inference using vLLM. It supports both blocking (`Omni`) and streaming (`AsyncOmni`) generation, batched prompts with configurable concurrency, and voice selection via preset name or reference audio file.
-
-When `mistral_common` has `SpeechRequest` support, prompt token IDs are built via `encode_speech_request`. Otherwise, the script falls back to manual token construction.
-
-## Usage Examples
-
-
-```bash
-# Basic single-prompt with cheerful_female voice preset
-python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
-    --write-audio --voice cheerful_female \
-    --model mistralai/Voxtral-4B-TTS-2603 \
-    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
-
-# 32 replicate prompts with cheerful_female voice preset
-python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
-    --num-prompts 32 --write-audio --voice cheerful_female \
-    --model mistralai/Voxtral-4B-TTS-2603 \
-    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
-
-# Streaming with neutral_female voice preset
-python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
-    --streaming --write-audio --voice neutral_female \
-    --model mistralai/Voxtral-4B-TTS-2603 \
-    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
-
-# 32 prompts, 8 concurrent requests per wave, streaming with neutral_female voice
-python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
-    --num-prompts 32 --concurrency 8 --streaming --write-audio --voice neutral_female \
-    --model mistralai/Voxtral-4B-TTS-2603 \
-    --text "That eerie silence after the first storm was just the calm before another round of chaos, wasn't it?"
-
-# Short debug prompt with reference audio
-# Note: Reference audio capability is not yet released.
-python3 examples/offline_inference/voxtral_tts/end2end.py \
-    --stage-configs-path vllm_omni/model_executor/stage_configs/voxtral_tts.yaml \
-    --write-audio \
-    --model mistralai/Voxtral-4B-TTS-2603 \
-    --text "This is a test message." \
-    --audio-path path/to/reference_audio.wav
-```
-
-## Arguments
-
-| Argument | Description |
-|---|---|
-| `--model PATH` | HuggingFace repo ID or local directory path (default: `mistralai/Voxtral-4B-TTS-2603`) |
-| `--text TEXT` | Text to synthesize (default: `"This is a test message."`) |
-| `--audio-path PATH` | Path to reference audio file for voice cloning |
-| `--output-dir DIR` | Directory to write output WAV files (default: `output_audio`) |
-| `--stage-configs-path PATH` | Path to stage configs YAML (currently it must be set for VoxtralTTS) |
-| `--num-prompts N` | Number of replicate prompts to run for measuring performance (default: 1) |
-| `--streaming` | Use streaming generation via `AsyncOmni` (default: blocking `Omni`) |
-| `--concurrency N` | Max concurrent requests per wave (must be used with `--streaming`, must evenly divide `--num-prompts`) |
-| `--voice NAME` | Voice preset to use instead of reference audio file (e.g., casual_female, casual_male, cheerful_female, neutral_female, neutral_male) |
-| `--write-audio` | Write generated audio to WAV files |
-| `--profiling-mode` | Enable profiling mode (reduces max tokens to 50) |
-| `--log-stats` | Enable detailed statistics logging |

examples/offline_inference/voxtral_tts/end2end.py

@@ -1,377 +0,0 @@
-"""
-This example shows how to use vLLM for running Voxtral TTS
-"""
-
-import asyncio
-import gc
-import logging
-import os
-import time
-import uuid
-from argparse import Namespace
-from pathlib import Path
-from typing import Any
-
-import numpy as np
-import soundfile as sf
-import torch
-from mistral_common.protocol.instruct.chunk import TextChunk
-
-try:
-    from mistral_common.protocol.speech.request import SpeechRequest
-    from mistral_common.tokens.tokenizers.mistral import MistralTokenizer
-except ImportError:
-    raise ImportError(
-        "Could not import SpeechRequest or MistralTokenizer from mistral_common. "
-        "Please pull the latest mistral-common code and install it with: "
-        "pip install -e ."
-    )
-from vllm import SamplingParams
-from vllm.utils.argparse_utils import FlexibleArgumentParser
-
-from vllm_omni import AsyncOmni
-from vllm_omni.entrypoints.omni import Omni
-
-logger = logging.getLogger(__name__)
-
-
-# ---- Streaming version ----
-async def run_streaming(inputs, sampling_params_list, model_name, args, output_dir):
-    async_omni = AsyncOmni(
-        model=model_name,
-        stage_configs_path=args.stage_configs_path,
-        log_stats=args.log_stats,
-    )
-
-    # Normalize to a list so batch and single-input share the same code path
-    if isinstance(inputs, list):
-        inputs_list = inputs
-    else:
-        inputs_list = [inputs]
-
-    total_audio_dur = 0.0
-    total_gen_time = 0.0
-    total_ttfa = 0.0
-    ttfa_count = 0
-    total_waits = 0
-    total_chunks = 0
-    results_lock = asyncio.Lock()
-
-    async def _generate_one(batch_idx, single_input):
-        nonlocal total_audio_dur, total_gen_time, total_ttfa, ttfa_count
-        nonlocal total_waits, total_chunks
-        request_id = str(uuid.uuid4())
-        all_audio_chunks = []
-        chunk_arrival_times = []
-        chunk_durations = []
-        chunk_idx = 0
-        gen_start = time.time()
-        ttfa = None
-        accumulated_sample = 0
-
-        async for stage_output in async_omni.generate(single_input, request_id, sampling_params_list):
-            mm_output = stage_output.multimodal_output
-            finished = stage_output.finished
-            if not mm_output or "audio" not in mm_output:
-                continue
-
-            now = time.time()
-            if ttfa is None:
-                ttfa = now - gen_start
-
-            audio_chunk = mm_output["audio"]
-            # audio_chunk is a 1-D tensor from decode_helper_async
-            if isinstance(audio_chunk, torch.Tensor):
-                if finished:
-                    # Last chunk may return whole audio instead of chunk delta.
-                    # Cut accumulated samples from previous chunks.
-                    audio_numpy = audio_chunk[accumulated_sample:].float().detach().cpu().numpy()
-                else:
-                    audio_numpy = audio_chunk.float().detach().cpu().numpy()
-            elif isinstance(audio_chunk, list):
-                # Audio_chunk list contain all previous chunks. Use chunk_idx to index
-                audio_numpy = audio_chunk[chunk_idx].float().detach().cpu().numpy()
-            else:
-                audio_numpy = audio_chunk
-
-            accumulated_sample += len(audio_numpy)
-            all_audio_chunks.append(audio_numpy)
-            chunk_arrival_times.append(now)
-            chunk_durations.append(len(audio_numpy) / 24000)
-            chunk_idx += 1
-
-        gen_elapsed = time.time() - gen_start
-
-        # Analyze wait/no-wait per chunk:
-        # A "wait" means the client has finished playing all previous audio
-        # before this chunk arrived, so there is a playback gap.
-        # buffer_time > 0: client still has audio to play (no wait)
-        # buffer_time <= 0: client ran out of audio and had to wait
-        chunk_labels = []
-        accumulated_audio_dur = 0.0
-        if chunk_arrival_times:
-            first_arrival = chunk_arrival_times[0]
-            for i in range(len(chunk_arrival_times)):
-                if i == 0:
-                    chunk_labels.append("no_wait")
-                else:
-                    playback_elapsed = chunk_arrival_times[i] - first_arrival
-                    buffer_time = accumulated_audio_dur - playback_elapsed
-                    if buffer_time <= 0:
-                        chunk_labels.append("wait")
-                    else:
-                        chunk_labels.append("no_wait")
-                accumulated_audio_dur += chunk_durations[i]
-
-        req_wait_count = sum(1 for chunk_label in chunk_labels if chunk_label == "wait")
-        req_wait_rate = req_wait_count / len(chunk_labels) if chunk_labels else 0.0
-
-        # Concatenate all chunks for this request's audio
-        if all_audio_chunks:
-            full_audio = np.concatenate(all_audio_chunks)
-            output_audio_dur = len(full_audio) / 24000
-            output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
-            if args.write_audio:
-                sf.write(output_path, full_audio, 24000)
-                print(
-                    f"Request {batch_idx}: saved {len(full_audio)} samples ({output_audio_dur:.2f}s) to {output_path}"
-                )
-            # Per-chunk wait details
-            for i, label in enumerate(chunk_labels):
-                dur_ms = chunk_durations[i] * 1000
-                arrival_ms = (chunk_arrival_times[i] - first_arrival) * 1000 if i > 0 else 0.0
-                print(
-                    f"  Request {batch_idx} chunk {i}: {label} | arrived={arrival_ms:.1f}ms | chunk_dur={dur_ms:.1f}ms"
-                )
-            print(
-                f"Request {batch_idx}: TTFA={ttfa:.4f}s | "
-                f"Generation={gen_elapsed:.4f}s | "
-                f"Audio={output_audio_dur:.2f}s | "
-                f"RTF={output_audio_dur / gen_elapsed:.4f} | "
-                f"WaitRate={req_wait_rate:.2%} ({req_wait_count}/{len(chunk_labels)})"
-            )
-            async with results_lock:
-                total_audio_dur += output_audio_dur
-                total_gen_time += gen_elapsed
-                total_waits += req_wait_count
-                total_chunks += len(chunk_labels)
-                if ttfa is not None:
-                    total_ttfa += ttfa
-                    ttfa_count += 1
-
-    # Launch requests in waves of `concurrency` size
-    concurrency = args.concurrency or len(inputs_list)
-    gen_start_all = time.time()
-    for wave_start in range(0, len(inputs_list), concurrency):
-        wave = inputs_list[wave_start : wave_start + concurrency]
-        wave_idx_offset = wave_start
-        print(f"\n--- Wave {wave_start // concurrency + 1} (requests {wave_start}-{wave_start + len(wave) - 1}) ---")
-        await asyncio.gather(*[_generate_one(wave_idx_offset + i, inp) for i, inp in enumerate(wave)])
-    generation_time = time.time() - gen_start_all
-    avg_ttfa = total_ttfa / ttfa_count if ttfa_count else float("nan")
-    overall_wait_rate = total_waits / total_chunks if total_chunks else float("nan")
-    print(
-        f"\nAll requests: Generation={generation_time:.4f}s | "
-        f"TotalAudio={total_audio_dur:.2f}s | "
-        f"Concurrency={concurrency} | "
-        f"AvgTTFA={avg_ttfa:.4f}s | "
-        f"RTF(total)={total_audio_dur / generation_time:.4f} | "
-        f"RTF(per-request)={total_audio_dur / total_gen_time:.4f} | "
-        f"WaitRate={overall_wait_rate:.2%} ({total_waits}/{total_chunks})"
-    )
-
-    async_omni.shutdown()
-    torch.cuda.empty_cache()
-    gc.collect()
-
-
-# ---- Non-streaming version ----
-def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir):
-    llm = Omni(
-        model=model_name,
-        log_stats=args.log_stats,
-        stage_configs_path=args.stage_configs_path,
-    )
-
-    if args.profiling_mode:
-        llm.start_profile()
-
-    vllm_start = time.time()
-    outputs = llm.generate(inputs, sampling_params_list)
-    vllm_elapsed = time.time() - vllm_start
-
-    if args.profiling_mode:
-        llm.stop_profile()
-        time.sleep(10)
-
-    print(f"vLLM run time: {vllm_elapsed:.4f}s")
-    output_audio_dur = 0.0
-
-    for batch_idx, o in enumerate(outputs):
-        audio_tensor = torch.cat(o.multimodal_output["audio"])
-        audio_array = audio_tensor.tolist()
-        output_audio_dur += float(len(audio_array)) / 24000
-        if args.write_audio:
-            output_path = os.path.join(output_dir, f"tts_output_{batch_idx}.wav")
-            sf.write(output_path, audio_array, 24000)
-            print(f"Audio saved to {output_path}")
-
-    print(f"Total audio duration: {output_audio_dur:.2f}s")
-    print(f"RTF: {output_audio_dur / vllm_elapsed:.4f}")
-
-    del llm
-    torch.cuda.empty_cache()
-    gc.collect()
-
-
-def parse_args() -> Namespace:
-    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with Voxtral TTS")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="mistralai/Voxtral-4B-TTS-2603",
-        help="Model name or path.",
-    )
-    parser.add_argument(
-        "--text",
-        type=str,
-        default="This is a test message.",
-        help="Text to synthesize.",
-    )
-    parser.add_argument(
-        "--audio-path",
-        type=str,
-        default=None,
-        help="Path to reference audio file for voice cloning.",
-    )
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        default="output_audio",
-        help="Directory to write output wav files.",
-    )
-    parser.add_argument(
-        "--stage-configs-path",
-        type=str,
-        default=None,
-        help="Path to stage configs YAML. Auto-resolved from model if not set.",
-    )
-    parser.add_argument(
-        "--num-prompts", type=int, default=1, help="Number of replicate prompts to run for measuring performance"
-    )
-    parser.add_argument(
-        "--profiling-mode",
-        default=False,
-        action="store_true",
-        help="Set max_num_tokens=2 to reduce profiling time",
-    )
-    parser.add_argument(
-        "--log-stats",
-        action="store_true",
-        default=False,
-        help="Enable writing detailed statistics (default: disabled)",
-    )
-    parser.add_argument(
-        "--write-audio",
-        action="store_true",
-        default=False,
-        help="Write audio output to files.",
-    )
-    parser.add_argument(
-        "--streaming",
-        action="store_true",
-        default=False,
-        help="Use streaming generation via AsyncOmni instead of blocking Omni.",
-    )
-    parser.add_argument(
-        "--concurrency",
-        type=int,
-        default=None,
-        help="Max concurrent requests per wave (default: all at once). Must evenly divide --num-prompts.",
-    )
-    parser.add_argument(
-        "--voice",
-        type=str,
-        default=None,
-        help="Voice to use instead of audio file.",
-    )
-    return parser.parse_args()
-
-
-def compose_request(
-    model_name: str,
-    text_chunk: TextChunk,
-    audio_prompt_file: str,
-    args: Any,
-) -> dict:
-    """Build the full TTS input dict (prompt_token_ids, multi_modal_data or additional_information)."""
-    inputs: dict[str, Any] = {}
-    if Path(model_name).is_dir():
-        mistral_tokenizer = MistralTokenizer.from_file(str(Path(model_name) / "tekken.json"))
-    else:
-        mistral_tokenizer = MistralTokenizer.from_hf_hub(model_name)
-    instruct_tokenizer = mistral_tokenizer.instruct_tokenizer
-
-    if args.voice is not None:
-        tokenized = instruct_tokenizer.encode_speech_request(SpeechRequest(input=text_chunk.text, voice=args.voice))
-        inputs["additional_information"] = {"voice": [args.voice]}
-        inputs["prompt_token_ids"] = tokenized.tokens
-    else:
-        with open(audio_prompt_file, "rb") as f:
-            ref_audio_bytes = f.read()
-        tokenized = instruct_tokenizer.encode_speech_request(
-            SpeechRequest(input=text_chunk.text, ref_audio=ref_audio_bytes)
-        )
-        audio = tokenized.audios[0]
-        inputs["multi_modal_data"] = {"audio": [(audio.audio_array, audio.sampling_rate)]}
-        inputs["prompt_token_ids"] = tokenized.tokens
-
-    return inputs
-
-
-def main(args: Any) -> None:
-    max_num_tokens = 50 if args.profiling_mode else 2500
-
-    model_name = args.model
-    output_dir = args.output_dir
-
-    if args.voice is None and args.audio_path is None:
-        raise ValueError("Either --voice or --audio-path must be provided.")
-
-    audio_prompt_file = args.audio_path
-    text_chunk = TextChunk(text=args.text)
-
-    if args.write_audio:
-        os.makedirs(output_dir, exist_ok=True)
-
-    inputs = compose_request(model_name, text_chunk, audio_prompt_file, args)
-
-    sampling_params = SamplingParams(
-        max_tokens=max_num_tokens,
-    )
-    sampling_params_list = [
-        sampling_params,
-        sampling_params,
-    ]
-
-    if args.num_prompts > 1:
-        inputs = [inputs] * args.num_prompts
-
-    if args.concurrency is not None:
-        assert args.streaming, "--concurrency must be used with --streaming on AsyncOmni"
-        assert args.num_prompts % args.concurrency == 0, (
-            f"--num-prompts ({args.num_prompts}) must be divisible by --concurrency ({args.concurrency})"
-        )
-
-    torch.cuda.empty_cache()
-    gc.collect()
-
-    if args.streaming:
-        asyncio.run(run_streaming(inputs, sampling_params_list, model_name, args, output_dir))
-    else:
-        run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir)
-
-
-if __name__ == "__main__":
-    args = parse_args()
-    main(args)

examples/online_serving/image_to_video/run_curl_hunyuan_video_15.sh

@@ -1,68 +0,0 @@
-#!/bin/bash
-# HunyuanVideo-1.5 image-to-video curl example using the async video job API.
-
-set -euo pipefail
-
-INPUT_IMAGE="${INPUT_IMAGE:-test_input.jpg}"
-BASE_URL="${BASE_URL:-http://localhost:8099}"
-OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_i2v.mp4}"
-POLL_INTERVAL="${POLL_INTERVAL:-2}"
-
-if [ ! -f "$INPUT_IMAGE" ]; then
-    echo "Input image not found: $INPUT_IMAGE"
-    echo "Provide an image via INPUT_IMAGE env var."
-    exit 1
-fi
-
-create_response=$(
-  curl -sS -X POST "${BASE_URL}/v1/videos" \
-    -H "Accept: application/json" \
-    -F "prompt=The camera follows the puppy as it runs forward on the grass, its four legs alternating steps, its tail held high and wagging side to side." \
-    -F "input_reference=@${INPUT_IMAGE}" \
-    -F "size=832x480" \
-    -F "num_frames=33" \
-    -F "fps=24" \
-    -F "num_inference_steps=30" \
-    -F "guidance_scale=6.0" \
-    -F "flow_shift=5.0" \
-    -F "seed=42"
-)
-
-video_id="$(echo "${create_response}" | jq -r '.id')"
-if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
-  echo "Failed to create video job:"
-  echo "${create_response}" | jq .
-  exit 1
-fi
-
-echo "Created video job ${video_id}"
-echo "${create_response}" | jq .
-
-while true; do
-  status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
-  status="$(echo "${status_response}" | jq -r '.status')"
-
-  case "${status}" in
-    queued|in_progress)
-      echo "Video job ${video_id} status: ${status}"
-      sleep "${POLL_INTERVAL}"
-      ;;
-    completed)
-      echo "${status_response}" | jq .
-      break
-      ;;
-    failed)
-      echo "Video generation failed:"
-      echo "${status_response}" | jq .
-      exit 1
-      ;;
-    *)
-      echo "Unexpected status response:"
-      echo "${status_response}" | jq .
-      exit 1
-      ;;
-  esac
-done
-
-curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
-echo "Saved video to ${OUTPUT_PATH}"

examples/online_serving/image_to_video/run_server_hunyuan_video_15.sh

@@ -1,31 +0,0 @@
-#!/bin/bash
-# HunyuanVideo-1.5 image-to-video online serving startup script
-#
-# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB
-# 720p: needs FP8 + VAE tiling, ~35 GB VRAM
-
-MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v}"
-PORT="${PORT:-8099}"
-FLOW_SHIFT="${FLOW_SHIFT:-5.0}"
-QUANTIZATION="${QUANTIZATION:-}"
-CACHE_BACKEND="${CACHE_BACKEND:-none}"
-
-echo "Starting HunyuanVideo-1.5 I2V server..."
-echo "Model: $MODEL"
-echo "Port: $PORT"
-echo "Flow shift: $FLOW_SHIFT"
-echo "Quantization: ${QUANTIZATION:-none}"
-echo "Cache backend: $CACHE_BACKEND"
-
-EXTRA_FLAGS=""
-if [ -n "$QUANTIZATION" ]; then
-    EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION"
-fi
-if [ "$CACHE_BACKEND" != "none" ]; then
-    EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND"
-fi
-
-vllm serve "$MODEL" --omni \
-    --port "$PORT" \
-    --flow-shift "$FLOW_SHIFT" \
-    $EXTRA_FLAGS

examples/online_serving/qwen3_omni/qwen3_omni_moe_thinking.yaml

@@ -9,9 +9,9 @@ stage_args:
   - stage_id: 0
     runtime:
       devices: "0,1"
+      max_batch_size: 1
     engine_args:
       model_stage: thinker
-      max_num_seqs: 1
       model_arch: Qwen3OmniMoeForConditionalGeneration
       worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler

examples/online_serving/text_to_video/run_curl_hunyuan_video_15.sh

@@ -1,60 +0,0 @@
-#!/bin/bash
-# HunyuanVideo-1.5 text-to-video curl example using the async video job API.
-
-set -euo pipefail
-
-BASE_URL="${BASE_URL:-http://localhost:8098}"
-OUTPUT_PATH="${OUTPUT_PATH:-hunyuan_video_15_t2v.mp4}"
-POLL_INTERVAL="${POLL_INTERVAL:-2}"
-
-create_response=$(
-  curl -sS -X POST "${BASE_URL}/v1/videos" \
-    -H "Accept: application/json" \
-    -F "prompt=A little girl wearing a straw hat runs through a summer meadow full of wildflowers. A wide shot is used, with the camera panning right to follow her." \
-    -F "size=832x480" \
-    -F "num_frames=33" \
-    -F "fps=24" \
-    -F "num_inference_steps=30" \
-    -F "guidance_scale=6.0" \
-    -F "flow_shift=5.0" \
-    -F "seed=42"
-)
-
-video_id="$(echo "${create_response}" | jq -r '.id')"
-if [ -z "${video_id}" ] || [ "${video_id}" = "null" ]; then
-  echo "Failed to create video job:"
-  echo "${create_response}" | jq .
-  exit 1
-fi
-
-echo "Created video job ${video_id}"
-echo "${create_response}" | jq .
-
-while true; do
-  status_response="$(curl -sS "${BASE_URL}/v1/videos/${video_id}")"
-  status="$(echo "${status_response}" | jq -r '.status')"
-
-  case "${status}" in
-    queued|in_progress)
-      echo "Video job ${video_id} status: ${status}"
-      sleep "${POLL_INTERVAL}"
-      ;;
-    completed)
-      echo "${status_response}" | jq .
-      break
-      ;;
-    failed)
-      echo "Video generation failed:"
-      echo "${status_response}" | jq .
-      exit 1
-      ;;
-    *)
-      echo "Unexpected status response:"
-      echo "${status_response}" | jq .
-      exit 1
-      ;;
-  esac
-done
-
-curl -sS -L "${BASE_URL}/v1/videos/${video_id}/content" -o "${OUTPUT_PATH}"
-echo "Saved video to ${OUTPUT_PATH}"

examples/online_serving/text_to_video/run_server_hunyuan_video_15.sh

@@ -1,31 +0,0 @@
-#!/bin/bash
-# HunyuanVideo-1.5 text-to-video online serving startup script
-#
-# 480p: ~35 GB VRAM (BF16), fits 1x A100 80GB
-# 720p: needs FP8 + VAE tiling, ~35 GB VRAM
-
-MODEL="${MODEL:-hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v}"
-PORT="${PORT:-8098}"
-FLOW_SHIFT="${FLOW_SHIFT:-5.0}"
-QUANTIZATION="${QUANTIZATION:-}"
-CACHE_BACKEND="${CACHE_BACKEND:-none}"
-
-echo "Starting HunyuanVideo-1.5 T2V server..."
-echo "Model: $MODEL"
-echo "Port: $PORT"
-echo "Flow shift: $FLOW_SHIFT"
-echo "Quantization: ${QUANTIZATION:-none}"
-echo "Cache backend: $CACHE_BACKEND"
-
-EXTRA_FLAGS=""
-if [ -n "$QUANTIZATION" ]; then
-    EXTRA_FLAGS="$EXTRA_FLAGS --quantization $QUANTIZATION"
-fi
-if [ "$CACHE_BACKEND" != "none" ]; then
-    EXTRA_FLAGS="$EXTRA_FLAGS --cache-backend $CACHE_BACKEND"
-fi
-
-vllm serve "$MODEL" --omni \
-    --port "$PORT" \
-    --flow-shift "$FLOW_SHIFT" \
-    $EXTRA_FLAGS

examples/online_serving/voxtral_tts/gradio_demo.py

@@ -1,269 +0,0 @@
-"""
-- Make sure to install the following for this example to function correctly:
-- `pip install -e .`
-- `pip install gradio==5.50 mistral_common=1.10.0`
-
-Example use case:
-
-python examples/online_serving/voxtral_tts/gradio_demo.py --host slurm-199-077 --port 8000
-
-"""
-
-import argparse
-import io
-import json
-import logging
-import socket
-import uuid
-from datetime import datetime
-from pathlib import Path
-from typing import Any
-
-import gradio as gr
-import httpx
-import numpy as np
-import soundfile as sf
-
-logger = logging.getLogger()
-
-LOGFORMAT = "%(asctime)s - %(levelname)s - %(message)s"
-TIMEFORMAT = "%Y-%m-%d %H:%M:%S"
-
-logging.basicConfig(level=logging.INFO, force=True, format=LOGFORMAT, datefmt=TIMEFORMAT)
-logger.setLevel(logging.INFO)
-
-
-_BASE_URL = f"{socket.gethostname()}:7860"
-
-
-def run_inference(
-    voice_name: str,
-    text_prompt: str,
-    base_url: str,
-    model: str,
-) -> tuple[int, np.ndarray]:
-    """Call /v1/audio/speech and return (sample_rate, audio_array)."""
-    text_prompt = text_prompt.strip()
-    if not text_prompt:
-        raise gr.Error("Please enter a text prompt.")
-
-    payload: dict[str, Any] = {
-        "input": text_prompt,
-        "model": model,
-        "response_format": "wav",
-        "voice": voice_name,
-    }
-
-    response = httpx.post(
-        f"{base_url}/audio/speech",
-        json=payload,
-        timeout=120.0,
-    )
-    response.raise_for_status()
-
-    audio_array, sr = sf.read(io.BytesIO(response.content), dtype="float32")
-    return sr, audio_array
-
-
-def _save_example(
-    outputs_dir: Path,
-    voice_name: str,
-    text_prompt: str,
-    sr: int,
-    audio_array: np.ndarray,
-) -> tuple[str, str]:
-    """
-    Save inputs/outputs for sharing.
-    Returns (share_id, saved_audio_path)
-    """
-    share_id = uuid.uuid4().hex
-
-    # Save generated audio
-    saved_audio_path = outputs_dir / f"{share_id}_gen.wav"
-    sf.write(str(saved_audio_path), audio_array, sr)
-
-    meta = {
-        "id": share_id,
-        "created_at": datetime.utcnow().isoformat(),
-        "voice_name": voice_name,
-        "text_prompt": text_prompt,
-        "generated_audio_path": str(saved_audio_path),
-    }
-
-    with open(outputs_dir / f"{share_id}.json", "w") as f:
-        json.dump(meta, f, ensure_ascii=False)
-
-    return share_id, str(saved_audio_path)
-
-
-def _load_from_share(
-    outputs_dir: Path | None,
-    request: gr.Request,
-) -> tuple[str | None, str, str | None, dict[str, Any], str]:
-    """
-    Called on page load. If ?share_id=... is present, load stored example.
-    Returns: (voice_name, text_prompt, output_audio, submit_btn_update, share_link_text)
-    """
-    if outputs_dir is None:
-        return "Neutral_Male", "", None, gr.update(interactive=False), ""
-
-    share_id = None
-    if request and request.query_params:
-        share_id = request.query_params.get("share_id")
-
-    if not share_id:
-        return "Neutral_Male", "", None, gr.update(interactive=False), ""
-
-    meta_path = outputs_dir / f"{share_id}.json"
-    if not meta_path.exists():
-        logger.warning("No stored example for share_id=%s", share_id)
-        return "Neutral_Male", "", None, gr.update(interactive=False), ""
-
-    with open(meta_path) as f:
-        meta = json.load(f)
-
-    voice = meta.get("voice_name", "Neutral_Male")
-    text_prompt = meta.get("text_prompt", "")
-    gen_path = meta.get("generated_audio_path")
-
-    if _BASE_URL:
-        share_link = f"http://{_BASE_URL}?share_id={share_id}"
-    else:
-        share_link = f"Use this query on your current URL: ?share_id={share_id}"
-
-    return voice, text_prompt, gen_path, gr.update(interactive=True), share_link
-
-
-def main(
-    model: str,
-    host: str,
-    port: str,
-    output_dir: str | None = None,
-) -> None:
-    base_url = f"http://{host}:{port}/v1"
-    logger.info(f"Using speech API at: {base_url}/audio/speech")
-
-    outputs_dir: Path | None = None
-    if output_dir is not None:
-        outputs_dir = Path(output_dir)
-        outputs_dir.mkdir(parents=True, exist_ok=True)
-
-    with gr.Blocks(title="Voxtral TTS", fill_height=True) as demo:
-        gr.Markdown("## Voxtral TTS")
-
-        with gr.Row():
-            with gr.Column():
-                voice_name = gr.Dropdown(
-                    ["Neutral_Male", "Neutral_Female", "Cheerful_Female", "Casual_Male", "Casual_Female"],
-                    label="Voice",
-                    value="Neutral_Male",
-                )
-                text_prompt = gr.Textbox(
-                    label="Text prompt",
-                    placeholder="Enter the text you want to synthesize...",
-                    lines=4,
-                )
-                with gr.Row():
-                    reset_btn = gr.Button("Clear")
-                    submit_btn = gr.Button("Generate audio", interactive=False)
-
-            with gr.Column():
-                output_audio = gr.Audio(
-                    label="Generated audio",
-                    show_download_button=True,
-                    interactive=False,
-                    autoplay=True,
-                    type="filepath",
-                )
-                share_link_box = gr.Textbox(
-                    label="Shareable link",
-                    interactive=False,
-                    show_copy_button=True,
-                    visible=outputs_dir is not None,
-                )
-
-        # --- UI logic: disable submit until text is non-empty ---
-        def _toggle_submit(text: str):
-            enabled = bool(text.strip())
-            return gr.update(interactive=enabled)
-
-        text_prompt.change(
-            fn=_toggle_submit,
-            inputs=[text_prompt],
-            outputs=submit_btn,
-        )
-
-        # --- Wiring inference + persistence to the button ---
-        def _on_submit(voice: str, text: str):
-            assert text.strip() != ""
-            sr, audio_array = run_inference(voice, text, base_url, model)
-            if outputs_dir is not None:
-                share_id, saved_audio_path = _save_example(
-                    outputs_dir,
-                    voice_name=voice,
-                    text_prompt=text,
-                    sr=sr,
-                    audio_array=audio_array,
-                )
-                share_link = f"{_BASE_URL}?share_id={share_id}"
-                return saved_audio_path, share_link
-            return (sr, audio_array), ""
-
-        submit_btn.click(
-            fn=_on_submit,
-            inputs=[voice_name, text_prompt],
-            outputs=[output_audio, share_link_box],
-        )
-
-        # --- Clear everything and disable submit again ---
-        def _on_reset():
-            return (
-                "Neutral_Male",  # voice_name
-                "",  # text_prompt
-                None,  # output_audio
-                gr.update(interactive=False),  # submit_btn
-                "",  # share_link_box
-            )
-
-        reset_btn.click(
-            fn=_on_reset,
-            inputs=[],
-            outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
-        )
-
-        demo.load(
-            fn=lambda request: _load_from_share(outputs_dir, request),
-            inputs=[],
-            outputs=[voice_name, text_prompt, output_audio, submit_btn, share_link_box],
-        )
-
-    launch_kwargs: dict[str, Any] = {
-        "server_name": "0.0.0.0",
-        "share": True,
-    }
-    if outputs_dir is not None:
-        launch_kwargs["allowed_paths"] = [str(outputs_dir)]
-    demo.launch(**launch_kwargs)
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Voxtral TTS Gradio Demo")
-    parser.add_argument("--model", type=str, default="mistralai/Voxtral-4B-TTS-2603", help="Name of model repo on HF")
-    parser.add_argument("--host", type=str, default="localhost", help="Name of host")
-    parser.add_argument("--port", type=str, default="8091", help="port number")
-    parser.add_argument(
-        "--output-dir",
-        type=str,
-        default=None,
-        help="Directory to save generated audio and share links. "
-        "If not provided, save/share functionality is disabled.",
-    )
-
-    args = parser.parse_args()
-
-    main(
-        model=args.model,
-        host=args.host,
-        port=args.port,
-        output_dir=args.output_dir,
-    )