p

Signed-off-by: khluu <khluu000@gmail.com>

.buildkite/bootstrap-intel-omni.sh

@@ -0,0 +1,51 @@
+#!/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

@@ -0,0 +1,16 @@
+# .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.17.0"
+        priority: 100
+        timeout_in_minutes: 60
+        soft_fail: true

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

@@ -0,0 +1,56 @@
+#!/bin/bash
+
+# This script build the XPU docker image and run the offline inference inside the container.
+set -ex
+
+omni_source_dir=$(git rev-parse --show-toplevel)
+
+base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
+image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
+container_name="xpu_${BUILDKITE_COMMIT}_$(
+    tr -dc A-Za-z0-9 </dev/urandom | head -c 10
+    echo
+)"
+
+cd "${omni_source_dir}"
+if [ -z "$(docker images -q "${base_image_name}")" ]; then
+    docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
+fi
+
+# Try building the docker image
+docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
+
+# Setup cleanup
+remove_docker_container() {
+    docker rm -f "${container_name}" || true
+    docker image rm -f "${image_name}" || true
+    docker system prune -f || true
+}
+trap remove_docker_container EXIT
+
+HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
+mkdir -p "${HF_CACHE}"
+HF_MOUNT="/root/.cache/huggingface"
+
+time timeout -k 30 30m docker run \
+    --device /dev/dri:/dev/dri \
+    --net=host \
+    --ipc=host \
+    -v /dev/dri/by-path:/dev/dri/by-path \
+    -v "${HF_CACHE}:${HF_MOUNT}" \
+    --security-opt seccomp=unconfined \
+    --entrypoint="" \
+    -e VLLM_LOGGING_LEVEL \
+    -e VLLM_OMNI_LOGGING_LEVEL \
+    -e HF_TOKEN \
+    -e ZE_AFFINITY_MASK \
+    --name "${container_name}" \
+    "${image_name}" \
+    bash -c '
+    set -e
+    echo $ZE_AFFINITY_MASK
+    pip install tblib==3.1.0
+    cd /workspace/vllm-omni
+    pytest -v -s -m "core_model and xpu and B60"
+    pytest -v -s -m "advanced_model and xpu and B60"
+'

.buildkite/test-amd.yaml

@@ -101,6 +101,7 @@ 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,6 +12,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -29,6 +30,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -46,6 +48,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -64,6 +67,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -81,6 +85,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -98,6 +103,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -116,6 +122,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -134,6 +141,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -152,6 +160,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -178,6 +187,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -197,6 +207,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -226,6 +237,11 @@ 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:
@@ -261,6 +277,11 @@ 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,37 +293,50 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  # - label: "Bagel Text2Img Model Test with H100"
-  #   timeout_in_minutes: 30
-  #   depends_on: upload-merge-pipeline
-  #   commands:
-  #     - 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
+  - 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

.buildkite/test-nightly.yml

@@ -30,6 +30,11 @@ 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:
@@ -58,6 +63,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -86,6 +92,11 @@ 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:
@@ -116,10 +127,11 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: ":full_moon: Omni Model Perf Test with H100"
+  - label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
     key: nightly-performance
     timeout_in_minutes: 180
     depends_on: upload-nightly-pipeline
@@ -129,6 +141,8 @@ steps:
       - export BENCHMARK_DIR=tests
       - pytest -s -v tests/perf/scripts/run_benchmark.py
       - buildkite-agent artifact upload "tests/*.json"
+      - python tools/nightly/buildkite_testcase_statistics.py -o tests/buildkite_testcase_statistics.html
+      - buildkite-agent artifact upload "tests/*.html"
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -147,6 +161,11 @@ 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:
@@ -158,17 +177,19 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-  - label: ":email: Nightly Perf Collection & Email"
+  - label: ":email: Nightly Collection & Email"
     key: nightly-perf-distribution
-    depends_on: nightly-performance
+    depends_on:
+    - nightly-performance
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl
       - export DEFAULT_INPUT_DIR=tests
       - export DEFAULT_OUTPUT_DIR=tests
       - buildkite-agent artifact download "tests/*.json" . --step nightly-performance
+      - buildkite-agent artifact download "tests/*.html" . --step nightly-performance
       - python tools/nightly/generate_nightly_perf_excel.py
-      - python tools/nightly/send_nightly_perf_email.py
+      - python tools/nightly/send_nightly_email.py --report-file "tests/*.xlsx, tests/*.html"
       - buildkite-agent artifact upload "tests/*.xlsx"
     agents:
       queue: "cpu_queue_premerge"

.buildkite/test-ready.yml

@@ -12,6 +12,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -28,21 +29,14 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model CPU offloading Test"
     depends_on: upload-ready-pipeline
     commands:
-      - |
-        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))
-        '
+      - timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -52,6 +46,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -68,6 +63,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -84,6 +80,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -101,6 +98,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -118,6 +116,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -145,6 +144,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -167,6 +167,7 @@ steps:
           propagate-environment: true
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -229,6 +230,7 @@ steps:
           shm-size: "8gb"
           environment:
             - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
@@ -269,39 +271,132 @@ 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
-  #         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
+  - 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

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

@@ -74,28 +74,21 @@ 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 import OmniLLM, create_ar_stage_config, create_dit_stage_config
+      from vllm_omni.entrypoints.omni import Omni
+      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+      from vllm import SamplingParams
 
-      # Create stage configurations
-      ar_config = create_ar_stage_config(
-          stage_id=0,
-          model_path="Qwen/Qwen3-0.6B",
-          input_modalities=["text"],
-          output_modalities=["text"]
+      omni = Omni(
+          model="Qwen/Qwen-Image",
+          stage_configs_path="/path/to/stage_configs.yaml",
       )
 
-      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")
+      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)
       ```
 
       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

@@ -9,7 +9,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
 </h3>
 
 <p align="center">
-| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
+| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
 </p>
 
 
@@ -17,6 +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] **[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.

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/api_client.py

@@ -6,6 +6,7 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
 Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
 """
 
+import asyncio
 import json
 from typing import Any
 
@@ -18,8 +19,8 @@ from .format import (
     audio_to_base64,
     base64_to_audio,
     base64_to_image_tensor,
-    base64_to_video,
     bytes_to_audio,
+    bytes_to_video,
     image_tensor_to_base64,
     image_tensor_to_png_bytes,
     video_to_base64,
@@ -31,10 +32,43 @@ from .types import AudioFormat
 logger = get_logger(__name__)
 
 
+async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            try:
+                return await response.json()
+            except aiohttp.ContentTypeError as e:
+                raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
+async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            return await response.read()
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
 class VLLMOmniClient:
-    def __init__(self, base_url: str, timeout: float | None = None):
+    def __init__(
+        self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
+    ):
         self.base_url = base_url
         self.timeout = aiohttp.ClientTimeout(total=timeout)
+        self.poll_interval = poll_interval
+        self.max_poll_duration = max_poll_duration
 
     async def generate_image(
         self,
@@ -262,33 +296,41 @@ class VLLMOmniClient:
                 content_type="image/png",
             )
 
-        url = self.base_url + "/videos"
         async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            # Start the video generation job
+            url = f"{self.base_url}/videos"
+            data = await url_json(session, url, "post", data=form)
+            if (job_id := data.get("id", None)) is None:
+                raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
+            if (job_status := data.get("status", None)) is None:
+                raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+
+            # Poll for video generation job completion
+            deadline = asyncio.get_running_loop().time() + self.max_poll_duration
+            url = f"{self.base_url}/videos/{job_id}"
+            while job_status not in {"completed", "failed"}:
+                await asyncio.sleep(self.poll_interval)
+
+                data = await url_json(session, url)
+                if (job_status := data.get("status", None)) is None:
+                    raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+                if asyncio.get_running_loop().time() >= deadline:
+                    raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
+
+            if job_status == "failed":
+                raise RuntimeError(f"Video job failed: {data}")
+
+            # Retrieve completed content
+            video_bytes = await url_bytes(session, f"{url}/content")
+
+            # Decode video and make a best effort at cleaning up server resources
             try:
-                async with session.post(url, data=form) as response:
-                    if not response.ok:
-                        error_text = await response.text()
-                        raise (ValueError if response.status < 500 else RuntimeError)(
-                            f"vLLM-Omni API returned status {response.status}: {error_text}"
-                        )
-
-                    try:
-                        data = await response.json()
-                    except aiohttp.ContentTypeError as e:
-                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
-            except aiohttp.ClientError as e:
-                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
-
-        if "data" not in data:
-            raise RuntimeError("API response missing 'data' field - expected OpenAI DALL-E format")
-        if not data["data"]:
-            raise RuntimeError("API returned empty data array")
-        try:
-            base64_str = data["data"][0]["b64_json"]
-        except (KeyError, IndexError):
-            raise RuntimeError("API response missing 'b64_json' field in first data item")
-
-        return base64_to_video(base64_str)
+                return bytes_to_video(video_bytes)
+            finally:
+                try:
+                    await url_json(session, url, "delete")
+                except Exception as exc:
+                    logger.warning("Failed to clean up video job %s: %s", job_id, exc)
 
     async def generate_understanding_chat_completion(
         self,

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/format.py

@@ -143,15 +143,7 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
     return f"data:{mime_type};base64,{base64_str}"
 
 
-def base64_to_video(base64_str: str) -> VideoInput:
-    if base64_str.startswith("data:video"):
-        _, base64_str = base64_str.split(",", 1)
-
-    try:
-        video_bytes = base64.b64decode(base64_str)
-    except Exception as e:
-        raise ValueError(f"Invalid base64 string: {e}")
-
+def bytes_to_video(video_bytes: bytes) -> VideoInput:
     video_buffer = BytesIO(video_bytes)
 
     try:
@@ -208,6 +200,18 @@ def base64_to_video(base64_str: str) -> VideoInput:
     return InputImpl.VideoFromComponents(components)
 
 
+def base64_to_video(base64_str: str) -> VideoInput:
+    if base64_str.startswith("data:video"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        video_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_video(video_bytes)
+
+
 def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
     waveform = audio["waveform"][0]  # Shape: (C, T)
     sample_rate = audio["sample_rate"]

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/utils/validators.py

@@ -70,13 +70,17 @@ def add_sampling_parameters_to_stage(
 
     stages = pipeline_spec["stages"]
     if isinstance(sampling_param_list, dict):
+        sampling_param_list = sampling_param_list.__class__(sampling_param_list)
         sampling_param_list.update(params_to_add)
     elif sampling_param_list is None:
         sampling_param_list = params_to_add.copy()
     else:
         for i, stage in enumerate(stages):
             if stage == stage_type:
-                sampling_param_list[i].update(params_to_add)
+                stage_param = sampling_param_list[i]
+                stage_param = stage_param.__class__(stage_param)
+                stage_param.update(params_to_add)
+                sampling_param_list[i] = stage_param
 
     return sampling_param_list
 

benchmarks/diffusion/README.md

@@ -43,6 +43,27 @@ The benchmark supports three dataset modes via `--dataset`:
 
 ### VBench dataset
 
+`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
+per-request generation fields. In this mode, all requests share CLI values:
+`--width --height --num-frames --fps --num-inference-steps`
+(pass `--width` and `--height` together).
+
+Example (`t2v`):
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+	--base-url http://localhost:8099 \
+	--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+	--task t2v \
+	--dataset vbench \
+	--num-prompts 50 \
+	--width 640 --height 480 \
+	--num-frames 81 --fps 16 \
+	--num-inference-steps 40
+```
+
+Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
+
 If you use i2v/i2i bench datasets and need auto-download support, you may need:
 
 ```bash
@@ -90,7 +111,7 @@ Common optional flags:
 Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
 
 - For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
-- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
+- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
 
 Precedence rules for `trace` (i.e., what actually gets sent):
 

benchmarks/diffusion/backends.py

@@ -35,6 +35,7 @@ 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
 
@@ -55,6 +56,7 @@ 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()
@@ -108,6 +110,18 @@ 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

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_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.
+    - `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.
   - 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_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 "  - 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 "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/vllm_omni/configs/qwen3_tts_bs1.yaml

@@ -10,8 +10,6 @@ stage_args:
     engine_args:
       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
@@ -43,8 +41,6 @@ stage_args:
     engine_args:
       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

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

@@ -11,8 +11,6 @@ stage_args:
     engine_args:
       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
@@ -44,8 +42,6 @@ stage_args:
     engine_args:
       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

docker/Dockerfile.npu

@@ -1,17 +1,17 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.14.0rc1
+ARG VLLM_ASCEND_TAG=v0.17.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 pip install -v -e .
+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
 
 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.14.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.17.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 pip install -v -e .
+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
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

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

@@ -13,9 +13,14 @@ nav:
     - examples/README.md
     - Offline Inference:
       - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
+      - CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
+      - Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
       - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
+      - Helios Video Generation: user_guide/examples/offline_inference/helios.md
+      - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
       - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
       - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
       - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
@@ -23,8 +28,11 @@ nav:
       - Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
       - Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
       - Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
+      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
       - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
+      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
+      - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
       - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
       - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
@@ -92,6 +100,7 @@ nav:
   - api/vllm_omni
 - CLI Reference: cli
 - Community:
+  - Governance: community/governance.md
   - community/*
   - Slack: https://slack.vllm.ai
   - Blog: https://blog.vllm.ai

docs/README.md

@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 
 For more information, checkout the following:
 
-- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)

docs/api/README.md

@@ -6,19 +6,13 @@ 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.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][]
+- [vllm_omni.entrypoints.omni_base.OmniBase][]
 
 ## Inputs
 
@@ -48,9 +42,7 @@ 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/community/governance.md

@@ -0,0 +1,52 @@
+# Governance
+
+vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
+
+## Values
+
+vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
+
+### Design Values
+
+1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
+2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
+3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
+4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
+5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
+
+### Collaboration Values
+
+1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
+2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
+
+## Project Maintainers
+
+### Lead Maintainers
+
+Lead maintainers are responsible for the overall direction and strategy of the project:
+
+- [@Gaohan123](https://github.com/Gaohan123)
+- [@hsliuustc0106](https://github.com/hsliuustc0106)
+- [@ywang96](https://github.com/ywang96)
+
+### Active Committers
+
+Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
+
+- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
+- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
+- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
+- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
+- [@princepride](https://github.com/princepride): Diffusion and Omni Support
+- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
+- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
+- [@wtomin](https://github.com/wtomin):
+- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
+
+## Meetings
+
+Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
+
+## Committer Nomination Process
+
+Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.

docs/configuration/README.md

@@ -14,6 +14,10 @@ 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/pd_disaggregation.md

@@ -0,0 +1,176 @@
+# 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"
+      max_batch_size: 16
+    engine_args:
+      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"
+      max_batch_size: 64
+    engine_args:
+      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"
+      max_batch_size: 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_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
+omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
 ```
 
 For online serving:
@@ -30,7 +30,7 @@ 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 architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0 # mark the unique id for each stage
     runtime: # The disaggregated configuration

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -1,4 +1,4 @@
-# stage config for running qwen2.5-omni with architecture of OmniLLM.
+# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
 stage_args:
   - stage_id: 0
     runtime:

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, `OmniLLM` class, etc.)
+- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, 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,11 +168,6 @@ 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_markers.md

@@ -17,6 +17,7 @@ Defined in `pyproject.toml`:
 | `gpu`              | Tests that run on GPU *                                   |
 | `cuda`             | Tests that run on CUDA *                                  |
 | `rocm`             | Tests that run on AMD/ROCm *                              |
+| `xpu`              | Tests that run on Intel XPU *                             |
 | `npu`              | Tests that run on NPU/Ascend *                            |
 | `H100`             | Tests that require H100 GPU  *                            |
 | `L4`               | Tests that require L4 GPU *                               |
@@ -55,7 +56,7 @@ def test_video_to_audio()
 This decorator is intended to make hardware-aware, cross-platform test authoring easier and more robust for CI/CD environments. The `hardware_test` decorator in `vllm-omni/tests/utils.py` performs the following actions:
 
 1. **Applies platform and resource markers**  
-   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `A2`, `A3`).
+   Adds the appropriate pytest markers for each specified hardware platform (e.g., `cuda`, `rocm`, `xpu`, `npu`) and resource type (e.g., `L4`, `H100`, `MI325`, `B60`, `A2`, `A3`).
    ```
    @pytest.mark.cuda
    @pytest.mark.L4
@@ -80,14 +81,14 @@ This decorator is intended to make hardware-aware, cross-platform test authoring
 - Single call for multiple platforms:
     ```python
     @hardware_test(
-        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
-        num_cards={"cuda": 2, "rocm": 2, "npu": 2},
+        res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
+        num_cards={"cuda": 2, "rocm": 2, "xpu": 2, "npu": 2},
     )
     ```
     or
     ```python
     @hardware_test(
-        res={"cuda": "L4", "rocm": "MI325", "npu": "A2"},
+        res={"cuda": "L4", "rocm": "MI325", "xpu": "B60", "npu": "A2"},
         num_cards=2,
     )
     ```

docs/contributing/ci/tests_style.md

@@ -24,8 +24,6 @@ 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.
@@ -75,11 +73,6 @@ 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)
@@ -170,7 +163,7 @@ vllm_omni/                                    tests/
 
 ### Naming Conventions
 
-- **Unit Tests**: Use `test_<module_name>.py` format. Example: `omni_llm.py` → `test_omni_llm.py`
+- **Unit Tests**: Use `test_<module_name>.py` format. Example: `stage_utils.py` → `test_stage_utils.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

@@ -815,6 +815,107 @@ 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
 
 

docs/contributing/model/adding_omni_model.md

@@ -330,54 +330,46 @@ Stage transitions are the mechanism by which outputs from one stage are converte
 
 ### Where Stage Transitions Are Called
 
-Stage transitions happen automatically in the orchestrator (`OmniLLM` class) during the generation loop. Here's the detailed flow:
+Stage transitions happen automatically in the runtime orchestrator. Here's the detailed flow:
 
-1. **Location**: `vllm_omni/entrypoints/omni_llm.py` in the `_run_generation()` method
+1. **Location**: `vllm_omni/engine/orchestrator.py` in `_forward_to_next_stage()`
 2. **Trigger**: When a stage completes processing and produces outputs
 3. **Execution Flow**:
    ```python
-   # In omni_llm.py, _run_generation() method (around line 345-460)
+   # 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]
 
-   # 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
+   # Save current stage outputs so stage_input_processors can consume them.
+   self.stage_clients[stage_id].set_engine_outputs([output])
 
-       # Store outputs from this stage
-       engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
-       stage.set_engine_outputs(engine_outputs)
+   # THIS IS WHERE STAGE TRANSITION HAPPENS
+   next_inputs = next_client.process_engine_inputs(
+       stage_list=self.stage_clients,
+       prompt=req_state.prompt,
+   )
 
-       # 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)
+   # 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)
    ```
 
 ### How Stage Transitions Work
 
 The stage transition process follows these steps:
 
-1. **Stage Completion**: When a stage finishes processing a request, it stores outputs via `stage.set_engine_outputs(engine_outputs)`
+1. **Stage Completion**: When a stage finishes processing a request, the orchestrator stores outputs via `stage_client.set_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 `process_engine_inputs()` method in `OmniStage` (`omni_stage.py`) handles the transition:
+3. **Input Processing**: The stage input processor configured in stage YAML (under `vllm_omni/model_executor/stage_input_processors/`) handles the transition:
    ```python
    def process_engine_inputs(
        self, stage_list: list[Any], prompt: OmniTokensPrompt | TextPrompt = None

docs/contributing/profiling.md

@@ -23,30 +23,32 @@ 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_llm.start_profile()
+omni.start_profile()
 
 # Only profile Stage 1
-omni_llm.start_profile(stages=[1])
+omni.start_profile(stages=[1])
 ```
 
 ```python
 # Stage 0 (Thinker) and Stage 2 (Audio Decoder) for qwen omni
-omni_llm.start_profile(stages=[0, 2])
+omni.start_profile(stages=[0, 2])
 ```
 
 **Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni import omni_llm
+from vllm_omni.entrypoints.omni import Omni
+
+omni = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
 
 profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
 # 1. Start profiling if enabled
 if profiler_enabled:
-    omni_llm.start_profile(stages=[0])
+    omni.start_profile(stages=[0])
 
 # Initialize generator
-omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+omni_generator = omni.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
 total_requests = len(prompts)
 processed_count = 0
@@ -57,21 +59,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 += len(stage_outputs.request_output)
+    processed_count += 1
 
     # 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_llm.stop_profile()
+        omni.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_llm.close()
+omni.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) and provide the OmniStage abstraction for different AR/DiT stages |
+| **EntryPoints**   | define the APIs for offline/online serving (APIServer, Omni/AsyncOmni), while `AsyncOmniEngine` and `Orchestrator` coordinate multi-stage AR/DiT execution |
 | **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 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.
+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`.
 
 ## 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_lm with HF model name.
+# Create an omni runtime with HF model name.
 from vllm_omni.entrypoints.omni import Omni
 
-omni_lm = Omni(model="Qwen/Qwen3-Omni-30B-A3B-Instruct")
+omni = 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_lm.generate(om_inputs, sampling_params_list)
+outputs = omni.generate(om_inputs, sampling_params_list)
 ```
 
 ## Online Serving

docs/design/feature/omni_connectors/mooncake_store_connector.md

@@ -72,3 +72,241 @@ Parameters:
 
 For more details, refer to the
 [Mooncake repository](https://github.com/kvcache-ai/Mooncake).
+
+---
+
+## Design
+
+### 1. Overview
+
+`MooncakeStoreConnector` is the storage-oriented remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake distributed store APIs and provides a uniform `put()` / `get()` abstraction for transferring stage payloads across nodes.
+
+Compared with `SharedMemoryConnector`, this connector is intended for multi-node deployments. Compared with `MooncakeTransferEngineConnector`, it is simpler and more object-oriented: data is written into a distributed store and fetched back by key, rather than transferred through an explicit remote-memory write protocol.
+
+Its primary role is to provide a general-purpose remote connector for stage payloads when:
+
+- the pipeline spans multiple hosts
+- a shared-memory connector is not possible
+- a simpler remote backend is preferred over the RDMA-oriented transfer engine
+
+### 2. Relationship with the OmniConnector System
+
+`MooncakeStoreConnector` implements `OmniConnectorBase`, so it integrates with the same connector lifecycle as the other backends:
+
+- `OmniConnectorFactory` instantiates it from `ConnectorSpec`
+- `load_omni_transfer_config()` maps edge config to the connector spec
+- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
+
+This means the rest of the pipeline does not need store-specific logic. It only relies on the generic connector contract.
+
+### 3. Design Goals
+
+The connector is designed around the following goals:
+
+- **Cross-node object transfer** using a single key-based abstraction
+- **Minimal connector-specific control plane**, since the store itself is the shared medium
+- **Compatibility with arbitrary Python payloads** through the Omni serializer
+- **Operational simplicity** compared with the transfer-engine-based remote connector
+
+The connector is intentionally not optimized for zero-copy tensor movement or direct remote-memory access.
+
+### 4. Core Design
+
+#### 4.1 Object-Oriented Transfer Model
+
+`MooncakeStoreConnector` treats the transport backend as a distributed object store:
+
+1. Serialize the Python object into bytes.
+2. Generate a unique connector key.
+3. Store the bytes in Mooncake.
+4. Fetch the bytes from Mooncake on the receiver side.
+5. Deserialize the bytes back into the original object.
+
+This design makes the connector conceptually close to a distributed key-value transport.
+
+#### 4.2 Key Construction
+
+The connector uses `OmniConnectorBase._make_key()` to derive the internal store key:
+
+```text
+{request_key}@{from_stage}_{to_stage}
+```
+
+This adds stage routing information to the logical request key and avoids collisions between different pipeline edges.
+
+#### 4.3 No Extra Metadata Path
+
+Unlike `SharedMemoryConnector` and `MooncakeTransferEngineConnector`, this connector returns:
+
+```python
+metadata = None
+```
+
+from `put()`.
+
+This is a meaningful design difference:
+
+- the store itself is the shared rendezvous point
+- consumers only need the same key
+- no transport handle, address, or side-channel metadata needs to be propagated
+
+As a result, the control plane is simpler for this connector than for the SHM and RDMA variants.
+
+### 5. Initialization
+
+#### 5.1 Required Mooncake Components
+
+The connector requires the Mooncake Python bindings to expose:
+
+- `MooncakeDistributedStore`
+- `ReplicateConfig`
+
+If these symbols are unavailable, construction fails immediately with `ImportError`. This keeps startup failures explicit and avoids silent fallback behavior.
+
+#### 5.2 Store Setup
+
+During `_init_store()`, the connector:
+
+1. creates `MooncakeDistributedStore()`
+2. calls `store.setup(...)`
+3. validates the return code
+4. creates a `ReplicateConfig`
+5. enables `with_soft_pin = True`
+
+The setup step is completed during connector construction, not lazily at first use.
+
+### 6. Put / Get Flow
+
+#### 6.1 Producer Flow: `put()`
+
+The producer-side flow is:
+
+1. Validate that the store has been initialized.
+2. Serialize the payload into bytes.
+3. Build the internal key with stage routing information.
+4. Call `store.put(key, serialized_data, self.pin)`.
+5. Update metrics and return success.
+
+The returned tuple is:
+
+```python
+(True, len(serialized_data), None)
+```
+
+#### 6.2 Consumer Flow: `get()`
+
+The consumer-side flow is:
+
+1. Validate that the store has been initialized.
+2. Build the same internal key.
+3. Poll `store.get(key)` for up to 20 retries.
+4. Sleep 50 ms between retries.
+5. If data is found, deserialize it and return `(data, payload_size)`.
+6. If all retries are exhausted, record a timeout and return `None`.
+
+This design gives the connector a bounded waiting model rather than an indefinite blocking get.
+
+### 7. Integration with Stage Communication
+
+All callers use the connector through the same `put()` / `get()` contract:
+
+- the sender calls `put()` to serialize and store the payload
+- the receiver calls `get()` to retrieve and deserialize it
+- no connector-specific metadata is required, since the store key is the rendezvous point
+
+Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, which makes the connector functional but not the highest-performance option for large payloads such as KV cache blocks.
+
+### 8. Data Flow in the Pipeline
+
+The end-to-end transfer model is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant MooncakeStoreConnector
+    participant MooncakeStore
+    participant ReceiverStage
+
+    SenderStage->>MooncakeStoreConnector: put(from_stage, to_stage, put_key, data)
+    MooncakeStoreConnector->>MooncakeStoreConnector: serialize object
+    MooncakeStoreConnector->>MooncakeStore: put(internal_key, bytes)
+    MooncakeStoreConnector-->>SenderStage: (success, size, None)
+
+    ReceiverStage->>MooncakeStoreConnector: get(from_stage, to_stage, get_key)
+    MooncakeStoreConnector->>MooncakeStore: get(internal_key)
+    MooncakeStoreConnector->>MooncakeStoreConnector: deserialize bytes
+    MooncakeStoreConnector-->>ReceiverStage: (data, size)
+```
+
+This is a store-mediated remote transfer model rather than a direct peer-to-peer transport model.
+
+### 9. Strengths and Trade-offs
+
+#### Strengths
+
+- Simple conceptual model: store by key, fetch by key.
+- No connector-specific metadata handoff is required.
+- Works naturally for remote multi-node stage transfer.
+- Easy to integrate into the existing connector abstraction.
+
+#### Trade-offs
+
+- Full serialization and deserialization are always required.
+- Large payloads are more expensive than in direct-memory transports.
+- Runtime behavior depends on external Mooncake services being available.
+- Cleanup semantics are weaker than request-scoped local buffer management.
+
+### 10. Important Implementation Characteristics
+
+#### 10.1 Cleanup Is a No-op
+
+`cleanup()` only logs a debug message and does not actively delete remote data. The current design assumes Mooncake-side lifecycle management rather than explicit per-request removal.
+
+This keeps the connector implementation simple but means request-level reclamation is not modeled inside the connector itself.
+
+#### 10.2 Close Releases the Store Handle
+
+Unlike `SharedMemoryConnector`, where `close()` is a no-op, `MooncakeStoreConnector.close()` performs a meaningful shutdown:
+
+1. Calls `self.store.close()` to release the Mooncake store handle.
+2. Sets `self.store = None` to prevent further operations.
+
+This ensures the connector releases its connection to the external Mooncake service on shutdown. Errors during close are logged but do not propagate.
+
+#### 10.3 Health Output Uses a Shared Schema
+
+`health()` reports:
+
+- `host`
+- `metadata_server`
+- `master`
+- metrics
+
+and also includes placeholder fields such as `protocol`, `pool_device`, and `pool_size` to keep a more uniform shape across connector health outputs. This reflects a system-level consistency choice rather than a store-specific design need.
+
+#### 10.4 Get Is Retry-Based, Not Event-Driven
+
+The receiver side polls the store with a short retry loop. This is simple and robust, but it means the connector is latency-sensitive to:
+
+- store visibility delay
+- network jitter
+- payload size
+
+If the payload is not visible within the retry window, the connector reports a timeout.
+
+### 11. Summary
+
+`MooncakeStoreConnector` is the remote, store-based connector in the OmniConnector stack. Its design is straightforward:
+
+- serialize the payload
+- store it in Mooncake under a stage-qualified key
+- fetch it by the same key on the receiver side
+- deserialize it back into the original object
+
+This connector fills an important gap in the system:
+
+1. It enables cross-node transfer without requiring shared memory.
+2. It keeps the stage communication model uniform.
+3. It provides a simpler operational alternative to the transfer-engine-based remote connector.
+
+Its simplicity is also its main trade-off: it favors a clean object-transfer model over the fast-path and peer-to-peer optimizations implemented by `MooncakeTransferEngineConnector`.

docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md

@@ -205,3 +205,589 @@ for test-specific setup instructions.
 
 For more details on the underlying engine, refer to the
 [Mooncake repository](https://github.com/kvcache-ai/Mooncake).
+
+---
+
+## Design
+
+### 1. Overview
+
+`MooncakeTransferEngineConnector` is the high-performance remote connector in `vllm_omni/distributed/omni_connectors`. It is built on top of Mooncake `TransferEngine` and combines:
+
+- a **direct data plane** for remote memory writes
+- a **ZMQ side channel** for metadata lookup, handshake, and completion signaling
+- a **managed local memory pool** for both send and receive buffers
+
+Unlike `MooncakeStoreConnector`, which treats the backend as a distributed store, `MooncakeTransferEngineConnector` is designed as a peer-to-peer transport. Its goal is to move large stage payloads efficiently while still fitting the common `put()` / `get()` API defined by `OmniConnectorBase`.
+
+It is the most performance-oriented connector in the current OmniConnector family and is intended for large remote payloads such as:
+
+- KV cache transfer
+- stage hidden-state payloads
+- streaming chunk payloads
+- other binary-heavy inter-stage artifacts
+
+### 2. Relationship with the OmniConnector System
+
+`MooncakeTransferEngineConnector` implements the same connector contract as the other backends:
+
+- `put(from_stage, to_stage, put_key, data)`
+- `get(from_stage, to_stage, get_key, metadata=None)`
+- `cleanup(request_id, ...)`
+- `health()`
+- `close()`
+
+It is integrated into the system through the standard connector plumbing:
+
+- `OmniConnectorFactory` constructs the connector from `ConnectorSpec`
+- `load_omni_transfer_config()` resolves the edge-level connector configuration
+- `get_connectors_config_for_stage()` and `resolve_omni_kv_config_for_stage()` inject the connector role
+- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
+
+The key system-level distinction is that this connector is **role-aware**:
+
+- sender instances expose data and listen for pull requests
+- receiver instances allocate buffers and actively pull data from the sender
+
+### 3. Design Goals
+
+The connector is designed around four primary goals:
+
+1. **High-throughput remote transfer**
+   Avoid store-mediated round trips and write directly into the receiver memory region.
+
+2. **Fast path for raw payloads**
+   Support `torch.Tensor`, `bytes`, and `ManagedBuffer` without forcing all traffic through full object serialization.
+
+3. **Unified connector abstraction**
+   Preserve the same `put()` / `get()` interface used by the rest of the OmniConnector stack.
+
+4. **Safe lifecycle management**
+   Manage allocation, reuse, cleanup, and failure recovery for a registered memory pool.
+
+### 4. Architecture Overview
+
+At a high level, the connector is composed of four main subsystems:
+
+```mermaid
+classDiagram
+    class OmniConnectorBase {
+        <<abstract>>
+        +put(from_stage, to_stage, put_key, data)
+        +get(from_stage, to_stage, get_key, metadata)
+        +cleanup(request_id)
+        +health()
+        +close()
+    }
+
+    class MooncakeTransferEngineConnector {
+        +supports_raw_data: bool
+        -engine: TransferEngine
+        -allocator: BufferAllocator
+        -pool: torch.Tensor
+        -zmq_ctx: zmq.Context
+        -_local_buffers: dict
+        -_sender_executor: ThreadPoolExecutor
+        -_listener_thread: threading.Thread
+        +put(...)
+        +get(...)
+        +update_sender_info(sender_host, sender_zmq_port)
+        +get_connection_info()
+        +cleanup(request_id, from_stage, to_stage)
+        +close()
+    }
+
+    class BufferAllocator {
+        -total_size: int
+        -alignment: int
+        -free_blocks: list
+        +alloc(size) int
+        +free(offset, size)
+    }
+
+    class ManagedBuffer {
+        -allocator: BufferAllocator
+        -offset: int
+        -size: int
+        -pool_tensor: torch.Tensor
+        +tensor
+        +as_tensor(dtype, shape) torch.Tensor
+        +to_bytes() bytes
+        +release()
+    }
+
+    class TransferEngine {
+        +initialize(host, handshake, protocol, device_name)
+        +register_memory(base_ptr, size)
+        +batch_transfer_sync_write(remote_session, src_addrs, dst_addrs, lengths)
+        +unregister_memory(base_ptr)
+        +get_rpc_port() int
+    }
+
+    class QueryRequest {
+        +request_id: str
+    }
+
+    class QueryResponse {
+        +request_id: str
+        +data_size: int
+        +is_fast_path: bool
+    }
+
+    class MooncakeAgentMetadata {
+        +remote_hostname: str
+        +remote_port: int
+        +request_id: str
+        +dst_addrs: list[int]
+        +lengths: list[int]
+    }
+
+    OmniConnectorBase <|-- MooncakeTransferEngineConnector
+    MooncakeTransferEngineConnector *-- BufferAllocator
+    MooncakeTransferEngineConnector *-- TransferEngine
+    ManagedBuffer --> BufferAllocator : releases to
+    ManagedBuffer --> "1" torch.Tensor : views
+    MooncakeTransferEngineConnector ..> ManagedBuffer : returns / retains
+    MooncakeTransferEngineConnector ..> QueryRequest : decodes
+    MooncakeTransferEngineConnector ..> QueryResponse : encodes
+    MooncakeTransferEngineConnector ..> MooncakeAgentMetadata : exchanges
+```
+
+#### 4.1 Transfer Engine
+
+Mooncake `TransferEngine` is responsible for the actual data-plane transfer. It registers local memory and performs synchronous remote writes through:
+
+```python
+batch_transfer_sync_write(...)
+```
+
+#### 4.2 Managed Memory Pool
+
+Each connector instance owns a large pre-registered memory pool:
+
+- CPU pinned memory when `memory_pool_device == "cpu"`
+- GPU memory when `memory_pool_device == "cuda"`
+
+This avoids repeated memory registration and allows each transfer to allocate subranges from one long-lived pool.
+
+#### 4.3 Buffer Manager
+
+Two helper classes control local memory ownership:
+
+- `BufferAllocator`
+  Manages aligned subrange allocation and free-list merging.
+
+- `ManagedBuffer`
+  Represents one live slice of the pool and exposes:
+  - `.tensor`
+  - `.as_tensor(dtype, shape)`
+  - `.to_bytes()`
+  - `.release()`
+
+#### 4.4 ZMQ Side Channel
+
+ZMQ is used for transport coordination, not for the data payload itself. It handles:
+
+- metadata query from receiver to sender
+- pull request submission
+- completion or error signaling
+- internal notification from worker threads back to the listener thread
+
+This split makes the control plane lightweight while keeping the bulk payload on the transfer engine data plane.
+
+### 5. Role Model
+
+#### 5.1 Sender Role
+
+A sender connector:
+
+- accepts `put()` calls
+- stores live transfer-ready buffers in `_local_buffers`
+- starts a ZMQ listener thread
+- responds to metadata queries and pull requests from receivers
+
+#### 5.2 Receiver Role
+
+A receiver connector:
+
+- does not bind the sender-side ZMQ listener
+- accepts `get()` calls
+- allocates receive buffers from its own pool
+- requests metadata or transfer service from the sender
+
+The role is not inferred dynamically. It is injected by the stage configuration layer:
+
+- incoming edge for a stage -> `role="receiver"`
+- outgoing edge for a stage -> `role="sender"`
+
+This is important because incorrect role assignment would break initialization semantics.
+
+#### 5.3 Host Auto-Detection
+
+The `host` configuration field supports the special value `"auto"`. When set, the connector auto-detects the local IP address that would be used for external communication (via a UDP socket probe to `8.8.8.8`). If that fails, it falls back to hostname resolution, and ultimately to `127.0.0.1`.
+
+This is useful in environments where the operator does not want to hard-code IP addresses in the connector config.
+
+#### 5.4 RDMA Device Filtering
+
+The `device_name` configuration field allows the operator to specify which RDMA NICs to use (comma-separated, e.g. `"mlx5_0,mlx5_1"`). If not set in config, the connector also checks the `RDMA_DEVICE_NAME` environment variable.
+
+This is important in environments with mixed InfiniBand/RoCE NICs, where not all devices are suitable for the transfer engine.
+
+### 6. Local Memory Management
+
+#### 6.1 Memory Pool Registration
+
+During initialization, the connector:
+
+1. allocates a large pool tensor
+2. records its base pointer
+3. registers that memory with Mooncake
+4. creates a `BufferAllocator` for subrange management
+
+This means every later transfer only allocates offsets inside the pre-registered pool rather than registering memory per request.
+
+#### 6.2 BufferAllocator
+
+`BufferAllocator` maintains a sorted free list of `(offset, size)` blocks and enforces alignment. Its responsibilities include:
+
+- aligned allocation
+- freeing previously allocated blocks
+- adjacent block merging
+- double-free detection
+- overlap detection to catch corruption
+
+This is a critical piece of the connector because both sender and receiver depend on long-lived pool reuse.
+
+#### 6.3 ManagedBuffer
+
+`ManagedBuffer` is the main fast-path data wrapper. It can:
+
+- expose the pool slice as a zero-copy 1D `uint8` tensor
+- reinterpret that slice as a typed tensor
+- copy out the contents as Python `bytes`
+- release the slice back to the allocator
+
+The connector uses `ManagedBuffer` in two different ways:
+
+- as a send-side holder to keep the pool slice alive
+- as a receive-side return type when `is_fast_path=True`
+
+### 7. Put Flow
+
+#### 7.1 High-Level Behavior
+
+`put()` is only valid in sender mode. Its job is to expose a payload for later remote pull by the receiver.
+
+The high-level flow is:
+
+1. validate connector state and role
+2. convert the input into a pool-backed transferable representation
+3. store the transfer metadata in `_local_buffers`
+4. return lightweight metadata describing how the receiver can fetch the data
+
+#### 7.2 Payload Type Handling
+
+`put()` supports three payload classes:
+
+**A. `ManagedBuffer`**
+
+If the buffer belongs to the same pool, the connector can use it directly without copying. This is the most efficient path.
+
+If the buffer comes from a different pool, the connector falls back to a copy path.
+
+**B. `torch.Tensor` or `bytes`**
+
+These payloads are copied into the local pool and marked as fast-path data:
+
+- no Omni object serialization is required
+- receiver can get a `ManagedBuffer` back
+
+**C. Generic Python object**
+
+Any other payload is serialized via `OmniSerializer.serialize(...)` and then copied into the pool.
+
+In this case:
+
+- `is_fast_path=False`
+- the receiver will deserialize back into a Python object
+
+#### 7.3 Sender Metadata
+
+The sender returns:
+
+```python
+{
+    "source_host": self.host,
+    "source_port": self.zmq_port,
+    "data_size": size,
+    "is_fast_path": is_fast_path,
+}
+```
+
+This metadata is intentionally lightweight. It tells the receiver:
+
+- where the sender-side control plane lives
+- how large the remote transfer will be
+- whether the payload should be returned as a `ManagedBuffer` or a deserialized object
+
+#### 7.4 Sender Buffer Table
+
+The sender stores each live payload in `_local_buffers` under the stage-qualified key. Each entry contains:
+
+- source addresses
+- lengths
+- the holder object
+- ownership information (`should_release`)
+- `is_fast_path`
+- creation time
+
+This table is the sender-side truth source for both metadata queries and pull requests.
+
+### 8. Get Flow
+
+#### 8.1 High-Level Behavior
+
+`get()` runs on the receiver side and performs four steps:
+
+1. resolve metadata
+2. allocate a destination buffer in the local pool
+3. request the sender to write into that destination buffer
+4. return either a `ManagedBuffer` or a deserialized object
+
+#### 8.2 Metadata Resolution Paths
+
+The `metadata` parameter in `get()` is optional. The connector supports two resolution modes depending on whether the caller supplies it.
+
+**With metadata**
+
+When the caller passes metadata, the connector uses it directly. The metadata carries:
+
+- `source_host` / `source_port` — sender ZMQ endpoint
+- `data_size` — payload byte count
+- `is_fast_path` — whether the receiver gets a `ManagedBuffer` or a deserialized object
+
+This mode is suitable when the control plane already forwards the sender's `put()` output to the receiver.
+
+**Without metadata**
+
+When `get(metadata=None)` is called, the connector queries the sender over ZMQ to discover the same fields (`data_size`, `is_fast_path`). The caller must first call:
+
+```python
+update_sender_info(sender_host, sender_zmq_port)
+```
+
+so that the connector knows where to send the query.
+
+This mode is suitable for polling-based flows (e.g. KV transfer, async chunk transfer) where the receiver does not have metadata from the control plane.
+
+#### 8.3 Destination Allocation
+
+Once metadata is resolved, the receiver:
+
+1. allocates a subrange from its own local pool
+2. wraps it in a `ManagedBuffer`
+3. builds a `MooncakeAgentMetadata` request containing:
+   - receiver hostname
+   - receiver RPC port
+   - request ID
+   - destination addresses
+   - transfer lengths
+
+This tells the sender exactly where to write the incoming data.
+
+#### 8.4 Transfer Completion
+
+The receiver then sends the pull request over ZMQ and waits for:
+
+- `TRANS_DONE`
+- or `TRANS_ERROR`
+
+If the transfer succeeds:
+
+- for `is_fast_path=True`, the receiver returns `(ManagedBuffer, size)`
+- for `is_fast_path=False`, the receiver copies to bytes, deserializes, releases the buffer, and returns `(object, size)`
+
+### 9. Sender-Side Listener Design
+
+#### 9.1 Listener Thread
+
+In sender mode, the connector starts `_zmq_listener_loop()` after initialization. The listener:
+
+- binds `tcp://{host}:{zmq_port}`
+- receives incoming requests
+- uses a poller for socket events and internal notifications
+- periodically reclaims stale buffers
+
+If the bind fails, initialization fails immediately. The code does not silently downgrade the connector role.
+
+#### 9.2 Worker Thread Pool
+
+The listener hands work to `_sender_executor` so that the listener thread does not block on transfer work.
+
+There are two request types:
+
+- metadata query -> `_handle_query_request(...)`
+- transfer request -> `_handle_pull_request(...)`
+
+#### 9.3 Query Handling
+
+For metadata queries, the sender looks up the request ID in `_local_buffers` and returns:
+
+- data size
+- fast-path flag
+
+This supports consumers that only know the sender endpoint but not the original sender metadata.
+
+#### 9.4 Pull Handling
+
+For a transfer request, the sender:
+
+1. locates the source addresses in `_local_buffers`
+2. constructs the remote session identifier
+3. calls `batch_transfer_sync_write(...)`
+4. replies `TRANS_DONE` or `TRANS_ERROR`
+
+On success, the sender immediately calls `cleanup(meta.request_id)` and frees the producer-side buffer if it owns it.
+
+This choice is important: it makes the connector effectively a **single-consumer transfer model** for each successful put/get pair.
+
+### 10. Fast Path Semantics
+
+This connector explicitly advertises:
+
+```python
+supports_raw_data = True
+```
+
+That means it can move raw payloads without forcing everything through the Omni object serializer.
+
+#### Fast Path
+
+For `torch.Tensor`, `bytes`, or pool-local `ManagedBuffer`:
+
+- sender returns `is_fast_path=True`
+- receiver returns a `ManagedBuffer`
+- caller is responsible for calling `release()`
+
+This avoids an unnecessary copy on the receiver side.
+
+#### Serialized Object Path
+
+For arbitrary Python objects:
+
+- sender serializes the payload
+- receiver converts the receive buffer to bytes
+- receiver deserializes the object
+- receive buffer is released internally
+
+This preserves a uniform object-oriented API while still allowing optimized raw-data transport when possible.
+
+### 11. Failure Handling and Cleanup
+
+#### 11.1 Timeouts and Socket Recovery
+
+The receiver caches ZMQ REQ sockets per thread, but invalidates them after failures. This avoids reusing sockets that may be stuck in a bad state after timeout or receive errors.
+
+Timeout is scaled based on payload size:
+
+- a base timeout
+- plus additional time for large payloads
+
+This is intended to reduce false timeouts for large remote writes.
+
+#### 11.2 Stale Buffer Reclamation
+
+The sender periodically reclaims old entries from `_local_buffers` using a TTL policy. This protects the memory pool from permanent leaks if a receiver crashes or never consumes a prepared payload.
+
+This is a practical recovery mechanism, although the code notes that TTL cleanup can still race with very long-running in-flight transfers.
+
+#### 11.3 Connector Shutdown
+
+`close()` is a full resource teardown routine. It:
+
+- stops the listener thread
+- shuts down the worker executor
+- releases all pending buffers
+- closes cached sockets
+- unregisters memory from the engine when supported
+- terminates the ZMQ context
+- drops the pool reference
+
+This makes `MooncakeTransferEngineConnector` the most lifecycle-aware connector in the current connector family.
+
+### 12. Current Implementation Constraints
+
+The current code documents several important topology constraints.
+
+#### 12.1 One Sender to One Receiver per Successful Transfer
+
+After a successful transfer, the sender-side buffer is cleaned up immediately. This means the same prepared payload is not retained for multiple independent receivers.
+
+#### 12.2 One Receiver to One Active Sender Endpoint
+
+The receiver only stores one `(sender_host, sender_zmq_port)` pair through `update_sender_info(...)`. So the metadata-query mode is currently single-sender at a time.
+
+#### 12.3 Explicit Buffer Ownership Matters
+
+When the connector allocates a pool slice internally, it is responsible for releasing it. When a caller passes an externally owned `ManagedBuffer`, the connector keeps it alive for transfer but does not assume ownership of its eventual release.
+
+These constraints are consistent with the current implementation and should be treated as design assumptions rather than incidental behavior.
+
+### 13. Data Flow in the Pipeline
+
+The end-to-end sender/receiver interaction is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant SenderConnector
+    participant ReceiverConnector
+    participant ReceiverStage
+
+    SenderStage->>SenderConnector: put(from_stage, to_stage, put_key, data)
+    SenderConnector->>SenderConnector: place payload in local memory pool
+    SenderConnector-->>SenderStage: metadata(source_host, source_port, data_size, is_fast_path)
+
+    ReceiverStage->>ReceiverConnector: get(..., metadata)
+    ReceiverConnector->>ReceiverConnector: allocate destination buffer
+    ReceiverConnector->>SenderConnector: ZMQ pull request with dst addr
+    SenderConnector->>ReceiverConnector: TransferEngine remote write
+    SenderConnector-->>ReceiverConnector: TRANS_DONE
+    ReceiverConnector-->>ReceiverStage: ManagedBuffer or deserialized object
+```
+
+For metadata-less polling, the flow simply adds a metadata query step before the pull request.
+
+### 14. Strengths and Trade-offs
+
+#### Strengths
+
+- Best remote-transfer design in the current connector stack for large payloads.
+- Supports raw-data fast path.
+- Keeps stage communication under the same connector abstraction.
+- Includes real lifecycle and memory-pool management.
+- Works for both stage payload transfer and KV transfer scenarios.
+
+#### Trade-offs
+
+- More complex than the store-based connector.
+- Correctness depends on role injection and endpoint coordination.
+- Caller must release fast-path receive buffers.
+- Current implementation is optimized for single-consumer transfer semantics.
+
+### 15. Summary
+
+`MooncakeTransferEngineConnector` is the high-performance peer-to-peer transport in the OmniConnector system. Its design combines:
+
+- a registered memory pool
+- a safe subrange allocator
+- a ZMQ control plane
+- a Mooncake transfer-engine data plane
+
+This allows the connector to support both:
+
+1. a **fast path** for raw tensors and bytes
+2. a **generic object path** for arbitrary Python payloads
+
+Within vLLM-Omni, it is the connector that most directly targets performance-sensitive remote transfer, especially for large payloads and KV cache movement. Its additional complexity is deliberate: it is the connector that turns the generic OmniConnector abstraction into a transport capable of efficient remote memory movement rather than simple object storage.

docs/design/feature/omni_connectors/shared_memory_connector.md

@@ -7,9 +7,10 @@ auto-configured when no explicit connector is specified for an edge.
 
 ## How It Works
 
-- Small payloads (< threshold): serialized and passed inline in metadata.
-- Large payloads (>= threshold): stored in shared memory; the SHM name is
-  returned in metadata.
+All payloads are serialized and stored in shared memory (`/dev/shm`); the SHM
+segment name is returned in metadata. The configuration exposes a
+`shm_threshold_bytes` field for a future inline-vs-SHM split, but the current
+implementation always uses shared memory regardless of payload size.
 
 ## Configuration
 
@@ -25,3 +26,234 @@ runtime:
 ## Notes
 
 - Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.
+
+---
+
+## Design
+
+### 1. Overview
+
+`SharedMemoryConnector` is the default same-node connector in `vllm_omni/distributed/omni_connectors`. It is designed for stage-to-stage transfer when producer and consumer processes run on the same host and can share `/dev/shm`.
+
+The connector provides a unified `put()` / `get()` API for arbitrary Python objects while keeping the control plane lightweight:
+
+- The payload is serialized by the connector.
+- The serialized bytes are placed in shared memory.
+- The queue/control plane only carries a small metadata handle.
+
+This makes `SharedMemoryConnector` the simplest connector in the OmniConnector family and the default fallback when an edge does not explicitly configure another backend.
+
+### 2. Relationship with the OmniConnector System
+
+`SharedMemoryConnector` implements `OmniConnectorBase`, so it follows the same lifecycle and API contract as the other connectors:
+
+- `put(from_stage, to_stage, put_key, data)`
+- `get(from_stage, to_stage, get_key, metadata=None)`
+- `cleanup(request_id)`
+- `health()`
+- `close()`
+
+Within the larger system:
+
+- `load_omni_transfer_config()` automatically fills missing edges with `SharedMemoryConnector`.
+- Callers interact with the connector exclusively through the `put()` / `get()` / `cleanup()` contract — the connector does not require caller-specific logic.
+
+Compared with the remote Mooncake-based connectors, `SharedMemoryConnector` is intentionally minimal and local-only.
+
+### 3. Design Goals
+
+The connector is built around the following goals:
+
+- **Low-friction local transfer** for single-node multi-process pipelines.
+- **Unified object semantics** for arbitrary Python payloads.
+- **Small control-plane overhead** by passing only metadata through queues.
+- **Zero external dependencies** beyond Python shared memory and the existing stage utilities.
+
+It is not intended to provide cross-node transfer, RDMA, or raw tensor zero-copy semantics across processes.
+
+### 4. Core Design
+
+#### 4.1 Serialization Model
+
+`SharedMemoryConnector` always starts from a Python object and serializes it through the shared Omni serializer:
+
+```python
+payload = self.serialize_obj(data)
+```
+
+This keeps the connector behavior consistent with the rest of the connector stack:
+
+- producer code does not need connector-specific serialization logic
+- consumer code always receives the original object after deserialization
+- the connector can reuse the same serializer used by other backends
+
+#### 4.2 Shared Memory as the Data Plane
+
+The actual data plane is a shared-memory segment created by:
+
+- `shm_write_bytes(...)`
+- `shm_read_bytes(...)`
+
+The connector stores a small metadata object such as:
+
+```python
+{
+    "shm": {"name": ..., "size": ...},
+    "size": ...
+}
+```
+
+This metadata is passed over the control plane and allows the downstream stage to locate the shared-memory segment.
+
+#### 4.3 Locking Model
+
+To avoid races between the producer and consumer, the connector uses a lock file per request:
+
+```text
+/dev/shm/shm_{put_key}_lockfile.lock
+```
+
+Locking is done with `fcntl.flock`:
+
+- producer uses `LOCK_EX`
+- consumer uses `LOCK_EX`
+
+Both sides acquire an exclusive lock. This ensures that the shared-memory segment is not read while it is still being written and makes the handoff safer in a multi-process environment.
+
+### 5. Put / Get Flow
+
+#### 5.1 Producer Flow: `put()`
+
+The producer-side flow is:
+
+1. Serialize the input object to bytes.
+2. Compute the payload size.
+3. Acquire the per-request lock file.
+4. Write the bytes into shared memory.
+5. Return lightweight metadata to the caller.
+
+The returned tuple is:
+
+```python
+(success, serialized_size, metadata)
+```
+
+where `metadata` contains the shared-memory handle needed by the consumer.
+
+#### 5.2 Consumer Flow: `get(metadata=...)`
+
+The primary consumer path is metadata-driven:
+
+1. Extract the shared-memory handle from `metadata`.
+2. Acquire the exclusive lock.
+3. Read the raw bytes from shared memory.
+4. Deserialize the bytes back into the original Python object.
+5. Remove the lock file if it still exists.
+
+This is the path used by the current stage-to-stage connector flow.
+
+#### 5.3 Compatibility Flow: `get(metadata=None)`
+
+The connector also keeps a compatibility path for callers that only know the key:
+
+1. Attempt to open the shared-memory segment by name via `SharedMemory(name=get_key)`.
+2. If the segment exists and has non-zero size, acquire the exclusive lock and read the bytes.
+3. Deserialize the bytes and return the object.
+
+If the segment does not exist or any exception occurs, the call returns `None` immediately. There is no retry loop in this path -- it is a single-attempt open.
+
+This path is mainly for older code paths and is not the preferred mode for the current connector pipeline.
+
+### 6. Key Implementation Characteristics
+
+#### 6.1 Threshold Exists, but the Current Code Always Uses SHM
+
+The class keeps a `shm_threshold_bytes` field and still exposes metrics for inline writes. However, the current implementation uses:
+
+```python
+if True:
+    ...
+```
+
+inside `put()`, which means the current code path always writes to shared memory.
+
+So the design still suggests a future split between:
+
+- small payloads inline
+- large payloads in shared memory
+
+but the current behavior is effectively:
+
+- all payloads go through shared memory
+
+This should be documented because it affects real runtime behavior.
+
+#### 6.2 Cleanup Is Currently Passive
+
+`cleanup()` is currently a no-op. The intended assumption is:
+
+- the consumer reads the segment
+- the underlying shared-memory helpers unlink it
+
+If the consumer never executes `get()`, the shared-memory segment may remain allocated. This means the connector relies on the normal success path for resource reclamation.
+
+#### 6.3 Close Is Currently Minimal
+
+`close()` is also a no-op. There is no connector-owned background thread, socket, or memory pool to tear down, so the lifecycle is simple. The trade-off is that `close()` does not scan or recover leaked shared-memory resources.
+
+### 7. Data Flow in the Pipeline
+
+The typical flow with `SharedMemoryConnector` is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant SharedMemoryConnector
+    participant QueueOrControlPlane
+    participant ReceiverStage
+
+    SenderStage->>SharedMemoryConnector: put(from_stage, to_stage, put_key, data)
+    SharedMemoryConnector->>SharedMemoryConnector: serialize object
+    SharedMemoryConnector->>SharedMemoryConnector: write bytes to /dev/shm
+    SharedMemoryConnector-->>SenderStage: metadata {shm, size}
+    SenderStage->>QueueOrControlPlane: forward connector metadata
+    QueueOrControlPlane->>ReceiverStage: task + connector metadata
+    ReceiverStage->>SharedMemoryConnector: get(from_stage, to_stage, get_key, metadata)
+    SharedMemoryConnector->>SharedMemoryConnector: read bytes from /dev/shm
+    SharedMemoryConnector->>SharedMemoryConnector: deserialize object
+    SharedMemoryConnector-->>ReceiverStage: (data, size)
+```
+
+This is a classic split-control-plane / data-plane design, but constrained to a single host.
+
+### 8. Strengths and Trade-offs
+
+#### Strengths
+
+- Very simple deployment model.
+- No external service dependency.
+- Fits naturally into the existing queue-driven orchestration flow.
+- Good default for local multi-process pipelines.
+
+#### Trade-offs
+
+- Same-node only.
+- Full object serialization and deserialization are still required.
+- Resource cleanup depends on the normal consumer path.
+- Shared memory capacity is limited by host configuration.
+
+### 9. Summary
+
+`SharedMemoryConnector` is the baseline local transport for the OmniConnector system. Its design is intentionally straightforward:
+
+- serialize object
+- place bytes in shared memory
+- pass metadata through the control plane
+- deserialize on the receiving side
+
+It plays two important roles in vLLM-Omni:
+
+1. It is the simplest production-ready connector for same-node stage pipelines.
+2. It serves as the automatic fallback connector when no explicit edge transport is configured.
+
+Although the current implementation is deliberately minimal, it provides the foundation for reliable local connector semantics and keeps the stage communication model uniform across the system.

docs/design/feature/omni_connectors/yuanrong_connector.md

@@ -94,8 +94,265 @@ stage_args:
 Parameters:
 
 - host: datasystem worker host.
-- port: datasystem worker port.
+- port: datasystem worker port (default: `35001` if omitted; the example above uses `31501` to match the worker startup command).
 - get_sub_timeout_ms: get timeout in milliseconds (0 for no timeout).
 
 For more details, refer to the
 [Yuanrong Datasystem repository](https://atomgit.com/openeuler/yuanrong-datasystem).
+
+---
+
+## Design
+
+### 1. Overview
+
+`YuanrongConnector` is the Datasystem-based remote connector in `vllm_omni/distributed/omni_connectors`. It uses Yuanrong Datasystem's distributed key-value client as the transport backend and exposes the same `put()` / `get()` interface as the other OmniConnectors.
+
+Like `MooncakeStoreConnector`, it is a store-oriented remote connector rather than a direct peer-to-peer transport. Its role is to let stage payloads move across nodes through a deterministic key-based storage abstraction while keeping the rest of the pipeline on the common connector API.
+
+It is intended for deployments that already use Yuanrong Datasystem and want a remote connector that integrates with the existing OmniConnector configuration and orchestration model.
+
+### 2. Relationship with the OmniConnector System
+
+`YuanrongConnector` implements `OmniConnectorBase`, so it participates in the same connector lifecycle as the other backends:
+
+- `OmniConnectorFactory` constructs it from a `ConnectorSpec`
+- stage edge configuration is resolved by `load_omni_transfer_config()`
+- All callers (batch forwarding, chunk transfer, KV transfer, etc.) interact with it through the same `put()` / `get()` contract
+
+This means the connector is not exposed directly to stage logic. Stages only interact with the generic connector contract, and the backend choice remains a configuration concern.
+
+### 3. Design Goals
+
+The connector is built around the following goals:
+
+1. **Cross-node payload transfer through Datasystem**
+   Reuse Yuanrong Datasystem as the remote exchange medium for stage data.
+
+2. **Uniform object transfer semantics**
+   Allow arbitrary Python objects to be transmitted through the shared Omni serializer.
+
+3. **Minimal connector-specific control plane**
+   Use deterministic keys so that consumers can retrieve data without an extra transport metadata handoff.
+
+4. **Operational reuse of existing infrastructure**
+   Fit into environments that already deploy Yuanrong Datasystem workers and etcd.
+
+The connector is not designed for direct remote-memory writes or tensor-specific fast-path transfer.
+
+### 4. Core Design
+
+#### 4.1 Store-Oriented Transfer Model
+
+`YuanrongConnector` treats the transport backend as a distributed object store:
+
+1. serialize the Python payload
+2. build a deterministic connector key
+3. write the serialized bytes into Datasystem
+4. read the bytes back on the receiver side
+5. deserialize them into the original object
+
+This is the same broad architectural class as `MooncakeStoreConnector`, but implemented on top of Yuanrong Datasystem APIs instead of Mooncake store APIs.
+
+#### 4.2 Deterministic Keying
+
+Unlike the default `_make_key()` in `OmniConnectorBase`, `YuanrongConnector` defines its own key format:
+
+```text
+{request_id}:{from_stage}_{to_stage}
+```
+
+This has two design implications:
+
+- the request identifier remains the primary lookup handle
+- stage routing information is embedded in the key so that the same logical request ID can safely appear on different edges
+
+The explicit override also makes the key format easier to align with Datasystem-side debugging and operational inspection.
+
+#### 4.3 No Extra Metadata Hand-off
+
+`put()` returns:
+
+```python
+(success, serialized_size, None)
+```
+
+and does not generate connector-specific metadata.
+
+This design works because the receiver can reconstruct the exact same key from:
+
+- `get_key`
+- `from_stage`
+- `to_stage`
+
+As a result, the connector does not require a separate side-channel metadata handoff.
+
+### 5. Initialization
+
+#### 5.1 Datasystem Client Dependency
+
+The connector requires the Datasystem Python bindings to expose:
+
+- `KVClient`
+- `SetParam`
+- `WriteMode`
+
+If any of these symbols are unavailable, connector construction fails immediately with `ImportError`. This keeps configuration errors explicit and avoids a partially initialized runtime.
+
+#### 5.2 Client Setup
+
+During `_init_client()`, the connector:
+
+1. reads `host` and `port`
+2. creates `KVClient(host, port)`
+3. calls `client.init()`
+
+At construction time it also creates a `SetParam` and fixes:
+
+```python
+self.set_param.write_mode = WriteMode.NONE_L2_CACHE_EVICT
+```
+
+This means the connector has a stable write policy for all writes and does not currently expose write-mode selection as a higher-level connector option.
+
+### 6. Put / Get Flow
+
+#### 6.1 Producer Flow: `put()`
+
+The producer-side flow is:
+
+1. verify that the Datasystem client has been initialized
+2. serialize the input object with the shared Omni serializer
+3. build the Datasystem key using the connector-specific `_make_key()`
+4. call `client.set(key, serialized_data, self.set_param.write_mode)`
+5. update metrics and return success
+
+The returned metadata is always `None`, because the Datasystem key itself is the lookup contract between producer and consumer.
+
+#### 6.2 Consumer Flow: `get()`
+
+The consumer-side flow is:
+
+1. verify that the Datasystem client has been initialized
+2. rebuild the same key with `from_stage`, `to_stage`, and `get_key`
+3. call:
+
+```python
+client.get([key], False, self.get_sub_timeout_ms)
+```
+
+4. take the first returned element if present
+5. deserialize it and return `(data, payload_size)`
+
+If the returned list is empty or contains no data for the key, `get()` returns `None`.
+
+### 7. Timeout and Retrieval Semantics
+
+The connector uses `get_sub_timeout_ms` as its read timeout. Unlike `MooncakeStoreConnector`, which performs an explicit retry loop in Python, `YuanrongConnector` delegates waiting behavior more directly to the Datasystem client call.
+
+This leads to a slightly different retrieval model:
+
+- `MooncakeStoreConnector`: retry-oriented polling in connector code
+- `YuanrongConnector`: single client call with backend-managed timeout behavior
+
+From the connector API perspective the result is the same, but operationally the waiting behavior is more dependent on Datasystem client semantics.
+
+### 8. Integration with Stage Communication
+
+All callers use the connector through the same `put()` / `get()` contract:
+
+- the sender calls `put()` to serialize and store the payload
+- the receiver calls `get()` to retrieve and deserialize it
+- no connector-specific metadata is required, since the Datasystem key is the rendezvous point
+
+Because `put()` returns `metadata=None`, the connector is naturally compatible with callers that do not forward metadata (e.g. polling-based flows). The trade-off is that all payloads incur full serialization and deserialization costs, and there is no raw tensor fast path, which makes the connector functional but not optimized for the largest payloads.
+
+### 9. Data Flow in the Pipeline
+
+The end-to-end transfer model is:
+
+```mermaid
+sequenceDiagram
+    participant SenderStage
+    participant YuanrongConnector
+    participant Datasystem
+    participant ReceiverStage
+
+    SenderStage->>YuanrongConnector: put(from_stage, to_stage, put_key, data)
+    YuanrongConnector->>YuanrongConnector: serialize object
+    YuanrongConnector->>Datasystem: set(key, bytes)
+    YuanrongConnector-->>SenderStage: (success, size, None)
+
+    ReceiverStage->>YuanrongConnector: get(from_stage, to_stage, get_key)
+    YuanrongConnector->>Datasystem: get([key], timeout)
+    YuanrongConnector->>YuanrongConnector: deserialize bytes
+    YuanrongConnector-->>ReceiverStage: (data, size)
+```
+
+This is a store-mediated remote connector design with deterministic key lookup and no explicit side-channel metadata exchange.
+
+### 10. Strengths and Trade-offs
+
+#### Strengths
+
+- Reuses existing Yuanrong Datasystem infrastructure.
+- Keeps the connector contract simple and uniform.
+- Requires no connector-specific metadata handoff.
+- Suitable for remote stage transfer in Datasystem-based deployments.
+
+#### Trade-offs
+
+- Always pays full serialization and deserialization cost.
+- Does not support raw bytes / tensor fast-path semantics.
+- Depends on external Datasystem worker availability.
+- Retrieval and timeout behavior are partly delegated to the backend client.
+
+### 11. Important Implementation Characteristics
+
+#### 11.1 Cleanup Is a No-op
+
+`cleanup()` only logs a debug message and does not explicitly remove data from Datasystem. The current design assumes backend-side lifecycle or garbage collection rather than request-scoped delete semantics inside the connector.
+
+This mirrors the same design trade-off seen in other store-based connectors: simplicity over explicit per-request data reclamation.
+
+#### 11.2 Close Only Releases the Local Client Handle
+
+`close()` does not perform a remote shutdown. It simply clears the local `client` reference and marks the connector as closed from the process perspective:
+
+```python
+self.client = None
+```
+
+This is appropriate for a client-based store connector, but it also means that backend resource lifecycle remains outside the connector's control.
+
+#### 11.3 Error and Timeout Metrics Are Coarse-Grained
+
+The connector tracks:
+
+- `puts`
+- `gets`
+- `bytes_transferred`
+- `errors`
+- `timeouts`
+
+In the current code, failed `put()` increments `errors`, while `get()` exceptions increment `timeouts`. This is operationally useful, but it does not distinguish between:
+
+- backend timeout
+- not-found result
+- transport failure
+- deserialization failure
+
+So the metrics should be read as high-level health indicators, not detailed root-cause diagnostics.
+
+### 12. Summary
+
+`YuanrongConnector` is the Datasystem-backed remote connector in the OmniConnector family. Its design is straightforward:
+
+- serialize payloads
+- store them under a deterministic stage-qualified key
+- retrieve them by the same key
+- deserialize them on the receiving side
+
+Within vLLM-Omni, it provides a clean integration point for Yuanrong Datasystem-based deployments while preserving the same connector abstraction used by the rest of the pipeline.
+
+Its design priorities are simplicity, infrastructure reuse, and API consistency rather than direct-memory transport optimization.

docs/design/feature/sequence_parallel.md

@@ -55,6 +55,8 @@ The `_sp_plan` mechanism allows SP **without modifying `forward()` logic**. The
 - Tensor operations happen at `nn.Module` boundaries
 - Predictable sharding/gathering patterns
 
+This is the ideal approach for integrating sequence parallelism into new models, as it is easier to maintain and ensure compatibility with other types of acceleration.
+
 **How it works:**
 1. Declare `_sp_plan` dict in your transformer class
 2. Framework automatically applies hooks when `sequence_parallel_size > 1`
@@ -201,6 +203,36 @@ class TransformerWithRoPE(nn.Module):
     }
 ```
 
+**Pattern 3: Shard RoPE for Dual Stream Attention**
+In some cases, different streams in attention may need to handle sequence parallelism differently. For example, we may want to shard the image embeddings, while replicating the text embeddings to correctly configure joint attention.
+
+```python
+class DualStreamTransformer(nn.Module):
+    """
+    Dual-stream model where we need to replicate the text components, but shard
+    the image components to correctly handle sequence parallelism.
+    """
+    _sp_plan = {
+        # In this case, the rope_preparer returns a tuple of len 4, where the
+        # first 2 items correspond to the text, and the second 2 correspond to
+        # visual inputs, so we only shard the second.
+        "rope_preparer": {
+            # Outputs 0, 1 (text) - NOT sharded (replicated)
+            # Outputs 2, 3 (image) - sharded
+            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),  # img_cos
+            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),  # img_sin
+        },
+        # Shard transformer block INPUT
+        "transformer_blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at output
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
+```
+
+NOTE: be careful to test adequately when refactoring classes that take this style of plan, as changing the order of the return values will break sequence parallelism.
+
 ### API Reference
 
 **SequenceParallelInput Parameters:**
@@ -240,7 +272,7 @@ class TransformerWithRoPE(nn.Module):
 
 ## Approach 2: Intrusive Modification (For Complex Cases)
 
-For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls. Importantly, when taking this approach, be careful to ensure that you correctly manage the `_sp_shard_depth`; if the sequence parallel shard depth is 0, Ulysses will not be used.
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`, manually insert shard/gather calls.
 
 
 **When to use:**
@@ -253,18 +285,15 @@ from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
 
 def forward(self, hidden_states, ...):
     if self.parallel_config.sequence_parallel_size > 1:
-        # <Increment the _sp_shard_depth on the fwd context>
         hidden_states = sp_shard(hidden_states, dim=1)
 
     # ... computation ...
 
     if self.parallel_config.sequence_parallel_size > 1:
         output = sp_gather(output, dim=1)
-        # <Decrement the _sp_shard_depth on the fwd context>
 
     return output
 ```
-Note that currently, `sp_shard` / `sp_gather` do *not* automatically manage the `_sp_shard_depth`; you need to be careful to manage it yourself.
 
 ---
 
@@ -439,6 +468,7 @@ Complete examples in the codebase:
 
 | Model | Path | Pattern | Notes |
 |-------|------|---------|-------|
+| **LongCat** | `vllm_omni/diffusion/models/longcat_image/longcat_image_transformer.py` | Dual-stream | Text components replicated, image components sharded |
 | **Qwen-Image** | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` | Dual-stream + preprocessing | auto_pad, separate RoPE |
 | **Wan2.2** | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` | Dual-Transformer + RoPE | Video transformer |
 | **Z-Image** | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` | Unified sequence | Concatenated input |
@@ -453,7 +483,7 @@ Complete examples in the codebase:
 Adding Sequence Parallel support to a transformer:
 
 1. ✅ **Choose approach** - Use `_sp_plan` for standard cases, intrusive modification for complex cases
-2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered?
+2. ✅ **Identify sharding boundaries** - Where should tensors be split/gathered? And which module boundaries need to be moved to facilitate this?
 3. ✅ **Extract inline operations** - Move `torch.cat`, `pad_sequence`, etc. to submodules
 4. ✅ **Define `_sp_plan`** - Declare shard/gather points as class attribute
 5. ✅ **Use `auto_pad` for variable lengths** - Support non-uniform sequences

docs/design/module/ar_module.md

@@ -85,9 +85,6 @@ classDiagram
     class InputProcessor {
         +process_inputs() EngineCoreRequest
     }
-    class OmniInputProcessor {
-        +process_inputs() OmniEngineCoreRequest
-    }
 
     class VLLMOutputProcessor {
         +process_outputs() OutputProcessorOutput
@@ -96,7 +93,6 @@ classDiagram
         +process_outputs() OutputProcessorOutput
         +_route_and_normalize()
     }
-    InputProcessor <|-- OmniInputProcessor
     VLLMOutputProcessor <|-- MultimodalOutputProcessor
 ```
 
@@ -105,7 +101,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**: `OmniInputProcessor` extends `vllm.v1.engine.input_processor.InputProcessor` to serialize prompt embeddings and additional information
+- **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`
 - **OutputProcessor**: `MultimodalOutputProcessor` extends `vllm.v1.engine.output_processor.OutputProcessor` to route and accumulate multimodal outputs
 
 ## 3. Scheduler Design
@@ -118,7 +114,7 @@ The following diagram illustrates the request flow through the AR module compone
 
 ```mermaid
 flowchart TD
-    A[OmniInputProcessor] -->|OmniEngineCoreRequest| B[OmniARScheduler]
+    A[InputProcessor stage-0 in AsyncOmniEngine] -->|EngineCoreRequest then upgraded to OmniEngineCoreRequest| B[OmniARScheduler]
     B -->|schedule: OmniNewRequestData| C[GPUARWorker]
     C -->|SchedulerOutput| D[GPUARModelRunner]
     D -->|execute_model: None| E[Model Forward Pass]
@@ -297,15 +293,17 @@ The input/output processing pipeline handles serialization, routing, and accumul
 ```mermaid
 sequenceDiagram
     participant Client
-    participant OmniInputProcessor
+    participant AsyncOmniEngine
+    participant InputProcessor
     participant Scheduler
     participant ModelRunner
     participant MultimodalOutputProcessor
     participant Client
 
-    Client->>OmniInputProcessor: prompt + prompt_embeds + additional_info
-    OmniInputProcessor->>OmniInputProcessor: Serialize tensors to payloads
-    OmniInputProcessor->>Scheduler: OmniEngineCoreRequest (with payloads)
+    Client->>AsyncOmniEngine: prompt + prompt_embeds + additional_info
+    AsyncOmniEngine->>InputProcessor: process_inputs()
+    InputProcessor->>Scheduler: EngineCoreRequest
+    AsyncOmniEngine->>AsyncOmniEngine: _upgrade_to_omni_request() + serialize_additional_information()
     Scheduler->>ModelRunner: OmniNewRequestData (with payloads)
     ModelRunner->>ModelRunner: Decode payloads → CPU tensors
     ModelRunner->>ModelRunner: Overlay prompt_embeds on inputs_embeds
@@ -318,43 +316,18 @@ sequenceDiagram
     MultimodalOutputProcessor->>Client: RequestOutput (with multimodal_output)
 ```
 
-### OmniInputProcessor
+### Stage-0 Input Processing
 
-`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:
+Stage-0 now uses upstream `InputProcessor` directly, and `AsyncOmniEngine` upgrades the request to `OmniEngineCoreRequest` while restoring omni-specific payloads.
 
 ```python
-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,
-    )
+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)
 ```
 
 ### MultimodalOutputProcessor
@@ -400,13 +373,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 `PromptEmbedsPayload` and `AdditionalInformationPayload` for efficient inter-stage data transfer
+2. **Payload serialization**: Uses serialized `additional_information` payloads together with prompt-embedding handoff for efficient inter-stage 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 prompt embeddings and additional information enable direct transfer between pipeline stages
+- **Payload support**: Serialized additional information and prompt embeddings 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

@@ -0,0 +1,203 @@
+# 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[0]
+output = outputs[0].request_output
 print(f"Trajectory timesteps shape: {output.metrics['trajectory_timesteps'].shape}")
 print(f"Trajectory latents shape: {output.latents.shape}")
 ```

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

@@ -33,18 +33,25 @@ 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 .
+pip install -v -e . --no-build-isolation
+# or VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+
 export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
@@ -61,22 +68,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.16.0
+# Pin vLLM version to 0.17.0
 cd /vllm-workspace/vllm
 git pull origin main
 git fetch origin --tags
-git checkout v0.16.0
+git checkout v0.17.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-ascend
+cd /vllm-workspace/vllm-ascend
 git pull origin main
 git fetch origin --tags
-git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
+git checkout v0.17.0
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch
-cd ../vllm-omni
+cd /vllm-workspace/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

@@ -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[0].images
+    images = outputs[0].request_output.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[0]
+        this_request_output = prompt_output.request_output
         this_images = this_request_output.images
         for i_image, image in enumerate(this_images):
             image.save(f"p{i_prompt}-img{i_image}.jpg")

docs/models/supported_models.md

@@ -50,6 +50,8 @@ th {
 |`NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` |
 |`MiMoAudioForConditionalGeneration` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` |
 |`Flux2Pipeline` | FLUX.2-dev | `black-forest-labs/FLUX.2-dev` |
+|`FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` |
+|`DreamIDOmniPipeline`| DreamID-Omni | `XuGuo699/DreamID-Omni` |
 
 
 ## List of Supported Models for NPU
@@ -65,6 +67,7 @@ th {
 |--------------|--------|-------------------|
 | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
+| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` |
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |

docs/serving/speech_api.md

@@ -1,6 +1,9 @@
 # Speech API
 
-vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation using Qwen3-TTS models.
+vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation. Supported TTS models include:
+
+- **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.
 
 Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
 
@@ -9,13 +12,22 @@ Each server instance runs a single model (specified at startup via `vllm serve <
 ### Start the Server
 
 ```bash
-# CustomVoice model (predefined speakers)
+# Qwen3-TTS: CustomVoice model (predefined speakers)
 vllm serve Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
     --stage-configs-path vllm_omni/model_executor/stage_configs/qwen3_tts.yaml \
     --omni \
     --port 8091 \
     --trust-remote-code \
     --enforce-eager
+
+# Fish Speech S2 Pro
+vllm-omni serve fishaudio/s2-pro \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager \
+    --gpu-memory-utilization 0.9
 ```
 
 ### Generate Speech
@@ -164,6 +176,42 @@ curl -X POST http://localhost:8091/v1/audio/voices \
   -F "name=custom_voice_1"
 ```
 
+## Streaming Text Input (WebSocket)
+
+The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally and generates audio per sentence as boundaries are detected.
+
+> Note: text input is always streamed incrementally. Audio output remains sentence-scoped:
+> use `stream_audio=false` for one binary frame per sentence, or `stream_audio=true` for one or more PCM chunks per sentence.
+
+### WebSocket Protocol
+
+Client -> Server:
+
+| Message | Description |
+|---------|-------------|
+| `{"type": "session.config", ...}` | Session configuration (sent once, first message) |
+| `{"type": "input.text", "text": "..."}` | Text chunk |
+| `{"type": "input.done"}` | End of input, flushes remaining buffer |
+
+Server -> Client:
+
+| Message | Description |
+|---------|-------------|
+| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
+| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
+| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
+| `{"type": "session.done", "total_sentences": N}` | Session complete |
+| `{"type": "error", "message": "..."}` | Non-fatal error |
+
+### Session Config Parameters
+
+All REST API parameters are supported, plus:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
+| `split_granularity` | string | "sentence" | Text splitting granularity |
+
 
 ```bash
 DELETE /v1/audio/voices/{name}
@@ -281,6 +329,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 
 ## Supported Models
 
+### Qwen3-TTS
+
 | Model | Task Type | Description |
 |-------|-----------|-------------|
 | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
@@ -289,6 +339,14 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 | `Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice` | CustomVoice | Smaller/faster variant |
 | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | Base | Smaller/faster variant for voice cloning |
 
+### Fish Speech S2 Pro
+
+| Model | Description |
+|-------|-------------|
+| `fishaudio/s2-pro` | 4B dual-AR TTS with DAC codec (44.1 kHz). Supports text-to-speech and voice cloning. |
+
+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.
+
 ## Error Responses
 
 ### 400 Bad Request

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -24,21 +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 |
-|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|:------------------:|:---------------:|
-| **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` |     ❌      |    ❌    |      ❌       |        ✅        |         ❌          |        ✅        |
+| 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        | ❌   |
 
 !!! 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.

docs/user_guide/diffusion/quantization/gguf.md

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

docs/user_guide/diffusion_acceleration.md

@@ -69,7 +69,7 @@ The following table shows which models are currently supported by each accelerat
 | **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **NextStep-1.1** | `stepfun-ai/NextStep-1.1` | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ |
 | **FLUX.2-klein** | `black-forest-labs/FLUX.2-klein-4B` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
-| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **FLUX.2-dev** | `black-forest-labs/FLUX.2-dev` | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
 
 ### VideoGen
 
@@ -77,6 +77,7 @@ The following table shows which models are currently supported by each accelerat
 |-------|------------------|:--------:|:---------:|:----------:|:--------------:|:------------:|:----:|:----:|
 | **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **LTX-2** | `Lightricks/LTX-2` | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
+| **DreamID-Omni** | `XuGuo699/DreamID-Omni` | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ |
 
 ### Quantization
 

docs/user_guide/examples/offline_inference/cosyvoice3.md

@@ -0,0 +1,74 @@
+# CosyVoice3
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/cosyvoice3>.
+
+
+## Setup
+
+Install dependencies:
+```
+uv pip install -e .
+```
+
+> **Note:** This includes required libraries such as `librosa`, `soundfile`,
+> `onnxruntime`, `x-transformers`, and `einops` via
+> `requirements/common.txt` and platform-specific requirements files.
+
+Download the model snapshot:
+```
+from huggingface_hub import snapshot_download
+snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
+```
+
+Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
+```json
+{
+    "model_type": "cosyvoice3",
+    "architectures": [
+        "CosyVoice3Model"
+    ]
+}
+```
+
+> **Why `config.json` is required:**
+> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
+> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
+> If no `config.json` is present, model type cannot be inferred automatically.
+> If your downloaded checkpoint already includes a valid `config.json` with
+> `model_type: "cosyvoice3"`, this manual step can be skipped.
+
+Run the offline verification script:
+```
+python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
+  --model pretrained_models/Fun-CosyVoice3-0.5B \
+  --tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
+```
+
+## Implementation Overview
+
+CosyVoice3 runs as a **2-stage Omni pipeline**.
+
+- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
+- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
+
+Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
+
+- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
+    - Tokenizes `prompt` and `prompt_text`.
+    - Extracts speech tokens and mel features from the prompt audio.
+    - Extracts a speaker embedding.
+- `CosyVoice3Model` implements both stages:
+    - Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
+    - Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
+
+Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
+
+- Stage 0 emits latent speech tokens.
+- Stage 1 consumes them via `custom_process_input_func` and outputs audio.
+
+## Example materials
+
+??? abstract "verify_e2e_cosyvoice.py"
+    ``````py
+    --8<-- "examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py"
+    ``````

docs/user_guide/examples/offline_inference/fish_speech.md

@@ -0,0 +1,97 @@
+# Fish Speech S2 Pro
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/fish_speech>.
+
+
+This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
+
+## Model Overview
+
+[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
+
+- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
+- **Fast AR**: Predicts residual codebook tokens from semantic tokens
+- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
+
+The model produces high-quality speech with voice cloning capabilities.
+
+## Prerequisites
+
+Install the `fish-speech` package for the DAC codec:
+
+```bash
+pip install fish-speech
+```
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Quick Start
+
+### Text-Only Synthesis
+
+```bash
+python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
+```
+
+Generated audio files are saved to `output_audio/` by default.
+
+### Voice Cloning
+
+Provide a reference audio file and its transcript to clone a voice:
+
+```bash
+python end2end.py \
+    --text "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Transcript of the reference audio."
+```
+
+## Streaming Mode
+
+Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
+
+```bash
+python end2end.py --text "Hello, this is a test." --streaming
+```
+
+Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
+
+## Configuration
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--text` | `"Hello, this is a test..."` | Text to synthesize |
+| `--ref-audio` | None | Path to reference audio for voice cloning |
+| `--ref-text` | None | Transcript of the reference audio |
+| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
+| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
+| `--output-dir` | `output_audio` | Output directory for WAV files |
+| `--streaming` | False | Enable streaming via AsyncOmni |
+| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
+
+## Notes
+
+- Output audio is 44.1 kHz mono WAV.
+- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
+- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.
+
+## Example materials
+
+??? abstract "diag_codebook.py"
+    ``````py
+    --8<-- "examples/offline_inference/fish_speech/diag_codebook.py"
+    ``````
+??? abstract "diag_reference_compare.py"
+    ``````py
+    --8<-- "examples/offline_inference/fish_speech/diag_reference_compare.py"
+    ``````
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/fish_speech/end2end.py"
+    ``````
+??? abstract "standalone_test.py"
+    ``````py
+    --8<-- "examples/offline_inference/fish_speech/standalone_test.py"
+    ``````

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[0].images
+    images = outputs[0].request_output.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[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -72,7 +72,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 ## Local CLI Usage

docs/user_guide/examples/offline_inference/x_to_video_audio.md

@@ -0,0 +1,58 @@
+# X-To-Video-Audio
+
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
+
+## Local CLI Usage
+### Download the Model locally
+Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
+
+```bash
+python download_dreamid_omni.py --output-dir ./dreamid_omni
+```
+After download, the model directory will look like this:
+
+```
+dreamid_omni/
+├── DreamID-Omni/
+│   ├── dreamid_omni.safetensors
+├── MMAudio/
+│   ├── ext_weights/
+│   │   ├── best_netG.pt
+│   │   ├── v1-16.pth
+├── Wan2.2-TI2V-5B/
+│   ├── google/*
+│   ├── models_t5_umt5-xxl-enc-bf16.pth
+│   ├── Wan2.2_VAE.pth
+│
+├── model_index.json # create by download_dreamid_omni.py
+```
+
+### Run the Inference
+```
+python x_to_video_audio.py \
+  --model /xx/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --output dreamid_omni.mp4
+```
+In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
+The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+
+Key arguments:
+- `--prompt`: text description (string).
+- `--model`: path to the model local directory.
+- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--image-path`: path to the input image list.
+- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
+- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--num-inference-steps`: number of denoising steps (defaults 45).
+- `--video-negative-prompt`: negative prompt for video generation.
+- `--audio-negative-prompt`: negative prompt for audio generation.
+- `--enable-cpu-offload`: enable CPU offload (defaults False).

docs/user_guide/examples/online_serving/fish_speech.md

@@ -0,0 +1,194 @@
+# Fish Speech S2 Pro
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/fish_speech>.
+
+
+## Model
+
+| Model | Description |
+|-------|-------------|
+| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) |
+
+## Gradio Demo
+
+An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
+
+```bash
+# Option 1: Launch server + Gradio together
+./run_gradio_demo.sh
+
+# Option 2: If server is already running
+python gradio_demo.py --api-base http://localhost:8091
+```
+
+Then open http://localhost:7860 in your browser.
+
+Features:
+
+- Text-to-speech synthesis
+- Voice cloning from uploaded audio or URL
+- Streaming mode (progressive PCM playback)
+
+## Launch the Server
+
+```bash
+vllm-omni serve fishaudio/s2-pro \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager \
+    --gpu-memory-utilization 0.9
+```
+
+Or use the convenience script:
+
+```bash
+./run_server.sh
+```
+
+## Send TTS Request
+
+### Using curl
+
+```bash
+# Basic TTS
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "response_format": "wav"
+    }' --output output.wav
+```
+
+### Voice Cloning
+
+Provide a reference audio (URL or base64 data URL) and its transcript:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is a cloned voice.",
+        "voice": "default",
+        "ref_audio": "https://example.com/reference.wav",
+        "ref_text": "Transcript of the reference audio."
+    }' --output cloned.wav
+```
+
+### Using Python
+
+```python
+import httpx
+
+# Basic TTS
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+### Using the CLI Client
+
+```bash
+cd examples/online_serving/fish_speech
+
+# Basic TTS
+python speech_client.py --text "Hello, how are you?"
+
+# Voice cloning
+python speech_client.py \
+    --text "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Transcript of the reference audio."
+
+# Streaming PCM output
+python speech_client.py --text "Hello world" --stream --output output.pcm
+```
+
+The CLI client supports:
+
+- `--api-base`: API base URL (default: `http://localhost:8091`)
+- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`)
+- `--text`: Text to synthesize (required)
+- `--ref-audio`: Reference audio for voice cloning (local path or URL)
+- `--ref-text`: Transcript of the reference audio
+- `--stream`: Enable streaming (PCM output)
+- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
+- `--output` (or `-o`): Output file path
+
+## Streaming
+
+Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 -
+```
+
+**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz).
+
+## API Parameters
+
+Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation.
+
+Key parameters for Fish Speech:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text to synthesize |
+| `voice` | string | "default" | Voice name (use "default" for Fish Speech) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning |
+| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) |
+| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
+| `stream` | bool | false | Stream raw PCM chunks |
+
+## Prerequisites
+
+Install the `fish-speech` package for the DAC codec:
+
+```bash
+pip install fish-speech
+```
+
+## Troubleshooting
+
+1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder
+2. **Connection refused**: Ensure the server is running on the correct port
+3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings
+4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM
+
+## Example materials
+
+??? abstract "gradio_demo.py"
+    ``````py
+    --8<-- "examples/online_serving/fish_speech/gradio_demo.py"
+    ``````
+??? abstract "run_gradio_demo.sh"
+    ``````sh
+    --8<-- "examples/online_serving/fish_speech/run_gradio_demo.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/fish_speech/run_server.sh"
+    ``````
+??? abstract "speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/fish_speech/speech_client.py"
+    ``````

docs/user_guide/examples/online_serving/image_to_video.md

@@ -29,6 +29,35 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+Local file storage behavior can be controlled via the following environment variables:
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -38,7 +67,7 @@ The script allows overriding:
 bash run_curl_image_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -X POST http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=A bear playing with yarn, smooth motion" \
   -F "negative_prompt=low quality, blurry, static" \
@@ -52,7 +81,23 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "guidance_scale_2=1.0" \
   -F "boundary_ratio=0.875" \
   -F "flow_shift=12.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_i2v_output.mp4
 ```
 
 ## Request Format
@@ -66,6 +111,18 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "input_reference=@/path/to/qwen-bear.png"
 ```
 
+### Alternative JSON-Safe Reference Input
+
+Use `image_reference` when you want to pass a URL or JSON-safe image reference
+instead of uploading a file. Do not send `input_reference` and
+`image_reference` together.
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F 'image_reference={"image_url":"https://example.com/qwen-bear.png"}'
+```
+
 ### Generation with Parameters
 
 ```bash
@@ -85,6 +142,63 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "seed=42"
 ```
 
+## Create Response Format
+
+`POST /v1/videos` returns a job record.
+
+```json
+{
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+  "prompt": "A bear playing with yarn, smooth motion",
+  "created_at": 1234567890
+}
+```
+
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
+
+```bash
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_i2v_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+```
+
 ## Example materials
 
 ??? abstract "run_curl_image_to_video.sh"

docs/user_guide/examples/online_serving/text_to_video.md

@@ -29,6 +29,34 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -38,7 +66,7 @@ The script allows overriding:
 bash run_curl_text_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -s http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   -F "width=832" \
@@ -50,7 +78,24 @@ curl -s http://localhost:8091/v1/videos \
   -F "guidance_scale=4.0" \
   -F "guidance_scale_2=4.0" \
   -F "boundary_ratio=0.875" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+  -F "flow_shift=5.0" \
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_output.mp4
 ```
 
 ## Request Format
@@ -85,8 +130,9 @@ curl -X POST http://localhost:8091/v1/videos \
 | Parameter             | Type   | Default | Description                                      |
 | --------------------- | ------ | ------- | ------------------------------------------------ |
 | `prompt`              | str    | -       | Text description of the desired video            |
+| `seconds`             | str    | None    | Clip duration in seconds                         |
+| `size`                | str    | None    | Output size in `WIDTHxHEIGHT` format             |
 | `negative_prompt`     | str    | None    | Negative prompt                                  |
-| `n`                   | int    | 1       | Number of videos to generate                     |
 | `width`               | int    | None    | Video width in pixels                            |
 | `height`              | int    | None    | Video height in pixels                           |
 | `num_frames`          | int    | None    | Number of frames to generate                     |
@@ -98,24 +144,62 @@ curl -X POST http://localhost:8091/v1/videos \
 | `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
 | `seed`                | int    | None    | Random seed (reproducible)                       |
 | `lora`                | object | None    | LoRA configuration                               |
-| `extra_body`          | object | None    | Model-specific extra parameters                  |
 
-## Response Format
+## Create Response Format
+
+`POST /v1/videos` returns a job record.
 
 ```json
 {
-  "created": 1234567890,
-  "data": [
-    { "b64_json": "<base64-mp4>" }
-  ]
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+  "prompt": "A cinematic view of a futuristic city at sunset",
+  "created_at": 1234567890
 }
 ```
 
-## Extract Video
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
 
 ```bash
-# Extract base64 from response and decode to video
-cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
 ```
 
 ## Example materials

examples/offline_inference/bagel/end2end.py

@@ -34,7 +34,7 @@ def parse_args():
         help="Path to input image for img2img.",
     )
 
-    # OmniLLM init args
+    # Omni runtime init args
     parser.add_argument("--log-stats", action="store_true", default=False)
     parser.add_argument("--init-sleep-seconds", type=int, default=20)
     parser.add_argument("--batch-timeout", type=int, default=5)
@@ -58,6 +58,11 @@ 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
@@ -103,6 +108,7 @@ 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,
         }
     )
 
@@ -163,25 +169,17 @@ def main():
 
     omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
 
-    for i, req_output in enumerate(omni_outputs):
+    img_idx = 0
+    for req_output in 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")
-
-        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}")
+        for j, img in enumerate(images):
+            save_path = f"output_{img_idx}_{j}.png"
+            img.save(save_path)
+            print(f"[Info] Saved image to {save_path}")
+        img_idx += 1
 
     print(omni_outputs)
 

examples/offline_inference/cosyvoice3/README.md

@@ -0,0 +1,64 @@
+# CosyVoice3
+
+## Setup
+
+Install dependencies:
+```
+uv pip install -e .
+```
+
+> **Note:** This includes required libraries such as `librosa`, `soundfile`,
+> `onnxruntime`, `x-transformers`, and `einops` via
+> `requirements/common.txt` and platform-specific requirements files.
+
+Download the model snapshot:
+```
+from huggingface_hub import snapshot_download
+snapshot_download('FunAudioLLM/Fun-CosyVoice3-0.5B-2512', local_dir='pretrained_models/Fun-CosyVoice3-0.5B')
+```
+
+Add `config.json` in `pretrained_models/Fun-CosyVoice3-0.5B/`:
+```json
+{
+    "model_type": "cosyvoice3",
+    "architectures": [
+        "CosyVoice3Model"
+    ]
+}
+```
+
+> **Why `config.json` is required:**
+> `AutoConfig.register("cosyvoice3", CosyVoice3Config)` only registers a class mapping.
+> The loader still needs `model_type: "cosyvoice3"` from `config.json` to select that class.
+> If no `config.json` is present, model type cannot be inferred automatically.
+> If your downloaded checkpoint already includes a valid `config.json` with
+> `model_type: "cosyvoice3"`, this manual step can be skipped.
+
+Run the offline verification script:
+```
+python examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py \
+  --model pretrained_models/Fun-CosyVoice3-0.5B \
+  --tokenizer pretrained_models/Fun-CosyVoice3-0.5B/CosyVoice-BlankEN
+```
+
+## Implementation Overview
+
+CosyVoice3 runs as a **2-stage Omni pipeline**.
+
+- **Stage 0** (`talker`) converts text + prompt audio to speech tokens.
+- **Stage 1** (`code2wav`) converts speech tokens + prompt features to waveform via flow matching and HiFiGAN.
+
+Key components live in `vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py`.
+
+- `CosyVoice3MultiModalProcessor` builds the multimodal inputs:
+    - Tokenizes `prompt` and `prompt_text`.
+    - Extracts speech tokens and mel features from the prompt audio.
+    - Extracts a speaker embedding.
+- `CosyVoice3Model` implements both stages:
+    - Stage 0 uses `CosyVoice3LM` and outputs speech tokens + conditioning features.
+    - Stage 1 runs the flow model (DiT-based CFM) and HiFiGAN to synthesize waveform.
+
+Stage wiring is configured in `vllm_omni/model_executor/stage_configs/cosyvoice3.yaml`.
+
+- Stage 0 emits latent speech tokens.
+- Stage 1 consumes them via `custom_process_input_func` and outputs audio.

examples/offline_inference/cosyvoice3/verify_e2e_cosyvoice.py

@@ -0,0 +1,196 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import argparse
+import os
+from pathlib import Path
+
+import librosa
+import numpy as np
+import soundfile as sf
+from vllm import SamplingParams
+from vllm.assets.audio import AudioAsset
+
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.model_executor.models.cosyvoice3.config import CosyVoice3Config
+from vllm_omni.model_executor.models.cosyvoice3.tokenizer import get_qwen_tokenizer
+from vllm_omni.model_executor.models.cosyvoice3.utils import extract_text_token
+
+
+def _ensure_mel_filters_asset() -> None:
+    repo_root = Path(__file__).resolve().parents[3]
+    filters_path = repo_root / "vllm_omni" / "model_executor" / "models" / "cosyvoice3" / "assets" / "mel_filters.npz"
+    if filters_path.exists():
+        return
+
+    source_url = "https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/mel_filters.npz"
+    raise FileNotFoundError(
+        "Missing CosyVoice3 mel filter asset:\n"
+        f"  {filters_path}\n"
+        "Download it with:\n"
+        f"  mkdir -p {filters_path.parent} && "
+        f"curl -L {source_url} -o {filters_path}"
+    )
+
+
+def run_e2e():
+    parser = argparse.ArgumentParser()
+    # ""FunAudioLLM/Fun-CosyVoice3-0.5B-2512
+    parser.add_argument(
+        "--model",
+        type=str,
+        required=True,
+        help="Path to CosyVoice3 model directory (e.g., pretrained_models/Fun-CosyVoice3-0.5B/).",
+    )
+    parser.add_argument("--stage-config", type=str, default="vllm_omni/model_executor/stage_configs/cosyvoice3.yaml")
+    parser.add_argument("--prompt", type=str, default="Hello, this is a test of the CosyVoice system capability.")
+    parser.add_argument(
+        "--prompt-text",
+        type=str,
+        default="You are a helpful assistant.<|endofprompt|>Testing my voices. Why should I not?",
+    )
+    parser.add_argument("--audio-path", type=str, default="prompt.wav")
+    parser.add_argument(
+        "--tokenizer",
+        type=str,
+        required=True,
+        help="Path to tokenizer directory (e.g., <model_path>/CosyVoice-BlankEN).",
+    )
+    args = parser.parse_args()
+    _ensure_mel_filters_asset()
+    # Ensure tokenizer directory exists
+    if not os.path.exists(args.tokenizer):
+        raise FileNotFoundError(f"{args.tokenizer} does not exist!")
+
+    # Ensure stage config exists
+    if not os.path.exists(args.stage_config):
+        raise FileNotFoundError(f"{args.stage_config} does not exist!")
+
+    print(f"Initializing cosyvoice E2E with model={args.model}")
+
+    # Initialize Omni
+    # This spins up the engine(s) based on the stage config
+    # We pass trust_remote_code=True same as Qwen examples
+    omni = Omni(
+        model=args.model,
+        stage_configs_path=args.stage_config,
+        trust_remote_code=True,
+        tokenizer=args.tokenizer,
+        log_stats=True,
+    )
+
+    sampling_cfg = {"top_p": 0.8, "top_k": 25, "eos_token_id": 6561 + 1}
+
+    print("Model initialized. Preparing inputs...")
+    if args.audio_path:
+        if not os.path.exists(args.audio_path):
+            raise FileNotFoundError(f"Audio file not found: {args.audio_path}")
+        # Load at native sample rate
+        audio_signal, sr = librosa.load(args.audio_path, sr=None)
+
+        # Validate sample rate before processing (similar to original CosyVoice)
+        min_sr = 16000
+        if sr < min_sr:
+            raise ValueError(
+                f"Audio sample rate {sr} Hz is too low. "
+                f"Minimum required: {min_sr} Hz. "
+                f"Please provide audio with sample rate >= {min_sr} Hz."
+            )
+
+        audio_data = (audio_signal.astype(np.float32), sr)
+    else:
+        audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
+
+    prompts = {
+        "prompt": args.prompt,
+        "multi_modal_data": {
+            "audio": audio_data,
+        },
+        "mm_processor_kwargs": {
+            "prompt_text": args.prompt_text,
+            "sample_rate": audio_data[1],
+        },
+    }
+
+    print(f"Generating for prompt: {args.prompt}")
+
+    config = CosyVoice3Config()
+    tokenizer = get_qwen_tokenizer(
+        token_path=args.tokenizer,
+        skip_special_tokens=config.skip_special_tokens,
+        version=config.version,
+    )
+    _, text_token_len = extract_text_token(args.prompt, tokenizer, config.allowed_special)
+    base_len = int(text_token_len)
+    min_len = int(base_len * config.min_token_text_ratio)
+    max_len = int(base_len * config.max_token_text_ratio)
+
+    # Build SamplingParams for each stage (GPT, S2Mel, Vocoder)
+    gpt_sampling = SamplingParams(
+        temperature=1.0,
+        top_p=sampling_cfg["top_p"],
+        top_k=sampling_cfg["top_k"],
+        repetition_penalty=2.0,
+        min_tokens=min_len,
+        max_tokens=max_len,
+        stop_token_ids=[sampling_cfg["eos_token_id"]],
+        # allowed_token_ids=list(range(6561+3)),
+        detokenize=False,
+    )
+    # Not used
+    s2mel_sampling = SamplingParams(
+        temperature=1.0,
+        top_p=1.0,
+        top_k=-1,
+        repetition_penalty=2.0,
+        max_tokens=256,
+        detokenize=False,
+    )
+
+    sampling_params_list = [gpt_sampling, s2mel_sampling]
+
+    # Start profiling (requires VLLM_TORCH_PROFILER_DIR env var)
+    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
+        print("Starting profiler...")
+        omni.start_profile()
+
+    # Generate (Omni orchestrator requires a per-stage SamplingParams list)
+    outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list[:2]))
+
+    # Stop profiling and get results
+    if os.environ.get("VLLM_TORCH_PROFILER_DIR"):
+        print("Stopping profiler...")
+        profile_results = omni.stop_profile()
+        print(f"Profile traces saved to: {profile_results}")
+
+    print(outputs)
+    # Verify outputs
+    print(f"Received {len(outputs)} outputs.")
+    for i, output in enumerate(outputs):
+        try:
+            ro = output.request_output
+            if ro is None:
+                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)
+
+            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()
+
+
+if __name__ == "__main__":
+    run_e2e()

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[0]
+    first_out = outputs[0].request_output
     req_out: OmniRequestOutput = first_out
 
     # Verify trajectory data (from custom pipeline)

examples/offline_inference/fish_speech/README.md

@@ -0,0 +1,75 @@
+# Fish Speech S2 Pro
+
+This directory contains an offline demo for running Fish Speech S2 Pro (`fishaudio/s2-pro`) with vLLM Omni. It supports text-only synthesis and voice cloning with reference audio.
+
+## Model Overview
+
+[Fish Speech S2 Pro](https://huggingface.co/fishaudio/s2-pro) is a 4B dual-AR text-to-speech model by FishAudio. It uses:
+
+- **Slow AR**: Generates semantic tokens autoregressively (Qwen3-based backbone)
+- **Fast AR**: Predicts residual codebook tokens from semantic tokens
+- **DAC Decoder**: Converts codec codes to 44.1 kHz audio waveform
+
+The model produces high-quality speech with voice cloning capabilities.
+
+## Prerequisites
+
+Install the `fish-speech` package for the DAC codec:
+
+```bash
+pip install fish-speech
+```
+
+## Setup
+
+Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
+
+## Quick Start
+
+### Text-Only Synthesis
+
+```bash
+python end2end.py --text "Hello, this is a test of the Fish Speech text to speech system."
+```
+
+Generated audio files are saved to `output_audio/` by default.
+
+### Voice Cloning
+
+Provide a reference audio file and its transcript to clone a voice:
+
+```bash
+python end2end.py \
+    --text "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Transcript of the reference audio."
+```
+
+## Streaming Mode
+
+Add `--streaming` to stream audio chunks progressively via `AsyncOmni` (requires `async_chunk: true` in the stage config):
+
+```bash
+python end2end.py --text "Hello, this is a test." --streaming
+```
+
+Each DAC decoder chunk is logged as it arrives. The final WAV file is written once generation completes.
+
+## Configuration
+
+| Argument | Default | Description |
+|----------|---------|-------------|
+| `--text` | `"Hello, this is a test..."` | Text to synthesize |
+| `--ref-audio` | None | Path to reference audio for voice cloning |
+| `--ref-text` | None | Transcript of the reference audio |
+| `--model` | `fishaudio/s2-pro` | HuggingFace model path |
+| `--stage-configs-path` | `fish_speech_s2_pro.yaml` | Path to stage configs YAML |
+| `--output-dir` | `output_audio` | Output directory for WAV files |
+| `--streaming` | False | Enable streaming via AsyncOmni |
+| `--stage-init-timeout` | 600 | Stage initialization timeout (seconds) |
+
+## Notes
+
+- Output audio is 44.1 kHz mono WAV.
+- The DAC codec weights (`codec.pth`) are loaded lazily from the model directory.
+- Voice cloning encodes the reference audio using the DAC codec to extract semantic codes, then prepends them as a system message.

examples/offline_inference/fish_speech/end2end.py

@@ -0,0 +1,258 @@
+"""Offline inference demo for Fish Speech S2 Pro via vLLM Omni.
+
+Generates speech from text using the fishaudio/s2-pro model.
+Supports both text-only synthesis and voice cloning with reference audio.
+
+Usage:
+    # Text-only synthesis
+    python end2end.py --text "Hello, this is a test."
+
+    # Voice cloning with reference audio
+    python end2end.py --text "Hello, this is a test." --ref-audio ref.wav --ref-text "Reference text."
+
+    # Streaming mode
+    python end2end.py --text "Hello, this is a test." --streaming
+"""
+
+import asyncio
+import logging
+import os
+import time
+
+import soundfile as sf
+import torch
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import AsyncOmni, Omni
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_MODEL = "fishaudio/s2-pro"
+DEFAULT_STAGE_CONFIG = os.path.join(
+    os.path.dirname(__file__),
+    "..",
+    "..",
+    "..",
+    "vllm_omni",
+    "model_executor",
+    "stage_configs",
+    "fish_speech_s2_pro.yaml",
+)
+
+
+def build_prompt(
+    text: str,
+    ref_audio_path: str | None = None,
+    ref_text: str | None = None,
+    model_name: str = DEFAULT_MODEL,
+) -> dict:
+    """Build a prompt for Fish Speech S2 Pro.
+
+    Uses Qwen3 chat template format with <|voice|> modality marker:
+      <|im_start|>user\n<|speaker:0|>text<|im_end|>\n<|im_start|>assistant\n<|voice|>
+
+    The model then generates semantic tokens autoregressively.
+    """
+    from transformers import AutoTokenizer
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+
+    # Build user message with speaker tag prefix.
+    user_text = f"<|speaker:0|>{text}"
+
+    # Use chat template for user turn + assistant generation prompt.
+    messages = [{"role": "user", "content": user_text}]
+    prompt_ids = tokenizer.apply_chat_template(
+        messages,
+        tokenize=True,
+        add_generation_prompt=True,
+    )
+
+    # Append <|voice|> token (151673) to signal voice generation mode.
+    voice_token_id = tokenizer.encode("<|voice|>", add_special_tokens=False)
+    prompt_ids = prompt_ids + voice_token_id
+
+    additional_information = {
+        "text": [text],
+        "max_new_tokens": [4096],
+    }
+
+    if ref_audio_path:
+        additional_information["ref_audio"] = [ref_audio_path]
+    if ref_text:
+        additional_information["ref_text"] = [ref_text]
+
+    return {
+        "prompt_token_ids": prompt_ids,
+        "additional_information": additional_information,
+    }
+
+
+def _save_wav(output_dir: str, request_id: str, mm: dict) -> None:
+    """Concatenate audio chunks and write to a wav file."""
+    audio_data = mm["audio"]
+    sr_raw = mm["sr"]
+    sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
+    sr = sr_val.item() if hasattr(sr_val, "item") else int(sr_val)
+    audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
+    out_wav = os.path.join(output_dir, f"output_{request_id}.wav")
+    sf.write(out_wav, audio_tensor.float().cpu().numpy().flatten(), samplerate=sr, format="WAV")
+    logger.info("Request %s: saved audio to %s (sr=%d)", request_id, out_wav, sr)
+
+
+def main(args):
+    """Run offline inference with Omni."""
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
+    model_name = args.model or DEFAULT_MODEL
+
+    inputs = [
+        build_prompt(
+            text=args.text,
+            ref_audio_path=args.ref_audio,
+            ref_text=args.ref_text,
+            model_name=model_name,
+        )
+    ]
+
+    omni = Omni(
+        model=model_name,
+        stage_configs_path=stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    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,
+        )
+    t_end = time.perf_counter()
+    logger.info("Total inference time: %.1f ms", (t_end - t_start) * 1000)
+
+
+async def main_streaming(args):
+    """Run offline inference with AsyncOmni for streaming audio chunks."""
+    output_dir = args.output_dir
+    os.makedirs(output_dir, exist_ok=True)
+
+    stage_configs_path = args.stage_configs_path or DEFAULT_STAGE_CONFIG
+    model_name = args.model or DEFAULT_MODEL
+
+    prompt = build_prompt(
+        text=args.text,
+        ref_audio_path=args.ref_audio,
+        ref_text=args.ref_text,
+        model_name=model_name,
+    )
+
+    omni = AsyncOmni(
+        model=model_name,
+        stage_configs_path=stage_configs_path,
+        log_stats=args.log_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    request_id = "0"
+    t_start = time.perf_counter()
+    t_prev = t_start
+    chunk_idx = 0
+
+    async for stage_output in omni.generate(prompt, request_id=request_id):
+        mm = stage_output.request_output.outputs[0].multimodal_output
+        if not stage_output.finished:
+            t_now = time.perf_counter()
+            audio = mm.get("audio")
+            n = len(audio) if isinstance(audio, list) else (0 if audio is None else 1)
+            dt_ms = (t_now - t_prev) * 1000
+            ttfa_ms = (t_now - t_start) * 1000
+            if chunk_idx == 0:
+                logger.info("Request %s: chunk %d samples=%d TTFA=%.1fms", request_id, chunk_idx, n, ttfa_ms)
+            else:
+                logger.info("Request %s: chunk %d samples=%d inter_chunk=%.1fms", request_id, chunk_idx, n, dt_ms)
+            t_prev = t_now
+            chunk_idx += 1
+        else:
+            t_end = time.perf_counter()
+            total_ms = (t_end - t_start) * 1000
+            logger.info("Request %s: done total=%.1fms chunks=%d", request_id, total_ms, chunk_idx)
+            _save_wav(output_dir, request_id, mm)
+
+
+def parse_args():
+    parser = FlexibleArgumentParser(
+        description="Fish Speech S2 Pro offline inference with vLLM Omni",
+    )
+    parser.add_argument(
+        "--text",
+        type=str,
+        default="Hello, this is a test of the Fish Speech text to speech system.",
+        help="Text to synthesize.",
+    )
+    parser.add_argument(
+        "--ref-audio",
+        type=str,
+        default=None,
+        help="Path to reference audio for voice cloning.",
+    )
+    parser.add_argument(
+        "--ref-text",
+        type=str,
+        default=None,
+        help="Reference text corresponding to the reference audio.",
+    )
+    parser.add_argument(
+        "--model",
+        type=str,
+        default=None,
+        help=f"HuggingFace model path (default: {DEFAULT_MODEL}).",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to stage configs YAML.",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="output_audio",
+        help="Output directory for generated wav files.",
+    )
+    parser.add_argument(
+        "--log-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics.",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=600,
+        help="Timeout for initializing stages in seconds.",
+    )
+    parser.add_argument(
+        "--streaming",
+        action="store_true",
+        default=False,
+        help="Stream audio chunks as they arrive via AsyncOmni.",
+    )
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    logging.basicConfig(level=logging.INFO)
+    args = parse_args()
+    if args.streaming:
+        asyncio.run(main_streaming(args))
+    else:
+        main(args)

examples/offline_inference/glm_image/end2end.py

@@ -247,6 +247,7 @@ 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
@@ -310,36 +311,35 @@ 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":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
+            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:
-                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)}")
+                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,6 +464,11 @@ 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,6 +173,11 @@ 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.")
@@ -207,6 +212,7 @@ 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
@@ -293,12 +299,11 @@ def main():
                 )
 
             if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_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.")
+                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.")
             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 is built on the Hunyuan v1 architecture, specifically the tencent/Hunyuan-A13B-Instruct model. It utilizes two tokenizer delimiter templates:
+The tencent/HunyuanImage-3.0-Instruct base model uses the tencent/Hunyuan-A13B-Instruct backbone. 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,6 +41,11 @@ 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()
 
 
@@ -52,7 +57,10 @@ def load_image(image_path: str) -> Image.Image:
 
 
 def main(args: argparse.Namespace) -> None:
-    omni = Omni(model=args.model)
+    omni = Omni(
+        model=args.model,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
 
     prompt_dict = {
         "prompt": args.prompt,
@@ -70,8 +78,8 @@ def main(args: argparse.Namespace) -> None:
     prompts = [prompt_dict]
     omni_outputs = omni.generate(prompts=prompts)
 
-    prompt_text = omni_outputs[0].request_output[0].prompt
-    generated_text = omni_outputs[0].request_output[0].outputs[0].text
+    prompt_text = omni_outputs[0].request_output.prompt
+    generated_text = omni_outputs[0].request_output.outputs[0].text
     print(f"Prompt: {prompt_text}")
     print(f"Text: {generated_text}")
 

examples/offline_inference/image_to_image/image_edit.py

@@ -325,6 +325,11 @@ 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()
 
 
@@ -389,6 +394,7 @@ 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")
 
@@ -463,12 +469,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[0].images
+    # omni.generate() returns list[OmniRequestOutput], extract images from request_output.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[0]
+    req_out = first_output.request_output
     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

@@ -176,6 +176,11 @@ 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()
 
 
@@ -281,6 +286,7 @@ 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:
@@ -338,8 +344,6 @@ 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,6 +43,11 @@ 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()
 
 
@@ -70,6 +75,7 @@ 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(
@@ -99,14 +105,12 @@ def main() -> None:
 
     lines: list[str] = []
     for stage_outputs in outputs:
-        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")
+        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")
 
     print("\n".join(lines))
 

examples/offline_inference/mammothmodal2_preview/run_mammothmoda2_t2i.py

@@ -148,19 +148,16 @@ 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_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)
+        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)
     return images
 
 

examples/offline_inference/mimo_audio/end2end.py

@@ -180,15 +180,16 @@ def main(args):
     # Get the query function and call it with appropriate parameters
     query_func = query_map[args.query_type]
 
-    omni_llm = Omni(
+    omni = Omni(
         model=model_name,
         stage_configs_path=args.stage_configs_path,
         log_stats=args.enable_stats,
-        log_file=("omni_llm_pipeline.log" if args.enable_stats else None),
+        log_file=("omni_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(
@@ -290,7 +291,7 @@ def main(args):
     prompts = [copy.deepcopy(query_result) for _ in range(args.num_prompts)]
 
     print("prompts", prompts)
-    omni_outputs = omni_llm.generate(prompts, sampling_params_list)
+    omni_outputs = omni.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:
@@ -298,45 +299,44 @@ 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":
-            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")
+            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":
-            for output in stage_outputs.request_output:
-                request_id = output.request_id
-                audio_tensor = output.outputs[0].multimodal_output.get("audio")
+            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,6 +434,11 @@ 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,13 +320,14 @@ def main(args):
         query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate)
     else:
         query_result = query_func()
-    omni_llm = Omni(
+    omni = 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
@@ -378,9 +379,11 @@ def main(args):
             prompt["modalities"] = output_modalities
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
-    if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+    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)
 
     # 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
@@ -389,43 +392,42 @@ 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":
-            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}")
+            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":
-            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}")
+            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 += len(stage_outputs.request_output)
-        if profiler_enabled and processed_count >= total_requests:
+        processed_count += 1
+        if profiler_enabled and hasattr(omni, "stop_profile") 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_llm.stop_profile()
+            omni.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_llm.close()
+    omni.close()
 
 
 def parse_args():
@@ -539,6 +541,11 @@ 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,11 +294,12 @@ def main(args):
     else:
         query_result = query_func()
 
-    omni_llm = Omni(
+    omni = 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(
@@ -354,8 +355,8 @@ def main(args):
 
     profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
     if profiler_enabled:
-        omni_llm.start_profile(stages=[0])
-    omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
+        omni.start_profile(stages=[0])
+    omni_generator = omni.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)
@@ -366,51 +367,50 @@ 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":
-            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}")
+            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":
-            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")
+            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 += len(stage_outputs.request_output)
+        processed_count += 1
         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_llm.stop_profile()
+            omni.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_llm.close()
+    omni.close()
 
 
 def parse_args():
@@ -527,6 +527,11 @@ 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,52 +231,47 @@ async def run_single_request(
             sampling_params_list=sampling_params_list,
             output_modalities=output_modalities,
         ):
-            if not isinstance(omni_output.request_output, list):
-                outputs_list = [omni_output.request_output]
-            else:
-                outputs_list = omni_output.request_output
+            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 = []
 
-            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)
+                    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()
@@ -387,6 +382,7 @@ 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
@@ -589,6 +585,11 @@ 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/end2end.py

@@ -51,16 +51,68 @@ 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))
 
-        tok, tcfg = _cache[model_name]
+            # 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]
         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)
@@ -319,14 +371,15 @@ 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):
-            for output in stage_outputs.request_output:
-                _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
+            output = stage_outputs.request_output
+            _save_wav(output_dir, output.request_id, output.outputs[0].multimodal_output)
 
 
 async def main_streaming(args):
@@ -340,6 +393,7 @@ 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):
@@ -481,6 +535,11 @@ 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,6 +90,11 @@ 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()
 
 
@@ -133,7 +138,10 @@ def main():
     print(f"{'=' * 60}\n")
 
     # Initialize Omni with Stable Audio model
-    omni = Omni(model=args.model)
+    omni = Omni(
+        model=args.model,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
+    )
 
     # Calculate audio end time
     audio_end_in_s = args.audio_start + args.audio_length
@@ -177,7 +185,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[0]
+    request_output = output.request_output
     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[0].images
+    images = outputs[0].request_output.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[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.images[0].save(f"{i}.jpg")
 ```
 
 !!! info
@@ -69,7 +69,7 @@ if __name__ == "__main__":
       }
     ])
     for i, output in enumerate(outputs):
-      image = output.request_output[0].images[0].save(f"{i}.jpg")
+      image = output.request_output.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[0]
+            req_out = output.request_output
             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,7 +14,6 @@ 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
 
 
@@ -218,6 +217,11 @@ 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()
 
 
@@ -305,6 +309,7 @@ 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,
     }
@@ -400,20 +405,18 @@ def main():
         else:
             print("[Profiler] No valid profiling data returned.")
 
-    # Extract images from OmniRequestOutput
-    # omni.generate() returns list[OmniRequestOutput], extract images from the first output
+    # omni.generate() returns list[OmniRequestOutput]
     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[0]
-    if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
-        raise ValueError("Invalid request_output structure or missing 'images' key")
+    req_out = first_output.request_output
+    if not hasattr(req_out, "images"):
+        raise ValueError("Invalid request_output structure or missing 'images'.")
 
     images = req_out.images
     if not images:

examples/offline_inference/text_to_video/text_to_video.py

@@ -132,6 +132,11 @@ 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.",
+    )
     return parser.parse_args()
 
 
@@ -184,6 +189,7 @@ def main():
         enforce_eager=args.enforce_eager,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
+        enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
     )
 
     if profiler_enabled:
@@ -240,8 +246,6 @@ 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/x_to_video_audio/download_dreamid_omni.py

@@ -0,0 +1,103 @@
+import argparse
+import fcntl
+import os
+import site
+import subprocess
+import tempfile
+import time
+from pathlib import Path
+
+from huggingface_hub import snapshot_download
+
+DEPENDENCY_REPO = "https://github.com/bytedance/DreamID-V.git"
+DEPENDENCY_BRANCH = "omni"
+CACHE_DIR = Path(tempfile.gettempdir()) / "vllm-omni-dependency"
+LOCK_FILE = CACHE_DIR / ".install.lock"
+DEPENDENCY_DIR = CACHE_DIR / "DreamID-Omni"
+
+
+def download_dependency():
+    CACHE_DIR.mkdir(parents=True, exist_ok=True)
+
+    with open(LOCK_FILE, "w") as f:
+        fcntl.flock(f, fcntl.LOCK_EX)
+        if not DEPENDENCY_DIR.exists():
+            print(f"Downloading DreamID-Omni to {DEPENDENCY_DIR} ...")
+            subprocess.run(
+                ["git", "clone", "--depth", "1", DEPENDENCY_REPO, "--branch", DEPENDENCY_BRANCH, str(DEPENDENCY_DIR)],
+                check=True,
+            )
+            print("Download finished.")
+        fcntl.flock(f, fcntl.LOCK_UN)
+
+    # write .pth to site-packages
+    site_packages = Path(site.getsitepackages()[0])
+    pth_file = site_packages / "vllm_omni_dependency.pth"
+    pth_file.write_text(str(DEPENDENCY_DIR))
+    print(f"Added {DEPENDENCY_DIR} to site-packages via {pth_file}")
+
+
+def timed_download(repo_id: str, local_dir: str, allow_patterns: list | None = None):
+    """Download files from HF repo and log time + destination."""
+    if os.path.exists(local_dir):
+        print(f"Directory {local_dir} already exists. Skipping download.")
+        return
+    print(f"Starting download from {repo_id} into {local_dir}")
+    start_time = time.time()
+
+    snapshot_download(
+        repo_id=repo_id,
+        local_dir=local_dir,
+        local_dir_use_symlinks=False,
+        allow_patterns=allow_patterns,
+    )
+
+    elapsed = time.time() - start_time
+    print(f"✅ Finished downloading {repo_id} in {elapsed:.2f} seconds. Files saved at: {local_dir}")
+
+
+def main(output_dir: str):
+    # Wan2.2
+    wan_dir = os.path.join(output_dir, "Wan2.2-TI2V-5B")
+    timed_download(
+        repo_id="Wan-AI/Wan2.2-TI2V-5B",
+        local_dir=wan_dir,
+        allow_patterns=["google/*", "models_t5_umt5-xxl-enc-bf16.pth", "Wan2.2_VAE.pth"],
+    )
+
+    # MMAudio
+    mm_audio_dir = os.path.join(output_dir, "MMAudio")
+    timed_download(
+        repo_id="hkchengrex/MMAudio",
+        local_dir=mm_audio_dir,
+        allow_patterns=["ext_weights/best_netG.pt", "ext_weights/v1-16.pth"],
+    )
+
+    dreamid_dir = os.path.join(output_dir, "DreamID-Omni")
+
+    timed_download(repo_id="XuGuo699/DreamID-Omni", local_dir=dreamid_dir)
+
+    # Now we construct the config file
+    import json
+
+    data = {
+        "_class_name": "DreamIDOmniPipeline",
+        "fusion": "DreamID-Omni/dreamid_omni.safetensors",
+    }
+
+    with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
+        json.dump(data, f, indent=2)
+
+    print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
+
+    # now we download the dependency code
+    download_dependency()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description="Download models from Hugging Face")
+    parser.add_argument(
+        "--output-dir", type=str, default="./dreamid_omni", help="Base directory to save downloaded models"
+    )
+    args = parser.parse_args()
+    main(args.output_dir)

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -0,0 +1,58 @@
+# X-To-Video-Audio
+
+The `DreamID-Omni` pipeline generates short videos from text, image and video.
+
+## Local CLI Usage
+### Download the Model locally
+Since DreamID-Omni combine multiple models, and without any config, so we need to download them locally.
+
+```bash
+python download_dreamid_omni.py --output-dir ./dreamid_omni
+```
+After download, the model directory will look like this:
+
+```
+dreamid_omni/
+├── DreamID-Omni/
+│   ├── dreamid_omni.safetensors
+├── MMAudio/
+│   ├── ext_weights/
+│   │   ├── best_netG.pt
+│   │   ├── v1-16.pth
+├── Wan2.2-TI2V-5B/
+│   ├── google/*
+│   ├── models_t5_umt5-xxl-enc-bf16.pth
+│   ├── Wan2.2_VAE.pth
+│
+├── model_index.json # create by download_dreamid_omni.py
+```
+
+### Run the Inference
+```
+python x_to_video_audio.py \
+  --model /xx/dreamid_omni \
+  --prompt "Two people walking together and singing happily" \
+  --image-path ./example0.png ./example1.png \
+  --audio-path ./example0.wav ./example1.wav \
+  --video-negative-prompt "jitter, bad hands, blur, distortion" \
+  --audio-negative-prompt "robotic, muffled, echo, distorted" \
+  --cfg-parallel-size 2 \
+  --num-inference-steps 45 \
+  --height 704 \
+  --width 1280 \
+  --output dreamid_omni.mp4
+```
+In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled.
+The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload.
+
+Key arguments:
+- `--prompt`: text description (string).
+- `--model`: path to the model local directory.
+- `--height/--width`: output resolution (defaults 704 * 1024).
+- `--image-path`: path to the input image list.
+- `--audio-path`: path to the input audio list, indicate the timbre of the output video.
+- `--cfg-parallel-size`: number of parallel cfg parallel (defaults 1).
+- `--num-inference-steps`: number of denoising steps (defaults 45).
+- `--video-negative-prompt`: negative prompt for video generation.
+- `--audio-negative-prompt`: negative prompt for audio generation.
+- `--enable-cpu-offload`: enable CPU offload (defaults False).

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -0,0 +1,148 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+import re
+import time
+
+import librosa
+from PIL import Image
+
+from vllm_omni.diffusion.data import DiffusionParallelConfig
+from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+
+IMAGE_EXTENSIONS = {".jpg", ".jpeg", ".png", ".bmp", ".webp", ".gif"}
+AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
+    parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
+    parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
+    parser.add_argument("--prompt", default=None, help="Text prompt.")
+
+    parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
+    parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
+    parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
+
+    parser.add_argument("--height", type=int, default=704, help="Video height.")
+    parser.add_argument("--width", type=int, default=1280, help="Video width.")
+    parser.add_argument("--num-inference-steps", type=int, default=45, help="Sampling steps.")
+    parser.add_argument("--solver-name", default="unipc", help="Solver name: unipc|dpm++|euler.")
+    parser.add_argument("--shift", type=float, default=5.0, help="Scheduler shift.")
+    parser.add_argument("--seed", type=int, default=103, help="Random seed for reproducible generation.")
+    parser.add_argument(
+        "--cfg-parallel-size",
+        type=int,
+        default=1,
+        choices=[1, 2],
+        help="Number of GPUs used for classifier free guidance parallel size.",
+    )
+    parser.add_argument(
+        "--video-negative-prompt",
+        default="jitter, bad hands, blur, distortion",
+        help="Negative prompt for video.",
+    )
+    parser.add_argument(
+        "--audio-negative-prompt",
+        default="robotic, muffled, echo, distorted",
+        help="Negative prompt for audio.",
+    )
+    parser.add_argument("--output", default="dreamid_output.mp4", help="Output video path.")
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        default=False,
+        help="Enable CPU offloading for diffusion models.",
+    )
+    return parser.parse_args()
+
+
+def load_image_and_audio(image_paths, audio_paths):
+    image = []
+    audio = []
+
+    for path in image_paths:
+        with Image.open(path) as img:
+            img = img.convert("RGB")
+            image.append(img)
+
+    for path in audio_paths:
+        audio_array, sr = librosa.load(path, sr=16000)
+        audio_array = audio_array[int(sr * 1) : int(sr * 3)]
+        audio.append(audio_array)
+    return image, audio
+
+
+def main() -> None:
+    args = parse_args()
+    if args.prompt is None and args.prompt_file is None:
+        raise ValueError("Either --prompt or --prompt-file must be provided.")
+
+    text_prompt = args.prompt
+    if args.prompt_file:
+        import json
+
+        with open(args.prompt_file) as f:
+            text_prompt = json.load(f)
+            text_prompt = re.sub(
+                r"\[SPEAKER_TIMESTAMPS_START\].*?\[SPEAKER_TIMESTAMPS_END\]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(
+                r"\[AUDIO_DESCRIPTION_START].*?\[AUDIO_DESCRIPTION_END]", "", text_prompt, flags=re.DOTALL
+            ).strip()
+            text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
+            text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
+
+    image, audio = load_image_and_audio(args.image_path, args.audio_path)
+
+    prompt = {
+        "prompt": text_prompt,
+        "video_negative_prompt": args.video_negative_prompt,
+        "audio_negative_prompt": args.audio_negative_prompt,
+        "multi_modal_data": {"image": image, "audio": audio},
+    }
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+        seed=args.seed,
+        extra_args={
+            "solver_name": args.solver_name,
+            "shift": args.shift,
+        },
+    )
+
+    parallel_config = DiffusionParallelConfig(
+        cfg_parallel_size=args.cfg_parallel_size,
+    )
+
+    omni = Omni(
+        model=args.model,
+        parallel_config=parallel_config,
+        model_type=args.model_type,
+        enable_cpu_offload=args.enable_cpu_offload,
+    )
+    start = time.perf_counter()
+    outputs = omni.generate(prompt, sampling_params)
+    elapsed = time.perf_counter() - start
+
+    if not outputs:
+        raise RuntimeError("No output returned from DreamID-Omni.")
+    output = outputs[0].request_output
+    generated_video = output[0].images[0][0]
+    generated_audio = output[0].images[0][1]
+    try:
+        from dreamid_omni.utils.io_utils import save_video
+    except Exception as e:
+        raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}")
+    output_path = args.output
+    save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000)
+    print(f"Saved generated video to {output_path}")
+    print(f"Total time: {elapsed:.2f}s")
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/fish_speech/README.md

@@ -0,0 +1,172 @@
+# Fish Speech S2 Pro
+
+## Model
+
+| Model | Description |
+|-------|-------------|
+| `fishaudio/s2-pro` | Fish Speech S2 Pro -- 4B dual-AR TTS model with DAC codec (44.1 kHz) |
+
+## Gradio Demo
+
+An interactive Gradio demo is available with text-to-speech synthesis, voice cloning, and streaming support.
+
+```bash
+# Option 1: Launch server + Gradio together
+./run_gradio_demo.sh
+
+# Option 2: If server is already running
+python gradio_demo.py --api-base http://localhost:8091
+```
+
+Then open http://localhost:7860 in your browser.
+
+Features:
+
+- Text-to-speech synthesis
+- Voice cloning from uploaded audio or URL
+- Streaming mode (progressive PCM playback)
+
+## Launch the Server
+
+```bash
+vllm-omni serve fishaudio/s2-pro \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
+    --omni \
+    --port 8091 \
+    --trust-remote-code \
+    --enforce-eager \
+    --gpu-memory-utilization 0.9
+```
+
+Or use the convenience script:
+
+```bash
+./run_server.sh
+```
+
+## Send TTS Request
+
+### Using curl
+
+```bash
+# Basic TTS
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "response_format": "wav"
+    }' --output output.wav
+```
+
+### Voice Cloning
+
+Provide a reference audio (URL or base64 data URL) and its transcript:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, this is a cloned voice.",
+        "voice": "default",
+        "ref_audio": "https://example.com/reference.wav",
+        "ref_text": "Transcript of the reference audio."
+    }' --output cloned.wav
+```
+
+### Using Python
+
+```python
+import httpx
+
+# Basic TTS
+response = httpx.post(
+    "http://localhost:8091/v1/audio/speech",
+    json={
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "response_format": "wav",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+### Using the CLI Client
+
+```bash
+cd examples/online_serving/fish_speech
+
+# Basic TTS
+python speech_client.py --text "Hello, how are you?"
+
+# Voice cloning
+python speech_client.py \
+    --text "Hello, this is a cloned voice." \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Transcript of the reference audio."
+
+# Streaming PCM output
+python speech_client.py --text "Hello world" --stream --output output.pcm
+```
+
+The CLI client supports:
+
+- `--api-base`: API base URL (default: `http://localhost:8091`)
+- `--model` (or `-m`): Model name (default: `fishaudio/s2-pro`)
+- `--text`: Text to synthesize (required)
+- `--ref-audio`: Reference audio for voice cloning (local path or URL)
+- `--ref-text`: Transcript of the reference audio
+- `--stream`: Enable streaming (PCM output)
+- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
+- `--output` (or `-o`): Output file path
+
+## Streaming
+
+Set `stream=true` with `response_format="pcm"` to receive raw PCM audio chunks as they are decoded:
+
+```bash
+curl -X POST http://localhost:8091/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "default",
+        "stream": true,
+        "response_format": "pcm"
+    }' --no-buffer | play -t raw -r 44100 -e signed -b 16 -c 1 -
+```
+
+**Note:** Fish Speech outputs at 44.1 kHz (unlike Qwen3-TTS which outputs at 24 kHz).
+
+## API Parameters
+
+Fish Speech uses the same `/v1/audio/speech` endpoint as Qwen3-TTS. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full parameter documentation.
+
+Key parameters for Fish Speech:
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | **required** | Text to synthesize |
+| `voice` | string | "default" | Voice name (use "default" for Fish Speech) |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `ref_audio` | string | null | Reference audio URL or base64 data URL for voice cloning |
+| `ref_text` | string | null | Transcript of reference audio (required for voice cloning) |
+| `max_new_tokens` | int | 4096 | Maximum tokens to generate |
+| `stream` | bool | false | Stream raw PCM chunks |
+
+## Prerequisites
+
+Install the `fish-speech` package for the DAC codec:
+
+```bash
+pip install fish-speech
+```
+
+## Troubleshooting
+
+1. **No audio output**: Make sure the `fish-speech` package is installed for the DAC decoder
+2. **Connection refused**: Ensure the server is running on the correct port
+3. **Flashinfer version mismatch**: Set `FLASHINFER_DISABLE_VERSION_CHECK=1` if you see version warnings
+4. **Out of memory**: Reduce `--gpu-memory-utilization` or use a GPU with more VRAM

examples/online_serving/fish_speech/gradio_demo.py

@@ -0,0 +1,272 @@
+"""Gradio demo for Fish Speech S2 Pro online serving via /v1/audio/speech API.
+
+Supports:
+  - Text-to-speech synthesis
+  - Voice cloning from reference audio (upload or URL)
+  - Streaming (progressive PCM chunks) and non-streaming modes
+
+Usage:
+    # Start the server first (see run_server.sh), then:
+    python gradio_demo.py --api-base http://localhost:8091
+
+    # Or use run_gradio_demo.sh to start both server and demo together.
+"""
+
+import argparse
+import base64
+import io
+
+import gradio as gr
+import httpx
+import numpy as np
+import soundfile as sf
+
+PCM_SAMPLE_RATE = 44100
+DEFAULT_API_BASE = "http://localhost:8091"
+
+
+def encode_audio_to_base64(audio_data: tuple) -> str:
+    """Encode Gradio audio input (sample_rate, numpy_array) to base64 data URL."""
+    sample_rate, audio_np = audio_data
+    if audio_np.dtype != np.int16:
+        if audio_np.dtype in (np.float32, np.float64):
+            audio_np = np.clip(audio_np, -1.0, 1.0)
+            audio_np = (audio_np * 32767).astype(np.int16)
+        else:
+            audio_np = audio_np.astype(np.int16)
+    buf = io.BytesIO()
+    sf.write(buf, audio_np, sample_rate, format="WAV")
+    wav_b64 = base64.b64encode(buf.getvalue()).decode("utf-8")
+    return f"data:audio/wav;base64,{wav_b64}"
+
+
+def build_payload(
+    text: str,
+    ref_audio: tuple | None,
+    ref_audio_url: str,
+    ref_text: str,
+    response_format: str = "wav",
+    stream: bool = False,
+) -> dict:
+    """Build the /v1/audio/speech request payload."""
+    if not text or not text.strip():
+        raise gr.Error("Please enter text to synthesize.")
+
+    payload: dict = {
+        "input": text.strip(),
+        "voice": "default",
+        "response_format": "pcm" if stream else response_format,
+        "stream": stream,
+    }
+
+    # Voice cloning: ref_audio takes priority over URL.
+    ref_url_stripped = ref_audio_url.strip() if ref_audio_url else ""
+    if ref_audio is not None:
+        payload["ref_audio"] = encode_audio_to_base64(ref_audio)
+    elif ref_url_stripped:
+        payload["ref_audio"] = ref_url_stripped
+
+    if "ref_audio" in payload:
+        if not ref_text or not ref_text.strip():
+            raise gr.Error("Voice cloning requires a transcript of the reference audio.")
+        payload["ref_text"] = ref_text.strip()
+
+    return payload
+
+
+def generate_speech(api_base: str, text: str, ref_audio, ref_audio_url, ref_text, response_format):
+    """Non-streaming: call /v1/audio/speech and return full audio."""
+    payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=False)
+
+    try:
+        with httpx.Client(timeout=300.0) as client:
+            resp = client.post(
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            )
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+    if resp.status_code != 200:
+        raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+
+    content_type = resp.headers.get("content-type", "")
+    if "application/json" in content_type:
+        try:
+            raise gr.Error(f"Server error: {resp.json()}")
+        except ValueError:
+            pass
+
+    try:
+        if response_format == "pcm":
+            audio_np = np.frombuffer(resp.content, dtype=np.int16).astype(np.float32) / 32767.0
+            return (PCM_SAMPLE_RATE, audio_np)
+        audio_np, sample_rate = sf.read(io.BytesIO(resp.content))
+        if audio_np.ndim > 1:
+            audio_np = audio_np[:, 0]
+        return (sample_rate, audio_np.astype(np.float32))
+    except Exception as e:
+        raise gr.Error(f"Failed to decode audio: {e}")
+
+
+def generate_speech_stream(api_base, text, ref_audio, ref_audio_url, ref_text, response_format, chunk_seconds=0.5):
+    """Streaming: yield PCM chunks for progressive playback."""
+    payload = build_payload(text, ref_audio, ref_audio_url, ref_text, response_format, stream=True)
+
+    min_chunk_samples = int(PCM_SAMPLE_RATE * chunk_seconds)
+    prebuffer_chunks = 2
+    pending: list[np.ndarray] = []
+    pending_len = 0
+    chunks_yielded = 0
+    prebuffer: list[np.ndarray] = []
+
+    try:
+        leftover = b""
+        with httpx.Client(timeout=300.0) as client:
+            with client.stream(
+                "POST",
+                f"{api_base}/v1/audio/speech",
+                json=payload,
+                headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"},
+            ) as resp:
+                if resp.status_code != 200:
+                    resp.read()
+                    raise gr.Error(f"Server error ({resp.status_code}): {resp.text}")
+                for chunk in resp.iter_bytes():
+                    if not chunk:
+                        continue
+                    raw = leftover + chunk
+                    usable = len(raw) - (len(raw) % 2)
+                    leftover = raw[usable:]
+                    if usable == 0:
+                        continue
+                    float_samples = np.frombuffer(raw[:usable], dtype=np.int16).copy().astype(np.float32) / 32767.0
+                    pending.append(float_samples)
+                    pending_len += len(float_samples)
+                    if pending_len >= min_chunk_samples:
+                        audio_chunk = np.concatenate(pending)
+                        pending.clear()
+                        pending_len = 0
+                        if chunks_yielded < prebuffer_chunks:
+                            prebuffer.append(audio_chunk)
+                            chunks_yielded += 1
+                            if chunks_yielded == prebuffer_chunks:
+                                yield (PCM_SAMPLE_RATE, np.concatenate(prebuffer))
+                                prebuffer.clear()
+                        else:
+                            yield (PCM_SAMPLE_RATE, audio_chunk)
+        remaining = prebuffer + pending
+        if remaining:
+            yield (PCM_SAMPLE_RATE, np.concatenate(remaining))
+    except httpx.TimeoutException:
+        raise gr.Error("Request timed out. The server may be busy.")
+    except httpx.ConnectError:
+        raise gr.Error(f"Cannot connect to server at {api_base}. Is the server running?")
+
+
+def on_stream_change(stream: bool):
+    """Lock format to PCM when streaming."""
+    if stream:
+        return gr.update(value="pcm", interactive=False)
+    return gr.update(value="wav", interactive=True)
+
+
+def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
+    """Build the Gradio interface."""
+    with gr.Blocks(title="Fish Speech S2 Pro Demo") as demo:
+        gr.Markdown("# Fish Speech S2 Pro - Text to Speech")
+        gr.Markdown(f"**Server:** `{api_base}` | **Model:** fishaudio/s2-pro | **Output:** 44.1kHz")
+
+        with gr.Row():
+            with gr.Column(scale=3):
+                text_input = gr.Textbox(
+                    label="Text to Synthesize",
+                    placeholder="Enter text here...",
+                    lines=4,
+                )
+
+                with gr.Accordion("Voice Cloning (optional)", open=False):
+                    gr.Markdown(
+                        "Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
+                    )
+                    ref_audio = gr.Audio(
+                        label="Reference Audio",
+                        type="numpy",
+                        sources=["upload", "microphone"],
+                    )
+                    ref_audio_url = gr.Textbox(
+                        label="Reference Audio URL (alternative to upload)",
+                        placeholder="https://example.com/reference.wav",
+                        lines=1,
+                    )
+                    ref_text = gr.Textbox(
+                        label="Reference Audio Transcript (required for cloning)",
+                        placeholder="Exact transcript of the reference audio...",
+                        lines=2,
+                    )
+
+                with gr.Row():
+                    response_format = gr.Dropdown(
+                        choices=["wav", "mp3", "flac", "pcm"],
+                        value="wav",
+                        label="Audio Format",
+                        scale=1,
+                    )
+                    stream_checkbox = gr.Checkbox(
+                        label="Stream output",
+                        value=False,
+                        info="Progressive PCM streaming",
+                        scale=1,
+                    )
+
+                generate_btn = gr.Button("Generate Speech", variant="primary", size="lg")
+
+            with gr.Column(scale=2):
+                audio_output = gr.Audio(
+                    label="Generated Audio",
+                    interactive=False,
+                    streaming=True,
+                    autoplay=True,
+                )
+                gr.Markdown(
+                    "### About\n"
+                    "- **Fish Speech S2 Pro** by FishAudio: 4B dual-AR model\n"
+                    "- **Voice cloning**: upload 10-30s reference + transcript\n"
+                    "- **Streaming**: real-time PCM output\n"
+                    "- **44.1kHz** output via DAC codec"
+                )
+
+        stream_checkbox.change(fn=on_stream_change, inputs=[stream_checkbox], outputs=[response_format])
+
+        all_inputs = [text_input, ref_audio, ref_audio_url, ref_text, response_format]
+
+        def dispatch(stream_enabled, *args):
+            if stream_enabled:
+                yield from generate_speech_stream(api_base, *args, chunk_seconds=stream_chunk_seconds)
+            else:
+                yield generate_speech(api_base, *args)
+
+        generate_btn.click(fn=dispatch, inputs=[stream_checkbox] + all_inputs, outputs=[audio_output])
+        demo.queue()
+    return demo
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Gradio demo for Fish Speech S2 Pro")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help=f"API base URL (default: {DEFAULT_API_BASE})")
+    parser.add_argument("--host", default="0.0.0.0", help="Gradio host (default: 0.0.0.0)")
+    parser.add_argument("--port", type=int, default=7860, help="Gradio port (default: 7860)")
+    parser.add_argument("--share", action="store_true", help="Share publicly")
+    parser.add_argument("--stream-chunk-seconds", type=float, default=0.5, help="Seconds per streaming chunk")
+    args = parser.parse_args()
+
+    print(f"Connecting to vLLM server at: {args.api_base}")
+    demo = build_interface(args.api_base, stream_chunk_seconds=args.stream_chunk_seconds)
+    demo.launch(server_name=args.host, server_port=args.port, share=args.share)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/fish_speech/run_gradio_demo.sh

@@ -0,0 +1,48 @@
+#!/bin/bash
+# Launch Fish Speech S2 Pro server + Gradio demo together.
+#
+# Usage:
+#   ./run_gradio_demo.sh
+#   CUDA_VISIBLE_DEVICES=0 PORT=8091 GRADIO_PORT=7860 ./run_gradio_demo.sh
+
+set -e
+
+MODEL="${MODEL:-fishaudio/s2-pro}"
+PORT="${PORT:-8091}"
+GRADIO_PORT="${GRADIO_PORT:-7860}"
+SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
+REPO_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+
+echo "Starting Fish Speech S2 Pro server (port $PORT)..."
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm-omni serve "$MODEL" \
+    --stage-configs-path "$REPO_ROOT/vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml" \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni &
+SERVER_PID=$!
+
+cleanup() {
+    echo "Stopping server (PID $SERVER_PID)..."
+    kill $SERVER_PID 2>/dev/null
+    wait $SERVER_PID 2>/dev/null
+}
+trap cleanup EXIT
+
+# Wait for server to be ready.
+echo "Waiting for server to start..."
+for i in $(seq 1 120); do
+    if curl -s "http://localhost:$PORT/health" > /dev/null 2>&1; then
+        echo "Server ready."
+        break
+    fi
+    sleep 2
+done
+
+echo "Starting Gradio demo (port $GRADIO_PORT)..."
+python "$SCRIPT_DIR/gradio_demo.py" \
+    --api-base "http://localhost:$PORT" \
+    --port "$GRADIO_PORT"

examples/online_serving/fish_speech/run_server.sh

@@ -0,0 +1,23 @@
+#!/bin/bash
+# Launch vLLM-Omni server for Fish Speech S2 Pro
+#
+# Usage:
+#   ./run_server.sh
+#   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
+
+set -e
+
+MODEL="${MODEL:-fishaudio/s2-pro}"
+PORT="${PORT:-8091}"
+
+echo "Starting Fish Speech S2 Pro server with model: $MODEL"
+
+FLASHINFER_DISABLE_VERSION_CHECK=1 \
+vllm-omni serve "$MODEL" \
+    --stage-configs-path vllm_omni/model_executor/stage_configs/fish_speech_s2_pro.yaml \
+    --host 0.0.0.0 \
+    --port "$PORT" \
+    --gpu-memory-utilization 0.9 \
+    --trust-remote-code \
+    --enforce-eager \
+    --omni

examples/online_serving/fish_speech/speech_client.py

@@ -0,0 +1,129 @@
+"""Client for Fish Speech S2 Pro via /v1/audio/speech endpoint.
+
+Examples:
+    # Basic TTS
+    python speech_client.py --text "Hello, how are you?"
+
+    # Voice cloning
+    python speech_client.py --text "Hello, how are you?" \
+        --ref-audio ref.wav --ref-text "This is the reference transcript."
+
+    # Streaming PCM output
+    python speech_client.py --text "Hello world" --stream --output output.pcm
+"""
+
+import argparse
+import base64
+import os
+
+import httpx
+
+DEFAULT_API_BASE = "http://localhost:8091"
+DEFAULT_API_KEY = "EMPTY"
+
+
+def encode_audio_to_base64(audio_path: str) -> str:
+    """Encode a local audio file to base64 data URL."""
+    if not os.path.exists(audio_path):
+        raise FileNotFoundError(f"Audio file not found: {audio_path}")
+    ext = audio_path.lower().rsplit(".", 1)[-1]
+    mime_map = {"wav": "audio/wav", "mp3": "audio/mpeg", "flac": "audio/flac", "ogg": "audio/ogg"}
+    mime_type = mime_map.get(ext, "audio/wav")
+    with open(audio_path, "rb") as f:
+        audio_b64 = base64.b64encode(f.read()).decode("utf-8")
+    return f"data:{mime_type};base64,{audio_b64}"
+
+
+def run_tts(args) -> None:
+    """Generate speech via /v1/audio/speech API."""
+    payload = {
+        "model": args.model,
+        "input": args.text,
+        "voice": "default",
+        "response_format": args.response_format,
+    }
+
+    # Voice cloning parameters.
+    if args.ref_audio:
+        if args.ref_audio.startswith(("http://", "https://")):
+            payload["ref_audio"] = args.ref_audio
+        else:
+            payload["ref_audio"] = encode_audio_to_base64(args.ref_audio)
+    if args.ref_text:
+        payload["ref_text"] = args.ref_text
+
+    if args.stream:
+        payload["stream"] = True
+        payload["response_format"] = "pcm"
+
+    print(f"Model: {args.model}")
+    print(f"Text: {args.text}")
+    if args.ref_audio:
+        print(f"Voice cloning: ref_audio={args.ref_audio}, ref_text={args.ref_text}")
+    print("Generating audio...")
+
+    api_url = f"{args.api_base}/v1/audio/speech"
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": f"Bearer {args.api_key}",
+    }
+
+    if args.stream:
+        output_path = args.output or "output.pcm"
+        with httpx.Client(timeout=300.0) as client:
+            with client.stream("POST", api_url, json=payload, headers=headers) as resp:
+                if resp.status_code != 200:
+                    print(f"Error: {resp.status_code}")
+                    print(resp.read().decode())
+                    return
+                total_bytes = 0
+                with open(output_path, "wb") as f:
+                    for chunk in resp.iter_bytes():
+                        f.write(chunk)
+                        total_bytes += len(chunk)
+                print(f"Streamed {total_bytes} bytes to: {output_path}")
+    else:
+        with httpx.Client(timeout=300.0) as client:
+            response = client.post(api_url, json=payload, headers=headers)
+
+        if response.status_code != 200:
+            print(f"Error: {response.status_code}")
+            print(response.text)
+            return
+
+        try:
+            text = response.content.decode("utf-8")
+            if text.startswith('{"error"'):
+                print(f"Error: {text}")
+                return
+        except UnicodeDecodeError:
+            pass
+
+        output_path = args.output or "output.wav"
+        with open(output_path, "wb") as f:
+            f.write(response.content)
+        print(f"Audio saved to: {output_path}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Fish Speech S2 Pro TTS client")
+    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
+    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
+    parser.add_argument("--model", "-m", default="fishaudio/s2-pro", help="Model name")
+    parser.add_argument("--text", required=True, help="Text to synthesize")
+    parser.add_argument("--ref-audio", default=None, help="Reference audio for voice cloning (path or URL)")
+    parser.add_argument("--ref-text", default=None, help="Transcript of reference audio")
+    parser.add_argument("--stream", action="store_true", help="Enable streaming (PCM output)")
+    parser.add_argument(
+        "--response-format",
+        default="wav",
+        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
+        help="Audio format (default: wav)",
+    )
+    parser.add_argument("--output", "-o", default=None, help="Output file path")
+    args = parser.parse_args()
+    run_tts(args)
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/image_to_video/README.md

@@ -26,6 +26,35 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+Local file storage behavior can be controlled via the following environment variables:
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -35,7 +64,7 @@ The script allows overriding:
 bash run_curl_image_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -X POST http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=A bear playing with yarn, smooth motion" \
   -F "negative_prompt=low quality, blurry, static" \
@@ -49,7 +78,23 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "guidance_scale_2=1.0" \
   -F "boundary_ratio=0.875" \
   -F "flow_shift=12.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_i2v_output.mp4
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_i2v_output.mp4
 ```
 
 ## Request Format
@@ -63,6 +108,18 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "input_reference=@/path/to/qwen-bear.png"
 ```
 
+### Alternative JSON-Safe Reference Input
+
+Use `image_reference` when you want to pass a URL or JSON-safe image reference
+instead of uploading a file. Do not send `input_reference` and
+`image_reference` together.
+
+```bash
+curl -X POST http://localhost:8091/v1/videos \
+  -F "prompt=A bear playing with yarn, smooth motion" \
+  -F 'image_reference={"image_url":"https://example.com/qwen-bear.png"}'
+```
+
 ### Generation with Parameters
 
 ```bash
@@ -81,3 +138,60 @@ curl -X POST http://localhost:8091/v1/videos \
   -F "flow_shift=12.0" \
   -F "seed=42"
 ```
+
+## Create Response Format
+
+`POST /v1/videos` returns a job record, not inline base64 video data.
+
+```json
+{
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-I2V-A14B-Diffusers",
+  "prompt": "A bear playing with yarn, smooth motion",
+  "created_at": 1234567890
+}
+```
+
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
+
+```bash
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_i2v_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+```

examples/online_serving/image_to_video/run_curl_image_to_video.sh

@@ -1,33 +1,75 @@
 #!/bin/bash
-# Wan2.2 image-to-video curl example (OpenAI-style multipart)
+# Wan2.2 image-to-video curl example using the async video job API.
+
+set -euo pipefail
 
 INPUT_IMAGE="${INPUT_IMAGE:-../../offline_inference/image_to_video/qwen-bear.png}"
+BASE_URL="${BASE_URL:-http://localhost:8099}"
 OUTPUT_PATH="${OUTPUT_PATH:-wan22_i2v_output.mp4}"
 NEGATIVE_PROMPT="${NEGATIVE_PROMPT:-}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
 
 if [ ! -f "$INPUT_IMAGE" ]; then
     echo "Input image not found: $INPUT_IMAGE"
     exit 1
 fi
 
-NEGATIVE_PROMPT_FLAG=""
-if [ -n "$NEGATIVE_PROMPT" ]; then
-    NEGATIVE_PROMPT_FLAG="-F negative_prompt=${NEGATIVE_PROMPT}"
+create_cmd=(
+  curl -sS -X POST "${BASE_URL}/v1/videos"
+  -H "Accept: application/json"
+  -F "prompt=A bear playing with yarn, smooth motion"
+  -F "input_reference=@${INPUT_IMAGE}"
+  -F "seconds=2"
+  -F "size=832x480"
+  -F "fps=16"
+  -F "num_inference_steps=40"
+  -F "guidance_scale=1.0"
+  -F "guidance_scale_2=1.0"
+  -F "boundary_ratio=0.875"
+  -F "flow_shift=12.0"
+  -F "seed=42"
+)
+
+if [ -n "${NEGATIVE_PROMPT}" ]; then
+  create_cmd+=(-F "negative_prompt=${NEGATIVE_PROMPT}")
 fi
 
-curl -X POST http://localhost:8099/v1/videos \
-  -H "Accept: application/json" \
-  -F "prompt=A bear playing with yarn, smooth motion" \
-  $NEGATIVE_PROMPT_FLAG \
-  -F "input_reference=@${INPUT_IMAGE}" \
-  -F "seconds=2" \
-  -F "size=832x480" \
-  -F "fps=16" \
-  -F "num_inference_steps=40" \
-  -F "guidance_scale=1.0" \
-  -F "guidance_scale_2=1.0" \
-  -F "boundary_ratio=0.875" \
-  -F "flow_shift=12.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
+create_response="$("${create_cmd[@]}")"
+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/qwen3_tts/README.md

@@ -349,6 +349,42 @@ curl -X POST http://localhost:8091/v1/audio/speech \
 - `speed` adjustment is not supported when streaming.
 - Requires the server stage config to have `async_chunk: true` (default in `qwen3_tts.yaml`).
 
+## Streaming Text Input (WebSocket)
+
+The `/v1/audio/speech/stream` WebSocket endpoint accepts text incrementally, buffers and splits it at sentence boundaries, and generates audio per sentence.
+
+When `stream_audio=true`, each sentence is emitted as `audio.start`, one or more binary PCM frames, and `audio.done`.
+
+### Quick Start
+
+```bash
+python streaming_speech_client.py \
+    --text "Hello world. How are you? I am fine."
+
+python streaming_speech_client.py \
+    --text "Hello world. How are you? I am fine." \
+    --simulate-stt --stt-delay 0.1
+```
+
+### WebSocket Protocol
+
+Client -> Server:
+
+```jsonc
+{"type": "session.config", "voice": "Vivian", "task_type": "CustomVoice", "language": "Auto", "split_granularity": "sentence", "stream_audio": true, "response_format": "pcm"}
+{"type": "input.text", "text": "Hello, how are you? "}
+{"type": "input.done"}
+```
+
+Server -> Client:
+
+```jsonc
+{"type": "audio.start", "sentence_index": 0, "sentence_text": "Hello, how are you?", "format": "pcm", "sample_rate": 24000}
+// binary PCM frame(s)
+{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}
+{"type": "session.done", "total_sentences": 1}
+```
+
 ## Limitations
 
 - **Single request**: Batch processing is not yet optimized for online serving.

examples/online_serving/qwen3_tts/streaming_speech_client.py

@@ -0,0 +1,249 @@
+"""WebSocket client for streaming text-input TTS.
+
+Connects to the /v1/audio/speech/stream endpoint, sends text incrementally
+(simulating real-time STT output), and saves per-sentence audio files.
+
+Usage:
+    # Send full text at once
+    python streaming_speech_client.py --text "Hello world. How are you? I am fine."
+
+    # Simulate STT: send text word-by-word with delay
+    python streaming_speech_client.py \
+        --text "Hello world. How are you? I am fine." \
+        --simulate-stt --stt-delay 0.1
+
+    # VoiceDesign task
+    python streaming_speech_client.py \
+        --text "Today is a great day. The weather is nice." \
+        --task-type VoiceDesign \
+        --instructions "A cheerful young female voice"
+
+    # Base task (voice cloning)
+    python streaming_speech_client.py \
+        --text "Hello world. How are you?" \
+        --task-type Base \
+        --ref-audio /path/to/reference.wav \
+        --ref-text "Transcript of reference audio"
+
+Requirements:
+    pip install websockets
+"""
+
+import argparse
+import asyncio
+import json
+import os
+
+try:
+    import websockets
+except ImportError:
+    print("Please install websockets: pip install websockets")
+    raise SystemExit(1)
+
+
+async def stream_tts(
+    url: str,
+    text: str,
+    config: dict,
+    output_dir: str,
+    simulate_stt: bool = False,
+    stt_delay: float = 0.1,
+) -> None:
+    """Connect to the streaming TTS endpoint and process audio responses."""
+    os.makedirs(output_dir, exist_ok=True)
+
+    async with websockets.connect(url) as ws:
+        # 1. Send session config
+        config_msg = {"type": "session.config", **config}
+        await ws.send(json.dumps(config_msg))
+        print(f"Sent session config: {config}")
+
+        # 2. Send text (either all at once or word-by-word)
+        async def send_text():
+            if simulate_stt:
+                words = text.split(" ")
+                for i, word in enumerate(words):
+                    chunk = word + (" " if i < len(words) - 1 else "")
+                    await ws.send(
+                        json.dumps(
+                            {
+                                "type": "input.text",
+                                "text": chunk,
+                            }
+                        )
+                    )
+                    print(f"  Sent: {chunk!r}")
+                    await asyncio.sleep(stt_delay)
+            else:
+                await ws.send(
+                    json.dumps(
+                        {
+                            "type": "input.text",
+                            "text": text,
+                        }
+                    )
+                )
+                print(f"Sent full text: {text!r}")
+
+            # 3. Signal end of input
+            await ws.send(json.dumps({"type": "input.done"}))
+            print("Sent input.done")
+
+        # Run sender and receiver concurrently
+        sender_task = asyncio.create_task(send_text())
+
+        response_format = config.get("response_format", "wav")
+        current_sentence_index = 0
+        current_chunks: list[bytes] = []
+
+        try:
+            while True:
+                message = await ws.recv()
+
+                if isinstance(message, bytes):
+                    current_chunks.append(message)
+                    print(f"  Received audio chunk for sentence {current_sentence_index}: {len(message)} bytes")
+                else:
+                    # JSON frame
+                    msg = json.loads(message)
+                    msg_type = msg.get("type")
+
+                    if msg_type == "audio.start":
+                        current_sentence_index = msg["sentence_index"]
+                        current_chunks = []
+                        print(f"  [sentence {msg['sentence_index']}] Generating: {msg['sentence_text']!r}")
+                    elif msg_type == "audio.done":
+                        filename = os.path.join(
+                            output_dir,
+                            f"sentence_{msg['sentence_index']:03d}.{response_format}",
+                        )
+                        with open(filename, "wb") as f:
+                            f.write(b"".join(current_chunks))
+                        print(
+                            f"  [sentence {msg['sentence_index']}] Done"
+                            f" bytes={msg.get('total_bytes', len(b''.join(current_chunks)))}"
+                            f" error={msg.get('error', False)}"
+                            f" -> {filename}"
+                        )
+                        current_chunks = []
+                    elif msg_type == "session.done":
+                        print(f"\nSession complete: {msg['total_sentences']} sentence(s) generated")
+                        break
+                    elif msg_type == "error":
+                        print(f"  ERROR: {msg['message']}")
+                    else:
+                        print(f"  Unknown message: {msg}")
+        finally:
+            sender_task.cancel()
+            try:
+                await sender_task
+            except asyncio.CancelledError:
+                pass  # Task cancellation is expected during shutdown
+
+    print(f"\nAudio files saved to: {output_dir}/")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Streaming text-input TTS client")
+    parser.add_argument(
+        "--url",
+        default="ws://localhost:8000/v1/audio/speech/stream",
+        help="WebSocket endpoint URL",
+    )
+    parser.add_argument(
+        "--text",
+        required=True,
+        help="Text to synthesize",
+    )
+    parser.add_argument(
+        "--output-dir",
+        default="streaming_tts_output",
+        help="Directory to save audio files (default: streaming_tts_output)",
+    )
+
+    # Session config options
+    parser.add_argument("--model", default=None, help="Model name")
+    parser.add_argument("--voice", default="Vivian", help="Speaker voice")
+    parser.add_argument(
+        "--task-type",
+        default="CustomVoice",
+        choices=["CustomVoice", "VoiceDesign", "Base"],
+        help="TTS task type",
+    )
+    parser.add_argument("--language", default="Auto", help="Language")
+    parser.add_argument("--instructions", default=None, help="Voice style instructions")
+    parser.add_argument(
+        "--response-format",
+        default="wav",
+        choices=["wav", "pcm", "flac", "mp3", "aac", "opus"],
+        help="Audio format",
+    )
+    parser.add_argument(
+        "--stream-audio",
+        action="store_true",
+        help="Receive one or more PCM chunks per sentence (requires --response-format pcm)",
+    )
+    parser.add_argument("--speed", type=float, default=1.0, help="Playback speed (0.25-4.0)")
+    parser.add_argument("--max-new-tokens", type=int, default=None, help="Max tokens")
+
+    # Base task options
+    parser.add_argument("--ref-audio", default=None, help="Reference audio")
+    parser.add_argument("--ref-text", default=None, help="Reference text")
+    parser.add_argument(
+        "--x-vector-only-mode",
+        action="store_true",
+        default=False,
+        help="Speaker embedding only mode",
+    )
+
+    # STT simulation
+    parser.add_argument(
+        "--simulate-stt",
+        action="store_true",
+        help="Simulate STT by sending text word-by-word",
+    )
+    parser.add_argument(
+        "--stt-delay",
+        type=float,
+        default=0.1,
+        help="Delay between words in STT simulation (seconds)",
+    )
+
+    args = parser.parse_args()
+
+    # Build session config (only include non-None values)
+    config = {}
+    for key in [
+        "model",
+        "voice",
+        "task_type",
+        "language",
+        "instructions",
+        "response_format",
+        "speed",
+        "max_new_tokens",
+        "ref_audio",
+        "ref_text",
+    ]:
+        val = getattr(args, key.replace("-", "_"), None)
+        if val is not None:
+            config[key] = val
+    if args.stream_audio:
+        config["stream_audio"] = True
+    if args.x_vector_only_mode:
+        config["x_vector_only_mode"] = True
+
+    asyncio.run(
+        stream_tts(
+            url=args.url,
+            text=args.text,
+            config=config,
+            output_dir=args.output_dir,
+            simulate_stt=args.simulate_stt,
+            stt_delay=args.stt_delay,
+        )
+    )
+
+
+if __name__ == "__main__":
+    main()

examples/online_serving/text_to_video/README.md

@@ -26,6 +26,35 @@ The script allows overriding:
 - `CACHE_BACKEND` (default: `none`)
 - `ENABLE_CACHE_DIT_SUMMARY` (default: `0`)
 
+## Async Job Behavior
+
+`POST /v1/videos` is asynchronous. It creates a video job and immediately
+returns metadata like the job ID and initial `queued` status. To get the final
+artifact, poll the job status and then download the completed file from the
+content endpoint.
+
+The main endpoints are:
+- `POST /v1/videos`: create a video generation job
+- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
+- `GET /v1/videos`: list stored video jobs
+- `GET /v1/videos/{video_id}/content`: download the generated video file
+- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
+
+## Storage
+
+Generated video files are stored on local disk by the async video API.
+Local file storage behavior can be controlled via the following environment variables:
+
+- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
+- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
+
+Example:
+
+```bash
+export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
+export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
+```
+
 ## API Calls
 
 ### Method 1: Using curl
@@ -35,7 +64,7 @@ The script allows overriding:
 bash run_curl_text_to_video.sh
 
 # Or execute directly (OpenAI-style multipart)
-curl -s http://localhost:8091/v1/videos \
+create_response=$(curl -s http://localhost:8091/v1/videos \
   -H "Accept: application/json" \
   -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   -F "width=832" \
@@ -47,7 +76,24 @@ curl -s http://localhost:8091/v1/videos \
   -F "guidance_scale=4.0" \
   -F "guidance_scale_2=4.0" \
   -F "boundary_ratio=0.875" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+  -F "flow_shift=5.0" \
+  -F "seed=42")
+
+video_id=$(echo "$create_response" | jq -r '.id')
+while true; do
+  status=$(curl -s "http://localhost:8091/v1/videos/${video_id}" | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
+
+curl -s "http://localhost:8091/v1/videos/${video_id}" | jq .
+curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o wan22_output.mp4
 ```
 
 ## Request Format
@@ -82,8 +128,9 @@ curl -X POST http://localhost:8091/v1/videos \
 | Parameter             | Type   | Default | Description                                      |
 | --------------------- | ------ | ------- | ------------------------------------------------ |
 | `prompt`              | str    | -       | Text description of the desired video            |
+| `seconds`             | str    | None    | Clip duration in seconds                         |
+| `size`                | str    | None    | Output size in `WIDTHxHEIGHT` format             |
 | `negative_prompt`     | str    | None    | Negative prompt                                  |
-| `n`                   | int    | 1       | Number of videos to generate                     |
 | `width`               | int    | None    | Video width in pixels                            |
 | `height`              | int    | None    | Video height in pixels                           |
 | `num_frames`          | int    | None    | Number of frames to generate                     |
@@ -95,22 +142,60 @@ curl -X POST http://localhost:8091/v1/videos \
 | `flow_shift`          | float  | None    | Scheduler flow shift (Wan2.2)                    |
 | `seed`                | int    | None    | Random seed (reproducible)                       |
 | `lora`                | object | None    | LoRA configuration                               |
-| `extra_body`          | object | None    | Model-specific extra parameters                  |
 
-## Response Format
+## Create Response Format
+
+`POST /v1/videos` returns a job record, not inline base64 video data.
 
 ```json
 {
-  "created": 1234567890,
-  "data": [
-    { "b64_json": "<base64-mp4>" }
-  ]
+  "id": "video_gen_123",
+  "object": "video",
+  "status": "queued",
+  "model": "Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+  "prompt": "A cinematic view of a futuristic city at sunset",
+  "created_at": 1234567890
 }
 ```
 
-## Extract Video
+## Retrieve, List, Download, and Delete
+
+### Retrieve a job
 
 ```bash
-# Extract base64 from response and decode to video
-cat response.json | jq -r '.data[0].b64_json' | base64 -d > wan22_output.mp4
+curl -s http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+### List jobs
+
+```bash
+curl -s http://localhost:8091/v1/videos | jq .
+```
+
+### Download the completed video
+
+```bash
+curl -L http://localhost:8091/v1/videos/${video_id}/content -o wan22_output.mp4
+```
+
+### Delete a job and its stored file
+
+```bash
+curl -X DELETE http://localhost:8091/v1/videos/${video_id} | jq .
+```
+
+## Poll Until Complete
+
+```bash
+while true; do
+  status=$(curl -s http://localhost:8091/v1/videos/${video_id} | jq -r '.status')
+  if [ "$status" = "completed" ]; then
+    break
+  fi
+  if [ "$status" = "failed" ]; then
+    echo "Video generation failed"
+    exit 1
+  fi
+  sleep 2
+done
 ```

examples/online_serving/text_to_video/run_curl_text_to_video.sh

@@ -1,20 +1,63 @@
 #!/bin/bash
-# Wan2.2 text-to-video curl example
+# Wan2.2 text-to-video curl example using the async video job API.
 
-OUTPUT_PATH="wan22_output.mp4"
+set -euo pipefail
 
-curl -X POST http://localhost:8098/v1/videos \
-  -H "Accept: application/json" \
-  -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-  -F "seconds=2" \
-  -F "size=832x480" \
-  -F "negative_prompt=色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
-  -F "fps=16" \
-  -F "num_inference_steps=40" \
-  -F "guidance_scale=4.0" \
-  -F "guidance_scale_2=4.0" \
-  -F "boundary_ratio=0.875" \
-  -F "flow_shift=5.0" \
-  -F "seed=42" | jq -r '.data[0].b64_json' | base64 -d > "${OUTPUT_PATH}"
+BASE_URL="${BASE_URL:-http://localhost:8098}"
+OUTPUT_PATH="${OUTPUT_PATH:-wan22_output.mp4}"
+POLL_INTERVAL="${POLL_INTERVAL:-2}"
 
+create_response=$(
+  curl -sS -X POST "${BASE_URL}/v1/videos" \
+    -H "Accept: application/json" \
+    -F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
+    -F "seconds=2" \
+    -F "size=832x480" \
+    -F "negative_prompt=色调艳丽，过曝，静态，细节模糊不清，字幕，风格，作品，画作，画面，静止，整体发灰，最差质量，低质量，JPEG压缩残留，丑陋的，残缺的，多余的手指，画得不好的手部，画得不好的脸部，畸形的，毁容的，形态畸形的肢体，手指融合，静止不动的画面，杂乱的背景，三条腿，背景人很多，倒着走" \
+    -F "fps=16" \
+    -F "num_inference_steps=40" \
+    -F "guidance_scale=4.0" \
+    -F "guidance_scale_2=4.0" \
+    -F "boundary_ratio=0.875" \
+    -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}"

pyproject.toml

@@ -167,6 +167,7 @@ markers = [
     "gpu: Tests that run on GPU (auto-added)",
     "cuda: Tests that run on CUDA (auto-added)",
     "rocm: Tests that run on AMD/ROCm (auto-added)",
+    "xpu: Tests that run on XPU (auto-added)",
     "npu: Tests that run on NPU/Ascend (auto-added)",
     # specified computation resources marks (auto-added)
     "H100: Tests that require H100 GPU",
@@ -176,9 +177,11 @@ markers = [
     "A3: Tests that require A3 NPU",
     "distributed_cuda: Tests that require multi cards on CUDA platform",
     "distributed_rocm: Tests that require multi cards on ROCm platform",
+    "distributed_xpu: Tests that require multi cards on XPU platform",
     "distributed_npu: Tests that require multi cards on NPU platform",
     "skipif_cuda: Skip if the num of CUDA cards is less than the required",
     "skipif_rocm: Skip if the num of ROCm cards is less than the required",
+    "skipif_xpu: Skip if the num of XPU cards is less than the required",
     "skipif_npu: Skip if the num of NPU cards is less than the required",
     # more detailed markers
     "slow: Slow tests (may skip in quick CI)",

requirements/common.txt

@@ -6,7 +6,7 @@ diffusers>=0.36.0
 accelerate==1.12.0
 gradio==5.50
 soundfile>=0.13.1
-cache-dit==1.2.0
+cache-dit==1.3.0
 tqdm>=4.66.0
 torchsde>=0.2.6
 openai-whisper>=20250625
@@ -17,3 +17,4 @@ einops>=0.8.1
 prettytable>=3.8.0
 aenum==3.1.16
 pyzmq>=25.0.0
+janus>=1.0.0

tests/benchmarks/test_serve_cli.py

@@ -5,10 +5,14 @@ import pytest
 
 from tests.conftest import OmniServerParams
 from tests.utils import hardware_test
+from vllm_omni.platforms import current_omni_platform
 
 models = ["Qwen/Qwen2.5-Omni-7B"]
 stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "qwen2_5_omni_ci.yaml")]
 
+if current_omni_platform.is_xpu():
+    stage_configs = [str(Path(__file__).parent.parent / "e2e" / "stage_configs" / "xpu" / "qwen2_5_omni_ci.yaml")]
+
 # Create parameter combinations for model and stage config
 test_params = [
     OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
@@ -17,7 +21,7 @@ test_params = [
 
 @pytest.mark.core_model
 @pytest.mark.benchmark
-@hardware_test(res={"cuda": "L4"}, num_cards=3)
+@hardware_test(res={"cuda": "L4", "xpu": "B60"}, num_cards=3)
 @pytest.mark.parametrize("omni_server", test_params, indirect=True)
 def test_bench_serve_chat(omni_server):
     command = [