Initial plan

Signed-off-by: Chenguang ZHENG <645327136@qq.com>

.buildkite/pipeline.yml

@@ -18,18 +18,55 @@ steps:
   #     queue: "ascend"
 
   - label: "Simple Unit Test"
-    depends_on: ~
+    depends_on: image-build
     commands:
-      - ".buildkite/scripts/simple_test.sh"
+    - |
+      pytest -v -s \
+        tests/entrypoints/ \
+        tests/diffusion/cache/ \
+        tests/diffusion/lora/ \
+        tests/model_executor/models/qwen2_5_omni/test_audio_length.py \
+        tests/worker/ \
+        tests/distributed/omni_connectors/test_kv_flow.py \
+        --cov=vllm_omni \
+        --cov-branch \
+        --cov-report=term-missing \
+        --cov-report=html \
+        --cov-report=xml
     agents:
-      queue: "cpu_queue_premerge"
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
 
   - label: "Diffusion Model Test"
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
-      - pytest -s -v tests/e2e/offline_inference/test_ovis_image.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Images API LoRA E2E"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
+      - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -47,6 +84,7 @@ steps:
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
     agents:
       queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
     plugins:
@@ -93,11 +131,28 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Diffusion Parallelism Test"
-    timeout_in_minutes: 25
+  - label: "Diffusion Sequence Parallelism Test"
+    timeout_in_minutes: 20
     depends_on: image-build
     commands:
       - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          shm-size: "8gb"
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Tensor Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: image-build
+    commands:
       - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -116,7 +171,7 @@ steps:
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
-      - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+      - pytest -s -v tests/diffusion/test_diffusion_worker.py
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
     plugins:
@@ -130,31 +185,13 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "Omni Model Test"
+
+  - label: "Benchmark Test"
     timeout_in_minutes: 15
     depends_on: image-build
     commands:
-      - export VLLM_LOGGING_LEVEL=DEBUG
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
-    agents:
-      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
-    plugins:
-      - docker#v5.2.0:
-          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
-          always-pull: true
-          propagate-environment: true
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
-  - label: "Omni Model Test with H100"
-    timeout_in_minutes: 30
-    depends_on: image-build
-    commands:
-      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_async_omni.py
+      - pytest -s -v tests/benchmarks/test_serve_cli.py
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -184,12 +221,69 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
+  - label: "Omni Model Test"
+    timeout_in_minutes: 15
+    depends_on: image-build
+    commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+    agents:
+      queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+    plugins:
+    - docker#v5.2.0:
+        image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+        always-pull: true
+        propagate-environment: true
+        environment:
+        - "HF_HOME=/fsx/hf_cache"
+        volumes:
+        - "/fsx/hf_cache:/fsx/hf_cache"
+
+  # - label: "Omni Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   commands:
+  #     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #     - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+  #     - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+  #     - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+  #     - pytest -s -v tests/e2e/online_serving/test_async_omni.py
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 2
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate
+
   - label: "Diffusion Image Edit Test with H100 (1 GPU)"
     timeout_in_minutes: 20
     depends_on: image-build
     commands:
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
+      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
     agents:
       queue: "mithril-h100-pool"
     plugins:
@@ -218,3 +312,38 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
+
+  # - label: "Bagel Text2Img Model Test with H100"
+  #   timeout_in_minutes: 30
+  #   depends_on: image-build
+  #   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: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 1
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate

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

@@ -116,6 +116,9 @@ if [[ $commands == *"--shard-id="* ]]; then
         --shm-size=16gb \
         --group-add "$render_gid" \
         --rm \
+        -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+        -e MIOPEN_DEBUG_CONV_GEMM=0 \
+        -e VLLM_ROCM_USE_AITER=1 \
         -e HIP_VISIBLE_DEVICES="${GPU}" \
         -e HF_TOKEN \
         -e AWS_ACCESS_KEY_ID \
@@ -148,6 +151,9 @@ else
           --shm-size=16gb \
           --group-add "$render_gid" \
           --rm \
+          -e MIOPEN_DEBUG_CONV_DIRECT=0 \
+          -e MIOPEN_DEBUG_CONV_GEMM=0 \
+          -e VLLM_ROCM_USE_AITER=1 \
           -e HF_TOKEN \
           -e AWS_ACCESS_KEY_ID \
           -e AWS_SECRET_ACCESS_KEY \

.buildkite/scripts/simple_test.sh

@@ -1,55 +0,0 @@
-#!/usr/bin/env bash
-
-set -euo pipefail
-
-SCRIPT_DIR="$(cd -- "$(dirname -- "${BASH_SOURCE[0]}")" && pwd)"
-REPO_ROOT="$(cd -- "${SCRIPT_DIR}/../.." && pwd)"
-UV_BIN="uv"
-PYTHON_BOOTSTRAP="${PYTHON:-python3}"
-VENV_DIR="${VENV_DIR:-${REPO_ROOT}/.venv-simple-test}"
-TARGET_PY_VERSION="3.12"
-
-if ! command -v "${PYTHON_BOOTSTRAP}" >/dev/null 2>&1; then
-  PYTHON_BOOTSTRAP="python"
-fi
-
-cd "${REPO_ROOT}"
-
-# Ensure pip-installed scripts (including uv) are on PATH
-PYTHON_SCRIPTS="$(${PYTHON_BOOTSTRAP} -c 'import sysconfig; print(sysconfig.get_path("scripts"))')"
-USER_BASE="$(${PYTHON_BOOTSTRAP} -m site --user-base 2>/dev/null || true)"
-PATH="${PYTHON_SCRIPTS}:${PATH}"
-if [[ -n "${USER_BASE}" ]]; then
-  PATH="${USER_BASE}/bin:${PATH}"
-fi
-
-if ! command -v "${UV_BIN}" >/dev/null 2>&1; then
-  "${PYTHON_BOOTSTRAP}" -m pip install --upgrade pip
-  "${PYTHON_BOOTSTRAP}" -m pip install uv
-  hash -r
-fi
-
-TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  "${UV_BIN}" python install "${TARGET_PY_VERSION}"
-  TARGET_PY_PATH="$(${UV_BIN} python find "${TARGET_PY_VERSION}" 2>/dev/null | head -n1 || true)"
-fi
-
-if [[ -z "${TARGET_PY_PATH}" ]]; then
-  echo "Failed to locate python ${TARGET_PY_VERSION}" >&2
-  exit 1
-fi
-
-if [[ ! -d "${VENV_DIR}" ]]; then
-  "${UV_BIN}" venv --python "${TARGET_PY_PATH}" "${VENV_DIR}"
-fi
-
-VENV_PYTHON="${VENV_DIR}/bin/python"
-[[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; }
-
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.12.0
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
-"${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/
-"${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/
-"${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py
-"${VENV_PYTHON}" -m pytest -v -s tests/worker/

.buildkite/test-amd.yaml

@@ -8,14 +8,32 @@ steps:
   grade: Blocking
   commands:
     - export GPU_ARCHS=gfx942
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
 
+- label: "Diffusion Images API LoRA E2E"
+  timeout_in_minutes: 20
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
+
+- label: "Diffusion Model CPU offloading Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_1
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+
 - label: "Diffusion Cache Backend Test"
   timeout_in_minutes: 15
   agent_pool: mi325_1
@@ -26,25 +44,30 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_cache_dit.py tests/e2e/offline_inference/test_teacache.py
 
-- label: "Diffusion Parallelism Test"
-  timeout_in_minutes: 15
+- label: "Diffusion Sequence Parallelism Test"
+  timeout_in_minutes: 20
   agent_pool: mi325_2
   depends_on: amd-build
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
     - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
 
+- label: "Diffusion Tensor Parallelism Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  commands:
+    - export GPU_ARCHS=gfx942
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - pytest -s -v tests/e2e/offline_inference/test_zimage_tensor_parallel.py
+
 - label: "Diffusion GPU Worker Test"
   timeout_in_minutes: 20
   agent_pool: mi325_2
@@ -52,9 +75,7 @@ steps:
   mirror_hardwares: [amdproduction]
   grade: Blocking
   commands:
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/diffusion/test_gpu_diffusion_worker.py
+    - pytest -s -v tests/diffusion/test_diffusion_worker.py
 
 - label: "Omni Model Test Qwen2-5-Omni"
   timeout_in_minutes: 15
@@ -66,12 +87,6 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
     - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
 
 - label: "Omni Model Test Qwen3-Omni"
@@ -83,9 +98,10 @@ steps:
   commands:
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py
+    - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+    - pytest -s -v tests/e2e/online_serving/test_async_omni.py
 
 - label: "Diffusion Image Edit Test"
   timeout_in_minutes: 15
@@ -97,10 +113,4 @@ steps:
     - export GPU_ARCHS=gfx942
     - export VLLM_LOGGING_LEVEL=DEBUG
     - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-    - export MIOPEN_DEBUG_CONV_DIRECT=0
-    - export MIOPEN_DEBUG_CONV_GEMM=0
-    - export VLLM_ROCM_USE_AITER=1
-    - export VLLM_ROCM_USE_AITER_MHA=1
-    - export VLLM_ROCM_USE_AITER_LINEAR=0
-    - export VLLM_ROCM_USE_AITER_RMSNORM=0
-    - pytest -s -v tests/e2e/online_serving/test_i2i_multi_image_input.py
+    - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py

.buildkite/test-nightly.yaml

@@ -0,0 +1,45 @@
+steps:
+  - label: ":docker: Build image"
+    key: image-build
+    commands:
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --file docker/Dockerfile.ci -t vllm-omni-ci ."
+      - "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+      - "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
+    agents:
+      queue: "cpu_queue_premerge"
+
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 180
+    depends_on: image-build
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

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

@@ -14,7 +14,7 @@
         soft_fail: false
         commands:
           - "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
-          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --progress plain ."
+          - "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target final --progress plain ."
           - "docker push {{ docker_image_amd }}"
         key: "amd-build"
         env:

.github/workflows/build_wheel.yml

@@ -7,6 +7,7 @@ on:
     paths-ignore:
       - 'docs/**'
       - 'examples/**'
+      - 'apps/**'
       - '.github/**'
       - '.buildkite/**'
       - '**.md'
@@ -20,6 +21,7 @@ on:
     paths-ignore:
       - 'docs/**'
       - 'examples/**'
+      - 'apps/**'
       - '.github/**'
       - '.buildkite/**'
       - '*.yml'

.github/workflows/comfyui-publish.yml

@@ -0,0 +1,27 @@
+name: ComfyUI Integration -- Publish to Comfy registry
+on:
+  workflow_dispatch:
+  push:
+    tags:
+      - "*"
+    paths:
+      - "apps/ComfyUI-vLLM-Omni/**"
+
+permissions:
+  issues: write
+
+defaults:
+  run:
+    working-directory: apps/ComfyUI-vLLM-Omni
+
+jobs:
+  publish-node:
+    name: Publish Custom Node to registry
+    runs-on: ubuntu-latest
+    steps:
+      - name: ♻️ Check out code
+        uses: actions/checkout@v4
+      - name: 📦 Publish Custom Node
+        uses: Comfy-Org/publish-node-action@main
+        with:
+          personal_access_token: ${{ secrets.REGISTRY_ACCESS_TOKEN }}

.github/workflows/comfyui-validate.yml

@@ -0,0 +1,40 @@
+name: ComfyUI Integration -- Validate backwards compatibility
+
+on:
+  pull_request:
+    branches:
+      - master
+      - main
+    paths:
+      - "apps/ComfyUI-vLLM-Omni/**"
+
+defaults:
+  run:
+    working-directory: apps/ComfyUI-vLLM-Omni
+
+jobs:
+  check-base-path:
+    runs-on: ubuntu-latest
+    outputs:
+      exists: ${{ steps.check.outputs.exists }}
+    defaults:
+      run:
+        working-directory: .
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          ref: ${{ github.event.pull_request.base.sha }}
+          fetch-depth: 1
+      - id: check
+        run: |
+          if [ -d "apps/ComfyUI-vLLM-Omni" ]; then
+            echo "exists=true" >> "$GITHUB_OUTPUT"
+          else
+            echo "exists=false" >> "$GITHUB_OUTPUT"
+          fi
+  validate:
+    needs: check-base-path
+    if: needs.check-base-path.outputs.exists == 'true'
+    runs-on: ubuntu-latest
+    steps:
+      - uses: comfy-org/node-diff@main

.gitignore

@@ -153,6 +153,14 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
+# Claude
+CLAUDE.md
+.claude/
+
+# Codex
+AGENTS.md
+.codex/
+
 # cursor
 .cursor/
 
@@ -214,6 +222,7 @@ data/
 datasets/
 *.csv
 *.json
+!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet
 

.pre-commit-config.yaml

@@ -29,7 +29,10 @@ repos:
         # only for staged files
 
   - repo: https://github.com/rhysd/actionlint
-    rev: v1.7.9
+    # v1.7.8+ sets `go 1.24.0` in go.mod, which older Go toolchains (and most
+    # current CI images) cannot parse. Pin to v1.7.7 until actionlint fixes the
+    # go.mod directive.
+    rev: v1.7.7
     hooks:
       - id: actionlint
         files: ^\.github/workflows/.*\.ya?ml$

README.md

@@ -9,14 +9,16 @@ 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://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/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
 </p>
 
+
 ---
 
 *Latest News* 🔥
 
-- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true).
+- [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.
+- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
 - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving.
 
 ---
@@ -67,6 +69,19 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
 We welcome and value any contributions and collaborations.
 Please check out [Contributing to vLLM-Omni](https://vllm-omni.readthedocs.io/en/latest/contributing/) for how to get involved.
 
+## Citation
+
+If you use vLLM-Omni for your research, please cite our [paper](https://arxiv.org/abs/2602.02204):
+
+```bibtex
+@article{yin2026vllmomni,
+  title={vLLM-Omni: Fully Disaggregated Serving for Any-to-Any Multimodal Models},
+  author={Peiqi Yin, Jiangyun Zhu, Han Gao, Chenguang Zheng, Yongxiang Huang, Taichang Zhou, Ruirui Yang, Weizhi Liu, Weiqing Chen, Canlin Guo, Didan Deng, Zifeng Mo, Cong Wang, James Cheng, Roger Wang, Hongsheng Liu},
+  journal={arXiv preprint arXiv:2602.02204},
+  year={2026}
+}
+```
+
 ## Join the Community
 Feel free to ask questions, provide feedbacks and discuss with fellow users of vLLM-Omni in `#sig-omni` slack channel at [slack.vllm.ai](https://slack.vllm.ai) or vLLM user forum at [discuss.vllm.ai](https://discuss.vllm.ai).
 

apps/ComfyUI-vLLM-Omni/.gitignore

@@ -0,0 +1,115 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# OSX useful to ignore
+*.DS_Store
+.AppleDouble
+.LSOverride
+
+# Thumbnails
+._*
+
+# Files that might appear in the root of a volume
+.DocumentRevisions-V100
+.fseventsd
+.Spotlight-V100
+.TemporaryItems
+.Trashes
+.VolumeIcon.icns
+.com.apple.timemachine.donotpresent
+
+# Directories potentially created on remote AFP share
+.AppleDB
+.AppleDesktop
+Network Trash Folder
+Temporary Items
+.apdisk
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+env/
+venv/
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+*.egg-info/
+.installed.cfg
+*.egg
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+
+# Sphinx documentation
+docs/_build/
+
+# IntelliJ Idea
+.idea
+*.iml
+*.ipr
+*.iws
+
+# PyBuilder
+target/
+
+# Cookiecutter
+output/
+python_boilerplate/
+cookiecutter-pypackage-env/
+
+# vscode settings
+.history/
+*.code-workspace
+
+# Frontend extension
+node_modules/
+.env
+.env.local
+.env.development.local
+.env.test.local
+.env.production.local
+npm-debug.log*
+yarn-debug.log*
+yarn-error.log*
+node.zip
+.vscode/
+.claude/
+.codemate/

apps/ComfyUI-vLLM-Omni/LICENSE

@@ -0,0 +1,15 @@
+Apache Software License 2.0
+
+Copyright (c) 2026, Zeyu Huang
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.

apps/ComfyUI-vLLM-Omni/README.md

@@ -0,0 +1,170 @@
+# vLLM-Omni
+
+vLLM-Omni offers a ComfyUI integration on top of its online serving API.
+It can send model inference requests to either a locally running vLLM-Omni service or a remote one.
+
+## Requirement
+
+- Python 3.12 or above
+- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements)
+- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet.
+- No need to install additional packages apart from those already required by ComfyUI.
+
+> [!TIP]
+> If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them.
+
+
+## Installation
+
+Copy this folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`.
+
+If you are running ComfyUI during copying, you should restart ComfyUI to load this extension.
+
+> [!TIP]
+> You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info.
+
+On the device and virtual environment you run ComfyUI, launch ComfyUI with
+```bash
+cd ComfyUI
+
+# The regular way
+python main.py
+
+# If you are mainly using this node, launch it faster with
+python main.py --cpu
+```
+
+On the device and virtual environment you run vLLM-Omni, start a model service with
+```bash
+vllm serve The_Model_ID_to_Serve --omni --port 8000
+```
+
+Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**.
+If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`.
+
+## Quickstart
+
+This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
+
+- **Generate Image** for text-to-image and image-to-image tasks
+- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **TTS** and **TTS Voice Clone** for TTS tasks
+
+This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
+
+> [!INFO]
+> The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
+
+To build a simple workflow yourself,
+
+- Drag a generation node onto the canvas.
+- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
+    - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
+    - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
+    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
+
+**The following features are tested**:
+
+- Single-node workflows for
+    - Multimodal Comprehension (e.g., Qwen Omni, BAGEL)
+    - Text-to-Image Generation (e.g., Qwen-Image)
+    - Image-to-Image Generation (e.g., Qwen-Image-Edit)
+    - TTS (e.g., Qwen TTS, including VoiceDesign, VoiceClone, CustomVoice)
+
+**The following features are not currently tested**. They may work or break. You are welcomed to test it out and offer comments.
+
+- Multi-node workflow that connects multiple model services together.
+
+## Screenshots and Examples
+
+### Multimodal comprehension (e.g., Qwen Omni series, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Annotated Example**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-comprehension.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> Although this node enables all-modality input, you should check whether the specific model you host and request for supports the modalities you connect to the node.
+
+You can configure per-stage sampling parameters for multi-stage models.
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-multi-stage.jpg" width=55%>
+  </picture>
+</p>
+
+### Text-to-image and image-to-image generation (e.g., Z-Image-Turbo, Qwen-Image-Edit, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Image Generation**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-image-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
+
+### TTS (e.g., Qwen TTS series)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
+
+<p align="center">
+  <picture>
+    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg">
+    <img alt="vLLM-Omni Main Architecture" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-tts.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> There is a dedicated node for VoiceClone tasks with reference audio input. Other simple text-to-speech tasks should use the regular TTS node.
+
+## Develop
+
+Follow the [development convention and rules of vLLM-Omni](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/).
+
+## Limitation and Non-Goals
+
+- Single server mode only. No automatic load balancing or failover.
+- Features set is bounded to vLLM-Omni's online service capability, including
+    - The types of models supported in online mode,
+    - The types of sampling parameters supported in the online mode,
+    - The ways to send files (primarily through full-length base64 in JSON payload),
+    - Figuring out errors in the payload (such as unsupported fields by a specific model) if the endpoint does not explicitly return an error,
+    - (The lack of) Authentication
+    - (The lack of) Progress indicator
+
+## Support
+
+If you are new to ComfyUI, please check out [its documentation](https://docs.comfy.org/) for usage instructions.
+
+If you are new to vLLM-Omni, please also check out [its documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/) for usage instructions.
+
+Whenever you find an issue or problem, please
+
+- First find out if this is an upstream limitation of vLLM-Omni's online serving mode, by [checking their documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/examples/).
+- [Open an issue](https://github.com/vllm-project/vllm-omni/issues) that clearly describes this ComfyUI or online service problem.
+
+## Acknowledgements
+
+Features
+
+- https://github.com/dougbtv/comfyui-vllm-omni/ The official reference implementation for ComfyUI integration with vLLM-Omni's DALL-E compatible image generation API.
+- https://github.com/Comfy-Org/ComfyUI/tree/master/comfy_extras ComfyUI's built-in node implementations.
+
+UI/UX design references
+
+- https://github.com/sgl-project/sglang/pull/15271 SGLang Diffusion's official ComfyUI integration for image and video generation.
+- https://github.com/SXQBW/ComfyUI-Qwen-Omni A third party ComfyUI integration for Qwen Omni series.
+- https://github.com/flybirdxx/ComfyUI-Qwen-TTS https://github.com/DarioFT/ComfyUI-Qwen3-TTS Third  party ComfyUI integrations for Qwen TTS series.

apps/ComfyUI-vLLM-Omni/__init__.py

@@ -0,0 +1,52 @@
+"""Top-level package for comfyui_vllm_omni."""  # noqa: N999  # This is not a python library intended to be imported
+
+__all__ = [
+    "NODE_CLASS_MAPPINGS",
+    "NODE_DISPLAY_NAME_MAPPINGS",
+    "WEB_DIRECTORY",
+]
+
+__author__ = """Zeyu Huang"""
+__email__ = "11222265+fhfuih@users.noreply.github.com"
+__version__ = "0.0.1"
+
+from .comfyui_vllm_omni.nodes import (
+    VLLMOmniARSampling,
+    VLLMOmniComprehension,
+    VLLMOmniDiffusionSampling,
+    VLLMOmniGenerateImage,
+    VLLMOmniQwenTTSParams,
+    VLLMOmniSamplingParamsList,
+    VLLMOmniTTS,
+    VLLMOmniVoiceClone,
+)
+
+# A dictionary that contains all nodes you want to export with their names
+NODE_CLASS_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
+    "VLLMOmniComprehension": VLLMOmniComprehension,
+    "VLLMOmniTTS": VLLMOmniTTS,
+    "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
+    # === Params ===
+    "VLLMOmniARSampling": VLLMOmniARSampling,
+    "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
+    "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
+    "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": "Generate Image",
+    "VLLMOmniComprehension": "Multimodality Comprehension",
+    "VLLMOmniTTS": "TTS (Text to Speech)",
+    "VLLMOmniVoiceClone": "TTS Voice Cloning",
+    # === Params ===
+    "VLLMOmniARSampling": "AR Sampling Params",
+    "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
+    "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
+    "VLLMOmniQwenTTSParams": "Qwen TTS Params",
+}
+
+WEB_DIRECTORY = "./web"

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/__init__.py

@@ -0,0 +1 @@
+# noqa: N999   # This is not a python library intended to be imported

apps/ComfyUI-vLLM-Omni/comfyui_vllm_omni/nodes.py

@@ -0,0 +1,569 @@
+from typing import Literal, cast
+
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .utils.api_client import VLLMOmniClient
+from .utils.logger import get_logger
+from .utils.models import lookup_model_spec
+from .utils.types import AudioFormat
+from .utils.validators import (
+    add_sampling_parameters_to_stage,
+    validate_model_and_sampling_params_types,
+)
+
+logger = get_logger(__name__)
+
+
+class _VLLMOmniGenerateBase:
+    """Base class for vLLM-Omni generation nodes with shared functionality."""
+
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model) -> str | Literal[True]:
+        """
+        Can only validate this model's own input. Cannot check inputs from other nodes.
+        See: https://docs.comfy.org/custom-nodes/backend/server_overview#validate_inputs
+        """
+        if not url:
+            return "URL must not be empty"
+        if not model:
+            return "Model must not be empty"
+        return True
+
+
+class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Tongyi-MAI/Z-Image-Turbo"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 512, "min": 64, "max": 2048}),
+                "height": ("INT", {"default": 512, "min": 64, "max": 2048}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "mask": ("MASK",),
+                # "video": ("VIDEO",),
+                # "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("IMAGE",)
+    RETURN_NAMES = ("image",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        mask: torch.Tensor | None = None,
+        audio: AudioInput | None = None,  # Hidden & unused
+        video: VideoInput | None = None,  # Hidden & unused
+        sampling_params: dict | list[dict] | None = None,
+        **kwargs,
+    ):
+        logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+        if image is None and mask is not None:
+            raise ValueError("Mask input provided without an image input.")
+
+        client = VLLMOmniClient(url)
+
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        # Prefer DALL-E compatible API for simple (one-stage) diffusion models
+        if (spec is None or spec["stages"] == ["diffusion"]) and not is_bagel:
+            sampling_params = cast(dict | None, sampling_params)
+            if audio is None and image is None and video is None:
+                # No multimodal input --- use DALL-E image generation
+                logger.info("Using DALL-E image generation endpoint")
+                output = await client.generate_image(
+                    model=model,
+                    prompt=prompt,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    sampling_params=sampling_params,
+                )
+                return (output,)
+            elif image is not None and audio is None and video is None:
+                # Image and text input --- use DALL-E image edit
+                logger.info("Using DALL-E image edit endpoint")
+                output = await client.edit_image(
+                    model=model,
+                    prompt=prompt,
+                    image=image,
+                    width=width,
+                    height=height,
+                    negative_prompt=negative_prompt,
+                    mask=mask,
+                    sampling_params=sampling_params,
+                )
+                return (output,)
+
+        logger.info("Using chat completion endpoint")
+        sampling_params = add_sampling_parameters_to_stage(
+            model, sampling_params, "diffusion", width=width, height=height
+        )
+        logger.debug("Edited sampling params: %s", sampling_params)
+
+        output = await client.generate_image_chat_completion(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+        )
+
+        return (output,)
+
+
+class VLLMOmniComprehension(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen2.5-Omni-7B"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "output_text": ("BOOLEAN", {"default": True}),
+                "output_audio": ("BOOLEAN", {"default": True}),
+                "use_audio_in_video": ("BOOLEAN", {"default": True}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "video": ("VIDEO",),
+                "audio": ("AUDIO",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("STRING", "AUDIO")
+    RETURN_NAMES = ("text_response", "audio_response")
+    FUNCTION = "generate"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, url, model, output_text, output_audio) -> str | Literal[True]:  # type: ignore[reportIncompatibleMethodOverride]
+        super_validation = super().VALIDATE_INPUTS(url, model)
+        if isinstance(super_validation, str):
+            return super_validation
+        if not output_text and not output_audio:
+            return "At least one of output_text or output_audio must be True."
+        return True
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        output_text: bool = True,
+        output_audio: bool = True,
+        use_audio_in_video: bool = True,
+        **kwargs,
+    ) -> tuple[str, AudioInput]:
+        logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        client = VLLMOmniClient(url)
+        spec, pattern = lookup_model_spec(model)
+        is_bagel = pattern is not None and "bagel" in pattern.lower()
+
+        if is_bagel:
+            # A lot of special handlings here...
+            if output_audio:
+                raise ValueError("BAGEL models do not support audio output.")
+            if audio is not None or video is not None:
+                raise ValueError("BAGEL models do not support audio or video input.")
+            (
+                text_response,
+                _,
+            ) = await client.generate_comprehension_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=None,
+                video=None,
+                sampling_params=sampling_params,
+                modalities=["text"],
+            )
+        else:
+            modalities = []
+            if output_text:
+                modalities.append("text")
+            if output_audio:
+                modalities.append("audio")
+
+            if use_audio_in_video and video is not None:
+                use_audio_in_video = True
+            else:
+                use_audio_in_video = False
+
+            (
+                text_response,
+                audio,
+            ) = await client.generate_comprehension_chat_completion(
+                model=model,
+                prompt=prompt,
+                image=image,
+                audio=audio,
+                video=video,
+                sampling_params=sampling_params,
+                modalities=modalities,
+                # == extra kwargs ==
+                mm_processor_kwargs={"use_audio_in_video": use_audio_in_video},
+            )
+
+        if text_response is None:
+            text_response = ""
+        if audio is None:
+            channels = 1
+            duration = 1
+            sample_rate = 44100
+            num_samples = int(round(duration * sample_rate))
+            waveform = torch.zeros((1, channels, num_samples), dtype=torch.float32)
+            audio = {"waveform": waveform, "sample_rate": sample_rate}
+
+        return (text_response, audio)
+
+
+class VLLMOmniTTS(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": (
+                    "STRING",
+                    {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice"},
+                ),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        model_specific_params: dict | None,
+        **kwargs,
+    ) -> tuple[AudioInput]:
+        logger.info("Got extra kwargs in TTS: %s", kwargs)
+
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
+        if not is_qwen_tts and extra_params_type == "qwen-tts":
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {**kwargs, **(model_specific_params or {})}
+        combined_params.pop("type", None)  # Internal fields in model_specific_params
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniVoiceClone(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Qwen/Qwen3-TTS-12Hz-1.7B-Base"}),
+                "input": ("STRING", {"multiline": True}),
+                "voice": ("STRING", {"default": "Vivian"}),
+                "response_format": (["mp3", "opus", "aac", "flac", "wav", "pcm"],),
+                "speed": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.25, "max": 4.0, "step": 0.01},
+                ),
+                "ref_audio": ("AUDIO",),
+                "ref_text": ("STRING", {"multiline": True}),
+                "x_vector_only_mode": ("BOOLEAN", {"default": False}),
+            },
+            "optional": {
+                "model_specific_params": ("TTS_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("AUDIO",)
+    RETURN_NAMES = ("audio",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        ref_audio: AudioInput,
+        ref_text: str,
+        x_vector_only_mode: bool,
+        model_specific_params: dict | None,
+        **kwargs,
+    ):
+        is_qwen_tts = "qwen3-tts" in model.lower()
+        extra_params_type = None if model_specific_params is None else model_specific_params["type"]
+        if not is_qwen_tts and extra_params_type == "qwen-tts":
+            raise ValueError(
+                "You have provided Qwen-specific TTS params."
+                "However, the model appears to not be a Qwen TTS model (no 'Qwen3-TTS' in model name)."
+            )
+
+        combined_params = {
+            "ref_audio": ref_audio,
+            "ref_text": ref_text,
+            "x_vector_only_mode": x_vector_only_mode,
+            **kwargs,
+            **(model_specific_params or {}),
+        }
+        combined_params.pop("type", None)  # Internal fields in model_specific_params
+
+        client = VLLMOmniClient(url)
+
+        audio = await client.generate_speech(
+            model=model,
+            input=input,
+            voice=voice,
+            response_format=response_format,
+            speed=speed,
+            **combined_params,
+        )
+        return (audio,)
+
+
+class VLLMOmniARSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "max_tokens": ("INT", {"default": 100, "min": 1, "max": 10000}),
+                "temperature": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.01},
+                ),
+                "top_p": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01},
+                ),
+                "repetition_penalty": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.01},
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("AR sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = {
+            "type": "autoregression",  # for internal use, removed before sending the request
+            **kwargs,
+        }
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniDiffusionSampling:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "n": (
+                    "INT",
+                    {
+                        "default": 1,
+                        "min": 0,
+                        "max": 10,
+                        "step": 1,
+                        "tooltip": "Number of images to generate",
+                    },
+                ),
+                "num_inference_steps": (
+                    "INT",
+                    {
+                        "default": 50,
+                        "min": 1,
+                        "max": 1000,
+                        "tooltip": "Number of denoising steps (higher = better quality, slower).",
+                    },
+                ),
+                "guidance_scale": (
+                    "FLOAT",
+                    {
+                        "default": 7.5,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.1,
+                        "tooltip": "Classifier-free guidance scale (higher = more prompt adherence).",
+                    },
+                ),
+                "true_cfg_scale": (
+                    "FLOAT",
+                    {
+                        "default": 1.0,
+                        "min": 0.0,
+                        "max": 20.0,
+                        "step": 0.5,
+                        "tooltip": "True CFG scale for advanced control (model-specific).",
+                    },
+                ),
+                "vae_use_slicing": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE slicing for reduced memory usage (slight quality trade-off)",
+                    },
+                ),
+                # === Put seed at last. ===
+                # Whenever a field named "seed" is present, ComfyUI adds another field called "control after generate"
+                "seed": (
+                    "INT",
+                    {
+                        "default": -1,
+                        "min": -1,
+                        "step": 1,
+                        "tooltip": "-1 means to not provide a seed.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("diffusion sampling params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def get_params(self, seed, **kwargs):
+        params = {
+            "type": "diffusion",  # for internal use, removed before sending the request
+            **kwargs,
+        }
+        if seed >= 0:
+            params["seed"] = seed
+        return (params,)
+
+
+class VLLMOmniSamplingParamsList:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "param1": ("SAMPLING_PARAMS",),
+            },
+            "optional": {
+                "param2": ("SAMPLING_PARAMS",),
+                "param3": ("SAMPLING_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("SAMPLING_PARAMS",)
+    RETURN_NAMES = ("param list",)
+    FUNCTION = "aggregate"
+    CATEGORY = "vLLM-Omni/Sampling Params"
+
+    def aggregate(self, param1: dict, param2: dict | None = None, param3: dict | None = None):
+        for i, p in enumerate((param1, param2, param3)):
+            if isinstance(p, list):
+                raise ValueError(
+                    f"Input {i} is a Multi-Stage Sampling Params List. "
+                    f"Expected a single sampling parameters node (either AR or Diffusion)."
+                )
+
+        params = [param1]
+        if param2 is not None:
+            params.append(param2)
+        if param3 is not None:
+            params.append(param3)
+        return (params,)
+
+
+class VLLMOmniQwenTTSParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "task_type": (
+                    ["CustomVoice", "VoiceDesign", "Base"],
+                    {"default": "CustomVoice"},
+                ),
+                "language": (
+                    ["Auto", "Chinese", "English", "Japanese", "Korean"],
+                    {"default": "Auto"},
+                ),
+                "instructions": ("STRING", {"multiline": True}),
+                "max_new_tokens": ("INT", {"default": 2048, "min": 1}),
+            }
+        }
+
+    RETURN_TYPES = ("TTS_PARAMS",)
+    RETURN_NAMES = ("Qwen TTS params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/TTS Params"
+
+    def get_params(self, **kwargs):
+        return ({"type": "qwen-tts", **kwargs},)

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

@@ -0,0 +1,473 @@
+"""
+An high-level API client adapter that forwards ComfyUI inputs to vLLM-Omni's REST API,
+and transforms the API responses back to ComfyUI formats.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+from typing import Any
+
+import aiohttp
+import av.error
+import torch
+from comfy_api.input import AudioInput, VideoInput
+
+from .format import (
+    audio_to_base64,
+    base64_to_audio,
+    base64_to_image_tensor,
+    bytes_to_audio,
+    image_tensor_to_base64,
+    image_tensor_to_png_bytes,
+    video_to_base64,
+)
+from .logger import get_logger, pretty_printer
+from .models import lookup_model_spec
+from .types import AudioFormat
+
+logger = get_logger(__name__)
+
+
+class VLLMOmniClient:
+    def __init__(self, base_url: str, timeout: float = 300.0):
+        self.base_url = base_url
+        self.timeout = aiohttp.ClientTimeout(total=timeout)
+
+    async def generate_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        sampling_params: dict | None = None,
+    ) -> torch.Tensor:
+        """Run text-to-image generation via DALLE API"""
+        await self._check_model_exist(model)
+
+        size = f"{width}x{height}"
+        payload = {
+            "model": model,
+            "prompt": prompt,
+            "size": size,
+            "response_format": "b64_json",
+        }
+        if negative_prompt:
+            payload["negative_prompt"] = negative_prompt
+        if sampling_params is not None:
+            # Only select specific sampling params
+            for k in (
+                "n",
+                "num_inference_steps",
+                "guidance_scale",
+                "true_cfg_scale",
+                "vae_use_slicing",
+            ):
+                if k in sampling_params and sampling_params[k] is not None:
+                    payload[k] = sampling_params[k]
+        logger.debug("img gen payload: %s", payload)
+
+        url = self.base_url + "/images/generations"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) 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}")
+                    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")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+                        logger.debug("Image #%d has shape %s", idx, tensor.shape)
+
+                    batch_tensor = torch.stack(image_tensors, dim=0)
+                    logger.debug("batch_tensor output has shape: %s", batch_tensor.shape)
+                    return batch_tensor
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def edit_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        mask: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+    ) -> torch.Tensor:
+        """Run image editing via DALLE API"""
+        await self._check_model_exist(model)
+
+        size = f"{width}x{height}"
+        image_filename = "image.png"  # Required for multipart form
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field(
+            "image",
+            image_tensor_to_png_bytes(image, image_filename),
+            filename=image_filename,
+            content_type="image/png",
+        )
+        form.add_field("prompt", prompt)
+        form.add_field("size", size)
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            # Only select specific sampling params
+            for k in ("n", "num_inference_steps", "guidance_scale", "true_cfg_scale"):
+                if k in sampling_params and sampling_params[k] is not None:
+                    form.add_field(k, str(sampling_params[k]))
+        if mask is not None:
+            mask_filename = "mask.png"
+            form.add_field(
+                "mask",
+                image_tensor_to_png_bytes(mask, mask_filename),
+                filename=mask_filename,
+                content_type="image/png",
+            )
+
+        url = self.base_url + "/images/edits"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            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}")
+
+                    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")
+
+                    image_tensors = []
+                    for idx, img in enumerate(data["data"]):
+                        if "b64_json" not in img:
+                            raise RuntimeError(f"API returned image #{idx} without 'b64_json' field")
+                        base64_str = img["b64_json"]
+                        tensor = base64_to_image_tensor(base64_str)
+                        image_tensors.append(tensor)
+
+                    return torch.stack(image_tensors, dim=0)
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def generate_image_chat_completion(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+    ) -> torch.Tensor:
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=negative_prompt,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=["image"],
+        )
+        choices = await self._generate_base_chat_completion(model, payload)
+
+        image_tensors = []
+        for idx, img_content in enumerate(choices[0]["message"]["content"]):
+            base64_str = img_content.get("image_url", {}).get("url", "")
+            if not base64_str:
+                raise RuntimeError(f"API returned image #{idx} without image url")
+            tensor = base64_to_image_tensor(base64_str)
+            image_tensors.append(tensor)
+
+        return torch.stack(image_tensors, dim=0)
+
+    async def generate_comprehension_chat_completion(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] = ["text", "audio"],
+        **extra_body,
+    ) -> tuple[str | None, AudioInput | None]:
+        # Response may contain two choices: one with text, one with audio
+        payload = VLLMOmniClient._prepare_chat_completion_messages(
+            model=model,
+            prompt=prompt,
+            negative_prompt=None,
+            image=image,
+            audio=audio,
+            video=video,
+            sampling_params=sampling_params,
+            modalities=modalities,
+            **extra_body,
+        )
+
+        choices = await self._generate_base_chat_completion(model, payload)
+        text_response = None
+        audio_base64 = None
+        for choice in choices:
+            try:
+                text_response = choice["message"]["content"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+            try:
+                audio_base64 = choice["message"]["audio"]["data"]
+            except (KeyError, TypeError):
+                # Either this case (text response) or the audio response case will be hit. Checking None's later.
+                pass
+        if audio_base64 is None and text_response is None:
+            raise RuntimeError(
+                "API response missing both '.message.audio' and 'message.content' fields."
+                f"The choices object is {choices}"
+            )
+        if audio_base64 is not None:
+            audio = base64_to_audio(audio_base64)
+            logger.debug(
+                "audio sample rate %d, audio shape %s, duration in second %f",
+                audio["sample_rate"],
+                audio["waveform"].shape,
+                audio["waveform"].shape[2] / audio["sample_rate"],
+            )
+        else:
+            audio = None
+        return text_response, audio
+
+    async def generate_speech(
+        self,
+        *,
+        model: str,
+        input: str,
+        voice: str,
+        response_format: AudioFormat,
+        speed: float,
+        **extra_params,
+    ) -> AudioInput:
+        await self._check_model_exist(model)
+
+        ref_audio: AudioInput | None = extra_params.pop("ref_audio", None)
+
+        payload = {
+            "model": model,
+            "input": input,
+            "voice": voice,
+            "response_format": response_format,
+            "speed": speed,
+            **extra_params,
+        }
+
+        if ref_audio is not None:
+            audio_base64 = audio_to_base64(ref_audio)
+            payload["ref_audio"] = audio_base64
+
+        logger.debug("Omni TTS payload: %s", pretty_printer.pformat(payload))
+
+        url = self.base_url + "/audio/speech"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) 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:
+                        audio_bytes = await response.read()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+
+                    try:
+                        audio = bytes_to_audio(audio_bytes)
+                    except av.error.InvalidDataError as e:
+                        raise ValueError(
+                            f"Invalid audio data received from vLLM-Omni: {e}"
+                            "Check if you have input unsupported arguments (such as 'voice')"
+                        )
+                    return audio
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+    async def _generate_base_chat_completion(self, model: str, payload: dict[str, Any]) -> list[dict[str, Any]]:
+        logger.debug("Omni payload: %s", pretty_printer.pformat(payload))
+        await self._check_model_exist(model)
+
+        url = self.base_url + "/chat/completions"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.post(
+                    url,
+                    json=payload,
+                    headers={"Content-Type": "application/json"},
+                ) 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}")
+
+                    logger.debug(
+                        "chat completion response: %s",
+                        pretty_printer.pformat(data),
+                    )
+
+                    try:
+                        return data["choices"]
+                    except (KeyError, TypeError):
+                        raise RuntimeError("Invalid JSON response from vLLM-Omni: missing 'choices' field")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+
+    async def _check_model_exist(self, model: str):
+        url = self.base_url + "/models"
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            try:
+                async with session.get(
+                    url,
+                    headers={"Content-Type": "application/json"},
+                ) 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} "
+                            f"when getting hosted model list: {error_text}"
+                        )
+
+                    try:
+                        data = await response.json()
+                    except aiohttp.ContentTypeError as e:
+                        raise RuntimeError(f"Invalid JSON response when getting hosted model list from vLLM-Omni: {e}")
+
+            except aiohttp.ClientError as e:
+                raise RuntimeError(f"Network error connecting to vLLM-Omni at {self.base_url}: {e}")
+        try:
+            model_list = data["data"]
+            model_found = next((True for m in model_list if m["id"] == model), False)
+        except (KeyError, TypeError):
+            raise RuntimeError(f"Invalid JSON response of the hosted model list: {data}")
+
+        if not model_found:
+            raise ValueError(f"Model {model} not served at {self.base_url}.")
+
+    @staticmethod
+    def _prepare_chat_completion_messages(
+        *,
+        model: str,
+        prompt: str,
+        negative_prompt: str | None,
+        image: torch.Tensor | None = None,
+        audio: AudioInput | None = None,
+        video: VideoInput | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        modalities: list[str] | None = None,  # diffusion don't have this field
+        **extra_body,
+    ):
+        message_content: list[dict] = [{"type": "text", "text": prompt}]
+        if image is not None:
+            message_content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": image_tensor_to_base64(image)},
+                }
+            )
+        if audio is not None:
+            message_content.append({"type": "audio_url", "audio_url": {"url": audio_to_base64(audio)}})
+        if video is not None:
+            message_content.append({"type": "video_url", "video_url": {"url": video_to_base64(video)}})
+        messages = [{"role": "user", "content": message_content}]
+
+        combined_extra_body: dict[str, Any] = {}
+        if sampling_params is not None:
+            spec, _ = lookup_model_spec(model)
+            is_single_sampling_param = isinstance(sampling_params, dict) or len(sampling_params) == 1
+
+            # Exclude internal key
+            if isinstance(sampling_params, dict):
+                sampling_params = {k: v for k, v in sampling_params.items() if k != "type"}
+            else:
+                sampling_params = [{k: v for k, v in sp.items() if k != "type"} for sp in sampling_params]
+
+            if (spec is None and is_single_sampling_param) or (spec is not None and spec["stages"] == ["diffusion"]):
+                # Diffusion format: extra_body directly contains sampling params.
+                # Validation should have taken care of matching sampling params' types.
+                # * Use this mode if the model is a simple one-stage diffusion model.
+                # * Fallback to this mode if model is not registered and a single sampling param is provided.
+                sampling_params = sampling_params if isinstance(sampling_params, dict) else sampling_params[0]
+                combined_extra_body: dict[str, Any] = {**sampling_params}
+                if "n" in combined_extra_body:
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body["n"]
+                    del combined_extra_body["n"]
+            else:
+                # Use AR style payload, extra_body has a sampling_params_list field
+                combined_extra_body: dict[str, Any] = {"sampling_params_list": sampling_params}
+
+        if negative_prompt:
+            combined_extra_body["negative_prompt"] = negative_prompt
+
+        if extra_body:
+            combined_extra_body.update(extra_body)
+
+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if combined_extra_body:
+            payload["extra_body"] = combined_extra_body
+        if modalities:
+            payload["modalities"] = modalities
+
+        spec, _ = lookup_model_spec(model)
+        if spec:
+            preprocessor = spec.get("payload_preprocessor", None)
+            if preprocessor is not None:
+                payload = preprocessor(payload)
+
+        return payload

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

@@ -0,0 +1,230 @@
+"""Image/tensor format helpers.
+
+The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@dougbtv).
+Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
+"""
+
+import base64
+import mimetypes
+from io import BytesIO
+
+import av
+import numpy as np
+import torch
+from comfy_api.input import AudioInput, VideoInput
+from comfy_extras import nodes_audio
+from PIL import Image
+
+from .logger import get_logger
+
+logger = get_logger(__name__)
+
+
+def base64_to_image_tensor(base64_str: str, mode: str = "RGB") -> torch.Tensor:
+    """
+    Convert base64-encoded image to ComfyUI image tensor.
+
+    Args:
+        base64_str: Base64-encoded image string
+        mode: PIL image mode (default RGB for transparency support)
+
+    Returns:
+        torch.Tensor with shape (1, H, W, C) in float32 [0, 1] range
+
+    Raises:
+        ValueError: If base64 string is invalid or image cannot be decoded
+    """
+    if base64_str.startswith("data:image"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        image_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    # Create BytesIO object for PIL
+    image_bytesio = BytesIO(image_bytes)
+
+    # Open with PIL and convert to desired mode
+    try:
+        pil_image = Image.open(image_bytesio)
+        pil_image = pil_image.convert(mode)
+    except Exception as e:
+        raise RuntimeError(f"Failed to open image: {e}")
+
+    image_array = np.asarray(pil_image).astype(np.float32) / 255.0
+    image_tensor = torch.from_numpy(image_array)
+    return image_tensor
+
+
+def image_tensor_to_png_bytes(tensor: torch.Tensor, filename: str = "image.png") -> BytesIO:
+    """
+    Convert ComfyUI image tensor to PNG BytesIO for multipart upload.
+
+    This function converts a ComfyUI IMAGE tensor to a PNG-encoded BytesIO object
+    suitable for multipart/form-data upload. The BytesIO object has its .name
+    attribute set, which is required by aiohttp for file uploads.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+
+    Returns:
+        BytesIO object containing PNG-encoded image with .name attribute set
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    if tensor.ndim != 4:
+        raise ValueError(f"Expected 4D tensor with shape (B, H, W, C), got {tensor.ndim}D tensor")
+
+    image_tensor = tensor[0]  # Shape: (H, W, C)
+    image_np = (image_tensor.cpu().numpy() * 255.0).astype(np.uint8)
+    pil_image = Image.fromarray(image_np)
+
+    # Save to BytesIO as image file
+    img_bytes = BytesIO()
+    # Set name attribute (required for multipart upload and mimetype detection)
+    img_bytes.name = filename
+    try:
+        pil_image.save(img_bytes)
+    except Exception as e:
+        raise RuntimeError(f"Failed to save image as file: {e}")
+
+    # Reset position to beginning
+    img_bytes.seek(0)
+
+    return img_bytes
+
+
+def image_tensor_to_base64(tensor: torch.Tensor, filename: str = "image.png") -> str:
+    """
+    Convert ComfyUI image tensor to base64-encoded image string.
+
+    Args:
+        tensor: ComfyUI IMAGE tensor with shape (B, H, W, C), dtype float32, range [0, 1]
+        filename: Name attribute to set on BytesIO (default: "image.png")
+        format: File format of the output image file buffer (default: "PNG")
+
+    Returns:
+        Base64-encoded image string
+
+    Raises:
+        ValueError: If tensor format is invalid (not 4D, wrong dtype, etc.)
+    """
+    img_bytes = image_tensor_to_png_bytes(tensor, filename)
+    img_bytes.seek(0)
+    byte_data = img_bytes.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def video_to_bytes(video: VideoInput, filename: str = "video.mp4") -> BytesIO:
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    video.save_to(output_buffer)
+    output_buffer.seek(0)
+    return output_buffer
+
+
+def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
+    video_buffer = video_to_bytes(video, filename)
+    video_buffer.seek(0)
+    byte_data = video_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+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"]
+    format = filename.rsplit(".", maxsplit=1)[1]
+    layout = "mono" if waveform.shape[0] == 1 else "stereo"
+
+    output_buffer = BytesIO()
+    output_buffer.name = filename
+    output_container = av.open(output_buffer, mode="w", format=format)
+    if format == "opus":
+        out_stream = output_container.add_stream("libopus", rate=sample_rate, layout=layout)
+        if quality == "64k":
+            out_stream.bit_rate = 64000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "96k":
+            out_stream.bit_rate = 96000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "192k":
+            out_stream.bit_rate = 192000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    elif format == "mp3":
+        out_stream = output_container.add_stream("libmp3lame", rate=sample_rate, layout=layout)
+        if quality == "V0":
+            out_stream.codec_context.qscale = 1  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "128k":
+            out_stream.bit_rate = 128000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+        elif quality == "320k":
+            out_stream.bit_rate = 320000  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    else:  # format == "flac":
+        out_stream = output_container.add_stream("flac", rate=sample_rate, layout=layout)
+
+    frame = av.AudioFrame.from_ndarray(
+        waveform.movedim(0, 1).reshape(1, -1).float().numpy(),
+        format="flt",
+        layout=layout,
+    )
+    frame.sample_rate = sample_rate
+    frame.pts = 0
+    output_container.mux(out_stream.encode(frame))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    # Flush encoder
+    output_container.mux(out_stream.encode(None))  # type: ignore # copy from ComfyUI comfy_api/latest/_ui.py
+    output_container.close()
+    output_buffer.seek(0)
+
+    return output_buffer
+
+
+def audio_to_base64(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> str:
+    audio_buffer = audio_to_bytes(audio, filename, quality)
+    audio_buffer.seek(0)
+    byte_data = audio_buffer.read()
+    base64_str = base64.b64encode(byte_data).decode("utf-8")
+    mime_type = mimetypes.guess_type(filename)[0] or "application/octet-stream"
+    return f"data:{mime_type};base64,{base64_str}"
+
+
+def bytes_to_audio(audio_bytes: bytes) -> AudioInput:
+    """
+    Convert audio bytes to ComfyUI audio tensor.
+
+    Args:
+        audio_bytes: Audio file bytes
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    audio_buffer = BytesIO(audio_bytes)
+    waveform, sample_rate = nodes_audio.load(audio_buffer)  # type: ignore # Although expect string argument, it calls av.open underneath, which supports BytesIO (file-like)
+    return {"waveform": waveform.unsqueeze(0), "sample_rate": sample_rate}
+
+
+def base64_to_audio(base64_str: str) -> AudioInput:
+    """
+    Convert base64-encoded audio to ComfyUI audio tensor.
+
+    Args:
+        base64_str: Base64-encoded audio string
+    Returns:
+        torch.Tensor with shape (B, C, T) in float32 range [-1, 1]
+    """
+    if base64_str.startswith("data:audio"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        # Decode base64 to bytes
+        audio_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_audio(audio_bytes)

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

@@ -0,0 +1,125 @@
+"""Centralized logger configuration for vLLM-Omni ComfyUI."""
+
+import logging
+import pprint
+import sys
+from typing import Any
+
+
+def get_logger(name: str) -> logging.Logger:
+    """
+    Get or create a logger with proper formatting.
+
+    Args:
+        name: Logger name (typically __name__ of the calling module)
+
+    Returns:
+        Configured logger instance
+    """
+    logger = logging.getLogger(name)
+
+    # Only configure if not already configured
+    if not logger.handlers:
+        logger.setLevel(logging.DEBUG)
+
+        # Create console handler
+        handler = logging.StreamHandler(sys.stdout)
+        handler.setLevel(logging.INFO)
+
+        # Create formatter
+        formatter = logging.Formatter(
+            fmt="(ComfyUI-vLLM-Omni) [%(levelname)s] %(asctime)s [%(filename)s:%(lineno)s] %(message)s",
+            datefmt="%Y-%m-%d %H:%M:%S",
+        )
+        handler.setFormatter(formatter)
+
+        # Add handler to logger
+        logger.addHandler(handler)
+
+        # Prevent propagation to root logger
+        logger.propagate = False
+
+    return logger
+
+
+class OmitBase64PrettyPrinter(pprint.PrettyPrinter):
+    """
+    A PrettyPrinter that redacts specific field names with '...'
+    wherever they appear in nested structures.
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _format(self, obj: Any, stream, indent: int, allowance: int, context, level: int) -> None:
+        # Check if this is a dict with redacted keys
+        if isinstance(obj, dict):
+            # Create a copy with redacted values
+            display_obj = {}
+            for key, value in obj.items():
+                if key == "data" or key == "url":
+                    if value.startswith("data:"):
+                        base64_header = value.split(",", 1)[0]
+                        display_obj[key] = f"{base64_header},***"
+                    elif value.startswith("http://") or value.startswith("https://"):
+                        display_obj[key] = value
+                    elif len(value) > 10:
+                        display_obj[key] = f"{value[:10]}***"
+                    else:
+                        display_obj[key] = value
+                else:
+                    display_obj[key] = value
+            obj = display_obj
+
+        # Handle list/tuple/set containing dicts that might have redacted keys
+        # (pprint will recursively call _format on nested items, so this
+        # handles arbitrary nesting automatically)
+
+        super()._format(obj, stream, indent, allowance, context, level)
+
+
+pretty_printer = OmitBase64PrettyPrinter()
+
+
+# ========== PPrint EXAMPLES ==========
+
+if __name__ == "__main__":
+    data = {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "text",
+                        "text": "What sound is it, and what is the drawing about?",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": "data:image/png;base64,iVBORw0KGgoAAAANSUhEUgAAAwAAAA"},
+                    },
+                    {
+                        "type": "audio_url",
+                        "audio_url": {"url": "data:audio/mpeg;base64,SUQzBAAAAAAAIlRTU0UAAAAOAAADT="},
+                    },
+                ],
+            },
+        ],
+        "extra_body": {"mm_processor_kwargs": {"use_audio_in_video": False}},
+        "modalities": ["text"],
+    }
+
+    # Create printer that redacts 'password' and 'token' fields
+
+    print("\nRedactingPrettyPrinter (hides secrets):")
+    pretty_printer.pprint(data)

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

@@ -0,0 +1,119 @@
+import re
+
+from .types import Modality, ModelMode, Spec
+
+
+def _bagel_payload_preprocessor(payload: dict) -> dict:
+    try:
+        for message in payload["messages"]:
+            for content in message["content"]:
+                if content["type"] == "text":
+                    content["text"] = "<|im_start|>" + content["text"] + "<|im_end|>"
+    except (KeyError, TypeError):
+        raise RuntimeError("Internal Error: malformatted BAGEL payload")
+    extra_body = payload.pop("extra_body", {})
+    return {**payload, **extra_body}
+
+
+def _qwen25_payload_preprocessor(payload: dict) -> dict:
+    if payload["messages"][0]["role"] != "system":
+        payload["messages"] = [
+            {
+                "role": "system",
+                "content": (
+                    "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group,"
+                    "capable of perceiving auditory and visual inputs, as well as generating text and speech."
+                ),
+            },
+            *payload["messages"],
+        ]
+    return payload
+
+
+_MODEL_PIPELINE_SPECS: dict[str, Spec] = {
+    r"BAGEL-7B-MoT": {
+        "stages": [
+            "diffusion"  # The vLLM-Omni interface treats it as a single-stage diffusion model
+        ],
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [Modality.TEXT, Modality.IMAGE],
+            }
+        ],
+        "payload_preprocessor": _bagel_payload_preprocessor,
+    },
+    r"Qwen2.5-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "payload_preprocessor": _qwen25_payload_preprocessor,
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+    r"Qwen3-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "modes": [
+            {
+                "mode": ModelMode.COMPREHENSION,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+}
+# Convert dict keys to regex patterns
+MODEL_PIPELINE_SPECS: dict[re.Pattern, Spec] = {}
+for k, v in _MODEL_PIPELINE_SPECS.items():
+    MODEL_PIPELINE_SPECS[re.compile(k)] = v
+del _MODEL_PIPELINE_SPECS
+
+
+def lookup_model_spec(model: str) -> tuple[Spec | None, str | None]:
+    try:
+        last_component = model.rstrip("/").rsplit("/", 1)[-1]
+    except IndexError:
+        last_component = model
+    for pattern, spec in MODEL_PIPELINE_SPECS.items():
+        if pattern.search(last_component):
+            return spec, pattern.pattern
+    return None, None
+
+
+# ============== DEMONSTRATION ==============
+
+if __name__ == "__main__":
+    test_paths = [
+        "Qwen/Qwen2.5-Omni-7B",
+        "MyModels/Qwen2.5-Omni-3B",
+        "/root/home/Qwen2.5-Omni-7B",
+        "Qwen/Qwen3-Omni",
+        "Qwen/Qwen3-Omni-30B-A3B-Instruct",
+        "Custom/Path/UnknownModel-Instruct",
+        "Not/Matching/Anything",
+    ]
+
+    test_payload = {"messages": [{"role": "user", "content": "prompt"}]}
+
+    print("Testing registry lookups:\n")
+    for path in test_paths:
+        spec, _ = lookup_model_spec(path)
+        if spec:
+            if preprocessor := spec.get("payload_preprocessor"):
+                result = preprocessor(test_payload)
+                print(f"✓ {path:<40} → {result}")
+            else:
+                print(f"✓ {path:<40} → No preprocessor")
+        else:
+            print(f"✗ {path:<40} → No match")

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

@@ -0,0 +1,39 @@
+from collections.abc import Callable
+from enum import Enum, auto
+from typing import (
+    Any,
+    Literal,
+    NotRequired,
+    TypeAlias,
+    TypedDict,
+)
+
+AudioFormat: TypeAlias = Literal["mp3", "opus", "aac", "flac", "wav", "pcm"]
+
+
+class ModelMode(Enum):
+    IMAGE_GENERATION = auto()
+    VIDEO_GENERATION = auto()
+    AUDIO_GENERATION = auto()
+    COMPREHENSION = auto()
+
+
+class Modality(Enum):
+    TEXT = auto()  # maybe not useful. Prompt is always required
+    IMAGE = auto()
+    VIDEO = auto()
+    AUDIO = auto()
+
+
+class ModelModeSpec(TypedDict):
+    mode: ModelMode
+    input_modalities: list[Modality]
+
+
+PayloadPreprocessor: TypeAlias = Callable[[dict[str, Any]], dict[str, Any]]
+
+
+class Spec(TypedDict):
+    stages: list[Literal["diffusion", "autoregression"]]
+    modes: list[ModelModeSpec]
+    payload_preprocessor: NotRequired[PayloadPreprocessor]

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

@@ -0,0 +1,84 @@
+from .logger import get_logger
+from .models import lookup_model_spec
+
+logger = get_logger(__name__)
+
+
+def validate_model_and_sampling_params_types(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None = None,
+):
+    # Check if model name exists
+    if not model_name:
+        raise ValueError("Model name must not be empty.")
+
+    # Skip if no spec or no sampling params
+    pipeline_spec, _ = lookup_model_spec(model_name)
+    if pipeline_spec is None:
+        logger.info("skipping sampling params check because spec is not found")
+        return
+    if sampling_param_list is None:
+        return
+
+    # Check the number of stages and their data types
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, list):
+        # Check that the lengths match
+        if len(stages) != len(sampling_param_list):
+            raise ValueError(
+                f"Sampling parameter list length {len(sampling_param_list)} does not match "
+                f"number of stages {len(stages)} for model {model_name}."
+            )
+        # Check that each stage's type match
+        for i, sp in enumerate(sampling_param_list):
+            if "type" not in sp:
+                raise RuntimeError("Internal error: unknown sampling parameter type")
+            if sp["type"] != stages[i]:
+                raise ValueError(
+                    f"Sampling parameter type ({sp['type']}) does not match "
+                    f"stage type ({stages[i]}) at index {i} for model {model_name}."
+                )
+    elif isinstance(sampling_param_list, dict):
+        if "type" not in sampling_param_list:
+            raise RuntimeError("Internal error: unknown sampling parameter type")
+        # Check that the provided single sampling param matches all stages
+        elif any(stage != sampling_param_list["type"] for stage in stages):
+            raise ValueError(
+                f"When passing a single sampling parameter node, all stages of the model must match "
+                f"the provided sampling parameter's type. "
+                f"However, the stages of model {model_name} are: {stages}. "
+                f"The provided sampling parameter is {sampling_param_list['type']}"
+            )
+
+
+def add_sampling_parameters_to_stage(
+    model_name: str,
+    sampling_param_list: dict | list[dict] | None,
+    stage_type: str,
+    /,
+    **params_to_add,
+) -> dict | list[dict]:
+    """
+    Given a model's name and the sampling parameter list to query this model,
+    add arbitrary additional parameters to the sampling parameters of all stages of the given type.
+    """
+    pipeline_spec, _ = lookup_model_spec(model_name)
+    if not pipeline_spec:
+        logger.warning(
+            f"Since the model {model_name} is not in our list, we cannot ensure if "
+            f"the fields ({tuple(params_to_add.keys())}) are added to the correct stage's sampling params. "
+            f"We will do it heuristically."
+        )
+        pipeline_spec = {"stages": ["diffusion"]}
+
+    stages = pipeline_spec["stages"]
+    if isinstance(sampling_param_list, dict):
+        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)
+
+    return sampling_param_list

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Annotated Example.json

@@ -0,0 +1 @@
+{"id":"1c99f525-0a37-45ba-a28a-7df7c3af66b4","revision":0,"last_node_id":12,"last_link_id":14,"nodes":[{"id":1,"type":"VLLMOmniComprehension","pos":[1191.2177053682556,144.66829928181377],"size":[400,268],"flags":{},"order":8,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":1},{"localized_name":"video","name":"video","shape":7,"type":"VIDEO","link":2},{"localized_name":"audio","name":"audio","shape":7,"type":"AUDIO","link":3},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":14},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"output_text","name":"output_text","type":"BOOLEAN","widget":{"name":"output_text"},"link":null},{"localized_name":"output_audio","name":"output_audio","type":"BOOLEAN","widget":{"name":"output_audio"},"link":null},{"localized_name":"use_audio_in_video","name":"use_audio_in_video","type":"BOOLEAN","widget":{"name":"use_audio_in_video"},"link":null}],"outputs":[{"localized_name":"text_response","name":"text_response","type":"STRING","links":[8]},{"localized_name":"audio_response","name":"audio_response","type":"AUDIO","links":[9]}],"properties":{"Node name for S&R":"VLLMOmniComprehension"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen2.5-Omni-7B","",true,true,true]},{"id":3,"type":"LoadVideo","pos":[729.5984141255855,-198.631920454299],"size":[282.798828125,233.0743408203125],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"file","name":"file","type":"COMBO","widget":{"name":"file"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"VIDEO","name":"VIDEO","type":"VIDEO","links":[2]}],"properties":{"Node name for S&R":"LoadVideo"},"widgets_values":["draw.mp4","image"]},{"id":4,"type":"LoadAudio","pos":[729.8037086965753,99.86963519703949],"size":[282.798828125,136],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":5,"type":"VLLMOmniARSampling","pos":[510.3517536642828,658.073751009259],"size":[270,178],"flags":{},"order":2,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[12]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":7,"type":"VLLMOmniARSampling","pos":[503.33235181647115,419.34158016181806],"size":[270,178],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"max_tokens","name":"max_tokens","type":"INT","widget":{"name":"max_tokens"},"link":null},{"localized_name":"temperature","name":"temperature","type":"FLOAT","widget":{"name":"temperature"},"link":null},{"localized_name":"top_p","name":"top_p","type":"FLOAT","widget":{"name":"top_p"},"link":null},{"localized_name":"repetition_penalty","name":"repetition_penalty","type":"FLOAT","widget":{"name":"repetition_penalty"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"AR sampling params","name":"AR sampling params","type":"SAMPLING_PARAMS","links":[5,13]}],"properties":{"Node name for S&R":"VLLMOmniARSampling"},"widgets_values":[100,1,1,1,-1,"randomize"]},{"id":8,"type":"VLLMOmniSamplingParamsList","pos":[820.6056617389042,426.38372037182273],"size":[263.066015625,66],"flags":{},"order":7,"mode":0,"inputs":[{"localized_name":"param1","name":"param1","type":"SAMPLING_PARAMS","link":5},{"localized_name":"param2","name":"param2","shape":7,"type":"SAMPLING_PARAMS","link":12},{"localized_name":"param3","name":"param3","shape":7,"type":"SAMPLING_PARAMS","link":13}],"outputs":[{"localized_name":"param list","name":"param list"Line truncated

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Image Generation.json

@@ -0,0 +1 @@
+{"id":"91f75acc-8040-40f6-865a-2e8a7cfd6672","revision":0,"last_node_id":11,"last_link_id":23,"nodes":[{"id":3,"type":"PreviewImage","pos":[1281.8455167767304,-69.02638461454333],"size":[305.6628685610174,307.0336884172169],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"images","name":"images","type":"IMAGE","link":22}],"outputs":[],"properties":{"Node name for S&R":"PreviewImage"},"widgets_values":[]},{"id":11,"type":"VLLMOmniGenerateImage","pos":[816.0962515873642,-137.51112584854445],"size":[400,278],"flags":{},"order":3,"mode":0,"inputs":[{"localized_name":"image","name":"image","shape":7,"type":"IMAGE","link":21},{"localized_name":"mask","name":"mask","shape":7,"type":"MASK","link":null},{"localized_name":"sampling_params","name":"sampling_params","shape":7,"type":"SAMPLING_PARAMS","link":23},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"prompt","name":"prompt","type":"STRING","widget":{"name":"prompt"},"link":null},{"localized_name":"negative_prompt","name":"negative_prompt","type":"STRING","widget":{"name":"negative_prompt"},"link":null},{"localized_name":"width","name":"width","type":"INT","widget":{"name":"width"},"link":null},{"localized_name":"height","name":"height","type":"INT","widget":{"name":"height"},"link":null}],"outputs":[{"localized_name":"image","name":"image","type":"IMAGE","links":[22]}],"properties":{"Node name for S&R":"VLLMOmniGenerateImage"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen-Image-Edit","Put this figure in a realistic mountain view","",512,512]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,202],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"n","name":"n","type":"INT","widget":{"name":"n"},"link":null},{"localized_name":"num_inference_steps","name":"num_inference_steps","type":"INT","widget":{"name":"num_inference_steps"},"link":null},{"localized_name":"guidance_scale","name":"guidance_scale","type":"FLOAT","widget":{"name":"guidance_scale"},"link":null},{"localized_name":"true_cfg_scale","name":"true_cfg_scale","type":"FLOAT","widget":{"name":"true_cfg_scale"},"link":null},{"localized_name":"vae_use_slicing","name":"vae_use_slicing","type":"BOOLEAN","widget":{"name":"vae_use_slicing"},"link":null},{"localized_name":"seed","name":"seed","type":"INT","widget":{"name":"seed"},"link":null}],"outputs":[{"localized_name":"diffusion sampling params","name":"diffusion sampling params","type":"SAMPLING_PARAMS","links":[23]}],"properties":{"Node name for S&R":"VLLMOmniDiffusionSampling"},"widgets_values":[4,50,1,1,false,42,"fixed"]},{"id":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: Task and Input","properties":{},"widgets_values":["vLLM-Omni nodes are categorized based on the output modality. The \"Generate Image\" node supports both text-to-image generation or image-to-image generation (a.k.a. image editing). The node will route to the correct endpoint depending on whether an input image is present or not."],"color":"#432","bgcolor":"#000"},{"id":4,"type":"LoadImage","pos":[496.31859627609606,-229.71277089860084],"size":[270,314],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"image","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"IMAGE","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"MASK","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]}],"links":[[21,4,0,11,0,"IMAGE"],[22,11,0,3,0,"IMAGE"],[23,8,0,11,2,"SAMPLING_PARAMS"]],"groups":[{"id":1,"title":"Input","bounding":[213.1706010087147,-313.20095750667554,560.8246471437027,407.4246532472182],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486894,"offset":[276.8427259980847,549.4885226647582]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni TTS.json

@@ -0,0 +1 @@
+{"id":"c0a11241-f0e6-4f5c-8488-b3d344a9df3d","revision":0,"last_node_id":7,"last_link_id":3,"nodes":[{"id":4,"type":"PreviewAudio","pos":[690.1759063763212,226.76444957981545],"size":[270,88],"flags":{},"order":5,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","link":2},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null}],"outputs":[],"properties":{"Node name for S&R":"PreviewAudio"},"widgets_values":[]},{"id":5,"type":"LoadAudio","pos":[-83.61093919682321,103.80510898465587],"size":[282.798828125,136],"flags":{},"order":0,"mode":0,"inputs":[{"localized_name":"audio","name":"audio","type":"COMBO","widget":{"name":"audio"},"link":null},{"localized_name":"audioUI","name":"audioUI","type":"AUDIO_UI","widget":{"name":"audioUI"},"link":null},{"localized_name":"choose file to upload","name":"upload","type":"AUDIOUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"AUDIO","name":"AUDIO","type":"AUDIO","links":[3]}],"properties":{"Node name for S&R":"LoadAudio"},"widgets_values":["Megan-Fox.mp3",null,null]},{"id":3,"type":"VLLMOmniVoiceClone","pos":[248.63466280064262,168.13112893706375],"size":[400,306],"flags":{},"order":4,"mode":0,"inputs":[{"localized_name":"ref_audio","name":"ref_audio","type":"AUDIO","link":3},{"localized_name":"model_specific_params","name":"model_specific_params","shape":7,"type":"TTS_PARAMS","link":1},{"localized_name":"url","name":"url","type":"STRING","widget":{"name":"url"},"link":null},{"localized_name":"model","name":"model","type":"STRING","widget":{"name":"model"},"link":null},{"localized_name":"input","name":"input","type":"STRING","widget":{"name":"input"},"link":null},{"localized_name":"voice","name":"voice","type":"STRING","widget":{"name":"voice"},"link":null},{"localized_name":"response_format","name":"response_format","type":"COMBO","widget":{"name":"response_format"},"link":null},{"localized_name":"speed","name":"speed","type":"FLOAT","widget":{"name":"speed"},"link":null},{"localized_name":"ref_text","name":"ref_text","type":"STRING","widget":{"name":"ref_text"},"link":null},{"localized_name":"x_vector_only_mode","name":"x_vector_only_mode","type":"BOOLEAN","widget":{"name":"x_vector_only_mode"},"link":null}],"outputs":[{"localized_name":"audio","name":"audio","type":"AUDIO","links":[2]}],"properties":{"Node name for S&R":"VLLMOmniVoiceClone"},"widgets_values":["http://localhost:8000/v1","Qwen/Qwen3-TTS-12Hz-1.7B-Base","Someone just spilled a cup of coffee on my jacket this morning in the subway. Now I have to wear this stained jacket for a whole day!","","mp3",1,"",true]},{"id":6,"type":"MarkdownNote","pos":[255.52667408769662,27.26842775454861],"size":[381.58204108637847,88],"flags":{},"order":2,"mode":0,"inputs":[],"outputs":[],"title":"Note: TTS Nodes","properties":{},"widgets_values":["Apart from the Voice Cloning node, there is also another TTS node for simple text-to-speech tasks (without extra inputs for reference audio)."],"color":"#432","bgcolor":"#000"},{"id":7,"type":"MarkdownNote","pos":[-196.99453003423977,598.3693693727904],"size":[382.4539509209102,88],"flags":{},"order":3,"mode":0,"inputs":[],"outputs":[],"title":"Note: Model-Specific Parameters","properties":{},"widgets_values":["TTS models often require some tailor-made parameters. If you need to customize these parameters, grab one that matches the model you are requesting from \"TTS Params\" subfolder."],"color":"#432","bgcolor":"#000"},{"id":2,"type":"VLLMOmniQwenTTSParams","pos":[-207.8925310560497,342.7942867650348],"size":[400,200],"flags":{},"order":1,"mode":0,"inputs":[{"localized_name":"task_type","name":"task_type","type":"COMBO","widget":{"name":"task_type"},"link":null},{"localized_name":"language","name":"language","type":"COMBO","widget":{"name":"language"},"link":null},{"localized_name":"instructions","name":"instructions","type":"STRING","widget":{"name":"instructions"},"link":null},{"localized_name":"max_new_tokens","name":"max_new_tokens","type":"INT","widget":{"name":"max_new_tokens"},"link":null}],"outputs":[{"localized_name":"Qwen TTS params","name":"Qwen TTS params","type":"TTS_PARAMS","links":[1]}],"properties":{"Node name for S&R":"VLLMOmniQwenTTSParams"},"widgets_values":["Base","Auto","Super angry",2048]}],"links":[[1,2,0,3,1,"TTS_PARAMS"],[2,3,0,4,0,"AUDIO"],[3,5,0,3,0,"AUDIO"]],"groups":[{"id":1,"title":"Model-Specific Parameters","bounding":[-228.38328407738007,264.42790274716276,443.9887293164586,439.6853171248897],"color":"#3f789e","font_size":24,"flags":{}}],"config":{},"extra":{"workflowRendererVersion":"LG","ds":{"scale":1.1469075819486911,"offset":[898.8819468322333,258.7179979718389]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/web/main.js

@@ -0,0 +1,21 @@
+/**
+ * @file This file is intended to add dynamic fields to vLLM-Omni nodes
+ * based on widget (in-node form fields) and input (connection link) values and changes.
+ * However, this functionality is currently disabled/commented out.
+ * Because it introduces too much complexity,
+ * and it may even conflict with the current backend (Python) validation for unknown reasons (pending ComfyUI upstream fixes).
+ */
+
+import { app } from "../../scripts/app.js";
+app.registerExtension({
+    name: "vllm.vllm_omni",
+    async beforeRegisterNodeDef(nodeType, nodeData, app) {
+        if (!nodeData.name.startsWith("VLLMOmni")) {
+            return
+        }
+        // Stub frontend plugin for now
+    },
+    async setup() {
+        console.info("vLLM-Omni Setup complete!")
+    },
+})

benchmarks/diffusion/diffusion_benchmark_serving.py

@@ -973,10 +973,10 @@ if __name__ == "__main__":
         "--base-url",
         type=str,
         default=None,
-        help="Base URL of the server (e.g., http://localhost:30000). Overrides host/port.",
+        help="Base URL of the server (e.g., http://localhost:8091). Overrides host/port.",
     )
     parser.add_argument("--host", type=str, default="localhost", help="Server host.")
-    parser.add_argument("--port", type=int, default=30000, help="Server port.")
+    parser.add_argument("--port", type=int, default=8091, help="Server port.")
     parser.add_argument("--model", type=str, default="default", help="Model name.")
     parser.add_argument(
         "--dataset",

docker/Dockerfile.ci

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

docker/Dockerfile.npu

@@ -1,5 +1,5 @@
-ARG VLLM_ASCEND_IMAGE=quay.nju.edu.cn/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.11.0rc2
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
@@ -8,7 +8,7 @@ WORKDIR ${APP_DIR}
 COPY . .
 
 # Install vllm-omni with dev dependencies
-RUN pip install --no-cache-dir -e ".[dev]"
+RUN pip install --no-cache-dir -e . --no-build-isolation
 
 ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
 

docker/Dockerfile.npu.a3

@@ -0,0 +1,15 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+
+COPY . .
+
+# Install vllm-omni with dev dependencies
+RUN pip install --no-cache-dir -e . --no-build-isolation
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []

docker/Dockerfile.rocm

@@ -1,9 +1,7 @@
-ARG BASE_IMAGE=rocm/vllm-dev:nightly_main_20251205
-FROM ${BASE_IMAGE}
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.15.0
+FROM ${BASE_IMAGE} AS final
 
 ARG COMMON_WORKDIR=/app
-ARG VLLM_VERSION=v0.12.0
-ARG PYTORCH_ROCM_ARCH="gfx942;gfx950"
 
 WORKDIR ${COMMON_WORKDIR}
 
@@ -13,30 +11,18 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-# Step 2: Reinstall vllm from source
-RUN python3 -m pip uninstall -y vllm && rm -rf vllm &&\
-    git clone https://github.com/vllm-project/vllm.git && \
-    cd vllm && \
-    git checkout ${VLLM_VERSION} && \
-    python3 -m pip install -r requirements/rocm.txt && \
-    python3 setup.py clean --all && \
-    PYTORCH_ROCM_ARCH=${PYTORCH_ROCM_ARCH} python3 setup.py develop && \
-    cd ../ && \
-    rm -rf vllm/.git
-
 RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
 
-# Step 3: Copy vllm-omni code and install without uv
+# Step 2: Copy vllm-omni code and install without uv
 COPY . ${COMMON_WORKDIR}/vllm-omni
-RUN cd ${COMMON_WORKDIR}/vllm-omni && python3 -m pip install --no-cache-dir ".[dev]"
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
 
-# Create python symlink
-# `GPU_ARCHS` is an environment variable that is used to set the GPU archs for the AITER.
-# This is needed to prevent the AITER automatic GPU arch detection from failing on MI325X.
-# The AITER version used in this dockerfile has issues with handling
-# the GPU archs of MI325X (CI machine) correctly. So we manually set the GPU archs here.
-# We reuse AITER_ROCM_ARCH from the base image to avoid duplication.
-ENV GPU_ARCHS=${AITER_ROCM_ARCH}
 RUN ln -sf /usr/bin/python3 /usr/bin/python
 
+CMD ["/bin/bash"]
+
 ENTRYPOINT []
+
+#Set entrypoint for vllm-openai official images
+FROM final AS vllm-openai
+ENTRYPOINT ["vllm", "serve", "--omni"]

docs/.nav.yml

@@ -7,19 +7,27 @@ nav:
   - Serving:
     - OpenAI-Compatible API:
       - Image Generation: serving/image_generation_api.md
+      - Image Edit: serving/image_edit_api.md
   - Examples:
     - examples/README.md
     - Offline Inference:
+      - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.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
+      - LoRA Inference Examples: user_guide/examples/offline_inference/lora_inference.md
       - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
+      - Qwen3-TTS Offline Inference: user_guide/examples/offline_inference/qwen3_tts.md
+      - 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
     - Online Serving:
+      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
       - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
+      - Online LoRA Inference (Diffusion): user_guide/examples/online_serving/lora_inference.md
       - Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
       - Qwen3-Omni: user_guide/examples/online_serving/qwen3_omni.md
+      - Qwen3-TTS Online Serving: user_guide/examples/online_serving/qwen3_tts.md
       - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
   - General:
     - usage/*
@@ -36,6 +44,7 @@ nav:
       - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
       - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
       - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+    - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
     - contributing/README.md
@@ -52,6 +61,7 @@ nav:
     - Feature Design:
       - design/feature/disaggregated_inference.md
       - design/feature/ray_based_execution.md
+      - design/feature/omni_connectors/
     - Module Design:
       - design/module/ar_module.md
       - design/module/dit_module.md

docs/api/README.md

@@ -10,6 +10,9 @@ Main entry points for vLLM-Omni inference and serving.
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalContentParser][]
 - [vllm_omni.entrypoints.chat_utils.OmniAsyncMultiModalItemTracker][]
 - [vllm_omni.entrypoints.chat_utils.parse_chat_messages_futures][]
+- [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.log_utils.OrchestratorMetrics][]
@@ -26,7 +29,9 @@ Main entry points for vLLM-Omni inference and serving.
 
 Input data structures for multi-modal inputs.
 
+- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
 - [vllm_omni.inputs.data.OmniEmbedsPrompt][]
+- [vllm_omni.inputs.data.OmniTextPrompt][]
 - [vllm_omni.inputs.data.OmniTokenInputs][]
 - [vllm_omni.inputs.data.OmniTokensPrompt][]
 - [vllm_omni.inputs.parse.parse_singleton_prompt_omni][]
@@ -36,7 +41,6 @@ Input data structures for multi-modal inputs.
 
 Engine classes for offline and online inference.
 
-- [vllm_omni.diffusion.diffusion_engine.BackgroundResources][]
 - [vllm_omni.diffusion.diffusion_engine.DiffusionEngine][]
 - [vllm_omni.engine.AdditionalInformationEntry][]
 - [vllm_omni.engine.AdditionalInformationPayload][]
@@ -54,41 +58,17 @@ Engine classes for offline and online inference.
 
 Core scheduling and caching components.
 
+- [vllm_omni.core.sched.omni_ar_scheduler.KVCacheTransferData][]
 - [vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler][]
 - [vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler][]
+- [vllm_omni.core.sched.output.OmniCachedRequestData][]
 - [vllm_omni.core.sched.output.OmniNewRequestData][]
-
-## Model Executor
-
-Model execution components.
-
-- [vllm_omni.model_executor.custom_process_mixin.CustomProcessMixin][]
-- [vllm_omni.model_executor.models.output_templates.OmniOutput][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni.Qwen2_5OmniForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_talker.Qwen2_5OmniTalkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniConditionalGenerationMixin][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker.Qwen2_5OmniThinkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavBigVGANModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavDiTModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavForConditionalGenerationVLLM][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_token2wav.Qwen2_5OmniToken2WavModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2EmbeddingModel][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2ForCausalLM][]
-- [vllm_omni.model_executor.models.qwen2_5_omni.qwen2_old.Qwen2Model][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_moe.Qwen3MoeForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni.Qwen3OmniMoeForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_code2wav.Qwen3OmniMoeCode2Wav][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniCodePredictorBaseModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp.Qwen3OmniMoeTalkerCodePredictor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_talker.Qwen3OmniMoeTalkerSharedExpertWrapper][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMForCausalLM][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3MoeLLMModel][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeConditionalGenerationMixin][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerForConditionalGeneration][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerMultiModalProcessor][]
-- [vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker.Qwen3OmniMoeThinkerProcessingInfo][]
+- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.VectorQuantization][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.preprocess][]
 
 ## Configuration
 
@@ -98,25 +78,39 @@ Configuration classes.
 - [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
 - [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
 - [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSSpeakerEncoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerCodePredictorConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSTalkerConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2Config][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_12hz.configuration_qwen3_tts_tokenizer_v2.Qwen3TTSTokenizerV2DecoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1Config][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderBigVGANConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
+- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
 
 ## Workers
 
 Worker classes and model runners for distributed inference.
 
-- [vllm_omni.diffusion.worker.gpu_diffusion_model_runner.GPUDiffusionModelRunner][]
-- [vllm_omni.diffusion.worker.gpu_diffusion_worker.GPUDiffusionWorker][]
-- [vllm_omni.diffusion.worker.gpu_diffusion_worker.WorkerProc][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorker][]
-- [vllm_omni.diffusion.worker.npu.npu_worker.NPUWorkerProc][]
+- [vllm_omni.diffusion.worker.diffusion_model_runner.DiffusionModelRunner][]
+- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
+- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
+- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
+- [vllm_omni.platforms.npu.worker.npu_generation_model_runner.NPUGenerationModelRunner][]
+- [vllm_omni.platforms.npu.worker.npu_generation_worker.NPUGenerationWorker][]
+- [vllm_omni.platforms.npu.worker.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_model_runner.XPUARModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_ar_worker.XPUARWorker][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_model_runner.XPUGenerationModelRunner][]
+- [vllm_omni.platforms.xpu.worker.xpu_generation_worker.XPUGenerationWorker][]
 - [vllm_omni.worker.gpu_ar_model_runner.ExecuteModelState][]
 - [vllm_omni.worker.gpu_ar_model_runner.GPUARModelRunner][]
 - [vllm_omni.worker.gpu_ar_worker.GPUARWorker][]
 - [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
 - [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
 - [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
-- [vllm_omni.worker.npu.npu_ar_model_runner.ExecuteModelState][]
-- [vllm_omni.worker.npu.npu_ar_model_runner.NPUARModelRunner][]
-- [vllm_omni.worker.npu.npu_ar_worker.NPUARWorker][]
-- [vllm_omni.worker.npu.npu_generation_model_runner.NPUGenerationModelRunner][]
-- [vllm_omni.worker.npu.npu_generation_worker.NPUGenerationWorker][]
-- [vllm_omni.worker.npu.npu_model_runner.OmniNPUModelRunner][]
+- [vllm_omni.worker.mixins.OmniWorkerMixin][]

docs/cli/README.md

@@ -22,3 +22,21 @@ If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen2.5-Omni-7B --omni --stage-configs-path /path/to/stage_configs_file
 ```
+
+
+## bench
+
+Run benchmark tests for online serving throughput.
+Available Commands:
+
+```bash
+vllm bench serve --omni \
+    --model Qwen/Qwen2.5-Omni-7B \
+    --host server-host \
+    --port server-port \
+    --random-input-len 32 \
+    --random-output-len 4  \
+    --num-prompts  5
+```
+
+See [vllm bench serve](./bench/serve.md) for the full reference of all available arguments.

docs/cli/bench/serve.md

@@ -0,0 +1,359 @@
+# vLLM-Omni Benchmark CLI Guide
+The vllm bench command launches the vLLM-Omni benchmark to evaluate the performance of multimodal models.
+
+## Notes
+We currently only support using the "openai-chat-omni" backend.
+
+## Basic Parameter Description
+You can use `vllm bench serve --omni --help=all` to get descriptions of all parameters. The commonly used parameters are described below:
+- `--omni`  
+  Enable Omni (multimodal) mode, supporting multimodal inputs and outputs such as images, videos, and audio.
+
+- `--backend`  
+  Specify the backend adapter as openai-chat-omni, using OpenAI Chat compatible API behavior as the protocol. Currently only openai-chat-omni is supported.
+
+- `--model`  
+  The model identifier to load, filled according to the models supported by vLLM-Omni.
+
+- `--endpoint`  
+  The API endpoint exposed externally, to which clients send their requests.
+
+- `--dataset-name`  
+  The name of the dataset used; random-mm indicates generating random multimodal inputs (images, videos, audio).
+
+- `--num-prompts`  
+  The total number of requests to send, an integer.
+
+- `--max-concurrency`  
+  "Maximum number of concurrent requests. This can be used "
+        "to help simulate an environment where a higher level component "
+        "is enforcing a maximum number of concurrent requests. While the "
+        "--request-rate argument controls the rate at which requests are "
+        "initiated, this argument will control how many are actually allowed "
+        "to execute at a time. This means that when used in combination, the "
+        "actual request rate may be lower than specified with --request-rate, "
+        "if the server is not processing requests fast enough to keep up."
+
+- `--request-rate`  
+  "Number of requests per second. If this is inf, "
+        "then all the requests are sent at time 0. "
+        "Otherwise, we use Poisson process or gamma distribution "
+        "to synthesize the request arrival times."
+
+- `--ignore-eos`  
+  "Set ignore_eos flag when sending the benchmark request."
+
+- `--metric-percentiles`  
+  Comma-separated list of percentiles for selected metrics. "
+        "To report 25-th, 50-th, and 75-th percentiles, use \"25,50,75\". "
+        "Default value is \"99\"."
+        "Use \"--percentile-metrics\" to select metrics.
+
+- `--percentile-metrics`  
+        "Comma-separated list of selected metrics to report percentiles."
+                    "This argument specifies the metrics to report percentiles."
+                    'Allowed metric names are "ttft", "tpot", "itl", "e2el", "audio_ttfp", "audio_rtf", "audio_duration". '
+
+- `--save-result`  
+Specify to save benchmark results to a json file
+
+- `--save-detailed`  
+"When saving the results, whether to include per request "
+        "information such as response, error, ttfs, tpots, etc."
+
+- `--result-dir`  
+ "Specify directory to save benchmark json results."
+        "If not specified, results are saved in the current directory."
+
+- `--result-filename`  
+"Specify the filename to save benchmark json results."
+        "If not specified, results will be saved in "
+        "{label}-{args.request_rate}qps-{base_model_id}-{current_dt}.json"
+
+- `--random-prefix-len`  
+  Number of fixed prefix tokens before the random context in a request.
+  The total input length is the sum of random-prefix-len and a random
+  context length sampled from [input_len * (1 - range_ratio),
+  input_len * (1 + range_ratio)].Only the random and random-mm modes
+  support this parameter.
+
+- `--random-input-len`  
+  Number of input tokens per request.Only the random and random-mm modes support this parameter.
+
+- `--random-output-len`  
+  Number of output tokens per request.Only the random and random-mm modes support this parameter.
+
+- `--random-range-ratio`  
+  Range ratio for sampling input/output length,
+  used only for random sampling. Must be in the range [0, 1) to define
+  a symmetric sampling range
+  [length * (1 - range_ratio), length * (1 + range_ratio)].
+  Only the random and random-mm modes support this parameter.
+
+- `--random-mm-base-items-per-request`  
+  Base number of multimodal items per request for random-mm.
+  Actual per-request count is sampled around this base using
+  --random-mm-num-mm-items-range-ratio.
+  Only the random-mm mode supports this parameter.
+
+- `--random-mm-limit-mm-per-prompt`  
+  Per-modality hard caps for items attached per request, e.g.
+  '{"image": 3, "video": 1, "audio": 1}'. The sampled per-request item
+  count is clamped to the sum of these limits. When a modality
+  reaches its cap, its buckets are excluded and probabilities are
+  renormalized.
+  Only the random-mm mode supports this parameter.
+
+- `--random-mm-num-mm-items-range-ratio`  
+  Range ratio r in [0, 1] for sampling items per request.
+  We sample uniformly from the closed integer range
+  [floor(n*(1-r)), ceil(n*(1+r))]
+  where n is the base items per request.
+  r=0 keeps it fixed; r=1 allows 0 items. The maximum is clamped
+  to the sum of per-modality limits from
+  --random-mm-limit-mm-per-prompt.
+  An error is raised if the computed min exceeds the max.
+  Only the random-mm mode supports this parameter.
+
+- `--random-mm-bucket-config`  
+  The bucket config is a dictionary mapping a multimodal item
+  sampling configuration to a probability.
+  Currently allows for 3 modalities: audio, images and videos.
+  A bucket key is a tuple of (height, width, num_frames)
+  The value is the probability of sampling that specific item.
+  Example:
+  --random-mm-bucket-config
+  "{(256, 256, 1): 0.5, (720, 1280, 16): 0.4, (0, 1, 5): 0.10}"
+  First item: images with resolution 256x256 w.p. 0.5
+  Second item: videos with resolution 720x1280 and 16 frames
+  Third item: audios with 1s duration and 5 channels w.p. 0.1
+  OBS.: If the probabilities do not sum to 1, they are normalized.
+  Only the random-mm mode supports this parameter
+
+## Usage Examples
+
+### Online Benchmark
+<details class="admonition abstract" markdown="1">
+<summary>Show more</summary>
+
+First start serving your model:
+
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni
+```
+
+Then run the benchmarking for sharegpt:
+
+```bash
+# download dataset
+# wget https://huggingface.co/datasets/anon8231489123/ShareGPT_Vicuna_unfiltered/resolve/main/ShareGPT_V3_unfiltered_cleaned_split.json
+vllm bench serve \
+  --omni \
+  --port 43845 \
+  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --num-prompts 2 \
+  --dataset-name sharegpt \
+  --dataset-path ShareGPT_V3_unfiltered_cleaned_split.json \
+  --percentile-metrics ttft,tpot,itl,e2el
+```
+If successful, you will see the following output:
+```text
+============ Serving Benchmark Result ============
+Successful requests:                     2
+Failed requests:                         0
+Benchmark duration (s):                  81.63
+Request throughput (req/s):              0.02
+Peak concurrent requests:                2.00
+----------------End-to-end Latency----------------
+Mean E2EL (ms):                          56966.13
+Median E2EL (ms):                        56966.13
+P99 E2EL (ms):                           81016.80
+================== Text Result ===================
+Total input tokens:                      36
+Total generated tokens:                  5926
+Output token throughput (tok/s):         72.60
+Peak output token throughput (tok/s):    103.00
+Peak concurrent requests:                2.00
+Total Token throughput (tok/s):          73.04
+---------------Time to First Token----------------
+Mean TTFT (ms):                          124.76
+Median TTFT (ms):                        124.76
+P99 TTFT (ms):                           156.10
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          481.30
+Median TPOT (ms):                        481.30
+P99 TPOT (ms):                           947.55
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           25.11
+Median ITL (ms):                         0.33
+P99 ITL (ms):                            25.17
+================== Audio Result ==================
+Total audio duration generated(s):       3.95
+Total audio frames generated:            94890
+Audio throughput(audio duration/s):      0.05
+==================================================
+```
+
+Or run the benchmarking for random:
+
+```bash
+vllm bench serve \
+  --omni \
+  --port 43845 \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --dataset-name random \
+  --num-prompts 2 \
+  --random-prefix-len 5 \
+  --random-input-len 10 \
+  --random-output-len 100 \
+  --percentile-metrics ttft,tpot,itl,e2el,audio_ttfp,audio_rtf \
+  --ignore-eos
+```
+
+If successful, you will see the following output:
+
+```text
+============ Serving Benchmark Result ============
+Successful requests:                     2
+Failed requests:                         0
+Benchmark duration (s):                  3.89
+Request throughput (req/s):              0.51
+Peak concurrent requests:                2.00
+----------------End-to-end Latency----------------
+Mean E2EL (ms):                          3824.76
+Median E2EL (ms):                        3824.76
+P99 E2EL (ms):                           3888.54
+================== Text Result ===================
+Total input tokens:                      30
+Total generated tokens:                  10101
+Output token throughput (tok/s):         2595.57
+Peak output token throughput (tok/s):    111.00
+Peak concurrent requests:                2.00
+Total Token throughput (tok/s):          2603.28
+---------------Time to First Token----------------
+Mean TTFT (ms):                          117.15
+Median TTFT (ms):                        117.15
+P99 TTFT (ms):                           142.69
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          0.73
+Median TPOT (ms):                        0.73
+P99 TPOT (ms):                           0.74
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           16.47
+Median ITL (ms):                         16.19
+P99 ITL (ms):                            52.55
+================== Audio Result ==================
+Total audio duration generated(s):       15.79
+Total audio frames generated:            379050
+Audio throughput(audio duration/s):      4.06
+---------------Time to First Packet---------------
+Mean AUDIO_TTFP (ms):                    3701.37
+Median AUDIO_TTFP (ms):                  3701.37
+P99 AUDIO_TTFP (ms):                     3762.25
+-----------------Real Time Factor-----------------
+Mean AUDIO_RTF:                          0.47
+Median AUDIO_RTF:                        0.47
+P99 AUDIO_RTF:                           0.48
+==================================================
+```
+Notes:
+We use audio generation time / audio duration to calculate RTF.
+
+</details>
+
+### Multi-Modal Benchmark
+
+<details class="admonition abstract" markdown="1">
+<summary>Show more</summary>
+
+Benchmark the performance of multi-modal requests in vLLM-Omni.
+
+Generate synthetic image、video、audio inputs alongside random text prompts to stress-test vision models without external datasets.
+
+Notes:
+
+- Works only with online benchmark via the OpenAI backend (`--backend openai-chat-omni`) and endpoint `/v1/chat/completions`.
+
+Start the server (example):
+
+```bash
+vllm serve Qwen/Qwen2.5-Omni-7B --omni
+```
+
+It is recommended to use the flag `--ignore-eos` to simulate real responses. You can set the size of the output via the arg `random-output-len`.
+
+Then run the benchmarking script:
+```bash
+vllm bench serve \
+    --omni \
+  --dataset-name random-mm \
+  --port 40849 \
+  --model /home/models/Qwen/Qwen3-Omni-30B-A3B-Instruct \
+  --endpoint /v1/chat/completions \
+  --backend openai-chat-omni \
+  --request-rate 1 \
+  --num-prompts 1 \
+  --random-input-len 10 \
+  --random-range-ratio 0.0 \
+  --random-mm-base-items-per-request 2 \
+  --random-mm-num-mm-items-range-ratio 0 \
+  --random-mm-limit-mm-per-prompt '{"image":1,"video":1, "audio": 1}' \
+  --random-mm-bucket-config '{"(32, 32, 1)": 0.5, "(0, 1, 1)": 0.1, "(32, 32, 2)":0.4}' \
+  --ignore-eos \
+  --percentile-metrics ttft,tpot,itl \
+  --random-output-len 2 \
+  --extra_body '{"modalities": ["text"]}'
+```
+
+If successful, you will see the following output:
+
+```text
+============ Serving Benchmark Result ============
+Successful requests:                     1
+Failed requests:                         0
+Request rate configured (RPS):           1.00
+Benchmark duration (s):                  1.21
+Request throughput (req/s):              0.83
+Peak concurrent requests:                1.00
+================== Text Result ===================
+Total input tokens:                      10
+Total generated tokens:                  3
+Output token throughput (tok/s):         2.49
+Peak output token throughput (tok/s):    3.00
+Peak concurrent requests:                1.00
+Total Token throughput (tok/s):          10.77
+---------------Time to First Token----------------
+Mean TTFT (ms):                          179.74
+Median TTFT (ms):                        179.74
+P99 TTFT (ms):                           179.74
+-----Time per Output Token (excl. 1st token)------
+Mean TPOT (ms):                          12.76
+Median TPOT (ms):                        12.76
+P99 TPOT (ms):                           12.76
+---------------Inter-token Latency----------------
+Mean ITL (ms):                           12.76
+Median ITL (ms):                         12.76
+P99 ITL (ms):                            25.24
+================== Audio Result ==================
+Total audio duration generated(s):       0.00
+Total audio frames generated:            0
+Audio throughput(audio duration/s):      0.00
+==================================================
+```
+
+Behavioral notes:
+
+- If the requested base item count cannot be satisfied under the provided per-prompt limits, the tool raises an error rather than silently clamping.
+
+How sampling works:
+
+- Determine per-request item count k by sampling uniformly from the integer range defined by `--random-mm-base-items-per-request` and `--random-mm-num-mm-items-range-ratio`, then clamp k to at most the sum of per-modality limits.
+- For each of the k items, sample a bucket (H, W, T) according to the normalized probabilities in `--random-mm-bucket-config`, while tracking how many items of each modality have been added.
+- If a modality (e.g., image) reaches its limit from `--random-mm-limit-mm-per-prompt`, all buckets of that modality are excluded and the remaining bucket probabilities are renormalized before continuing.
+This should be seen as an edge case, and if this behavior can be avoided by setting `--random-mm-limit-mm-per-prompt` to a large number. Note that this might result in errors due to engine config `--limit-mm-per-prompt`.
+- The resulting request contains synthetic image data in `multi_modal_data` (OpenAI Chat format). When `random-mm` is used with the OpenAI Chat backend, prompts remain text and MM content is attached via `multi_modal_data`.
+</details>

docs/configuration/README.md

@@ -2,7 +2,7 @@
 
 This section lists the most common options for running vLLM-Omni.
 
-For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.12.0/configuration/index.html)
+For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
 
 Currently, the main options are maintained by stage configs for each model.
 

docs/configuration/stage_configs.md

@@ -40,7 +40,7 @@ stage_args:
     engine_args: # Engine arguments for a certain engine
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker # The specific worker used
+      worker_type: ar # The specific worker used
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
       gpu_memory_utilization: 0.8 # The gpu memory allocation for the stage within a single chip
       enforce_eager: true  # Now we only support eager mode
@@ -66,7 +66,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -92,7 +92,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

docs/configuration/stage_configs/qwen2_5_omni.yaml

@@ -8,7 +8,7 @@ stage_args:
     engine_args:
       model_stage: thinker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true  # Now we only support eager mode
@@ -34,7 +34,7 @@ stage_args:
     engine_args:
       model_stage: talker
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
+      worker_type: ar
       scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
       gpu_memory_utilization: 0.8
       enforce_eager: true
@@ -60,7 +60,7 @@ stage_args:
     engine_args:
       model_stage: code2wav
       model_arch: Qwen2_5OmniForConditionalGeneration
-      worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
+      worker_type: generation
       scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
       gpu_memory_utilization: 0.15
       enforce_eager: true

docs/contributing/model/adding_diffusion_model.md

@@ -44,7 +44,7 @@ vllm_omni/
 
 # Step-by-step-implementation
 ## Step 1: Model Implementation
-The diffusion pipeline’s implementation follows **HuggingFace Diffusers**, and components that do not need modification can be imported directly.
+The diffusion pipeline’s implementation follows **HuggingFace Diffusers**.
 ### 1.1 Define the Pipeline Class
 Define the pipeline class, e.g., `QwenImageEditPipeline`, and initialize all required submodules, either from HuggingFace `diffusers` or custom implementations. In `QwenImageEditPipeline`, only `QwenImageTransformer2DModel` is re-implemented to support optimizations such as Ulysses-SP. When adding new models in the future, you can either reuse this re-implemented `QwenImageTransformer2DModel` or extend it as needed.
 
@@ -73,10 +73,75 @@ def get_qwen_image_edit_post_process_func(
 The forward function of `QwenImageEditPipeline` follows the HuggingFace `diffusers` design for the most part. The key differences are:
 + As described in the overview, arguments are passed through `OnniDiffusionRequest`, so we need to get user parameters from it accordingly.
 ```python
-prompt = req.prompt if req.prompt is not None else prompt
+prompt = req.prompt
 ```
 + pre/post-processing are handled by the framework elsewhere, so skip them.
 
+### 1.4 Replace some ops or layers in DiT component
+
+vLLM-Omni provides a set of optimized operators with better performance and built-in support for parallelism, including attention, rotary embeddings (RoPE), and linear layers.
+
+Below is an example showing how to replace standard Transformer attention and FFN layers with vLLM-Omni implementations:
+
+```python
+from vllm_omni.diffusion.attention.layer import Attention
+from vllm_omni.diffusion.layers.rope import RotaryEmbedding
+
+class MyAttention(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.attn = Attention()
+        self.to_qkv = QKVParallelLinear()
+        self.to_out = RowParallelLinear()
+        self.rope = RotaryEmbedding(is_neox_style=False)
+
+    def forward(self, hidden_states):
+        qkv, _ = self.to_qkv(hidden_states)
+        q, k, v = qkv.split(...)
+        q, k = self.rope(...)
+        attn_output = self.attn(q, k, v)
+        output = self.to_out(attn_output)
+
+class MyFFN(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.fc1 = ColumnParallelLinear()
+        self.fc2 = RowParallelLinear()
+        self.act = F.gelu
+
+    def forward(self, hidden_states):
+        hidden, _ = self.fc1(hidden_states)
+        hidden = self.act(hidden)
+        output = self.fc2(hidden)
+        return output
+```
+
+In this example:
+
++ Attention uses vLLM-Omni’s optimized attention kernel together with parallel QKV projection and RoPE.
+
++ Linear layers are replaced with column- and row-parallel variants to enable tensor parallelism.
+
++ The FFN follows a standard two-layer structure and can be further optimized (e.g., using fused or merged projections) if needed.
+
+
+### 1.5 Provide a `_repeated_blocks` in DiT model
+`_repeated_blocks` is the small and frequently-repeated block(s) of a model -- typically a transformer layer.
+
+It's used for torch compile optimizations.
+```python
+_repeated_blocks = ["QwenImageTransformerBlock"]
+```
+
+
+### 1.6 (Optional) implement sequence parallelism
+vLLM-Omni has a non-intrusive `_sp_plan` that enable sequence parallel without modifying `forward()` logic.
+You can refer to [How to parallelize a new model](../../user_guide/diffusion/parallelism_acceleration.md)
+
+
+### 1.7 (Optional) integrate with Cache-Dit
+vLLM-Omni supports acceleration via [Cache-Dit](../../user_guide/diffusion/cache_dit_acceleration.md). Most models compatible with Diffusers can use Cache-Dit seamlessly. For new models, you can extend support by modifying`cache_dit_backend.py`
+
 ## Step 2: Extend OmniDiffusionRequest Fields
 User-provided inputs are ultimately passed to the model’s forward method through OmniDiffusionRequest, so we add the required fields here to support the new model.
 ```python
@@ -139,6 +204,165 @@ Key point for writing the example:
 
 + Save or display the generated results so users can validate the integration.
 
+## Step 5: TeaCache Coefficient Estimation (Optional)
+
+If your model supports TeaCache acceleration, you need to estimate the polynomial coefficients for optimal caching performance.
+
+### 5.1 Add Extractor Function
+
+First, implement an extractor function in `vllm_omni/diffusion/cache/teacache/extractors.py`. The extractor extracts the modulated input and defines how to run transformer blocks:
+
+```python
+def extract_your_model_context(
+    module: nn.Module,
+    hidden_states: torch.Tensor,
+    timestep: torch.Tensor,
+    **kwargs: Any,
+) -> CacheContext:
+    # 1. Preprocessing
+    temb = module.time_embed(timestep)
+
+    # 2. Extract modulated input (for cache decision)
+    modulated_input = module.transformer_blocks[0].norm1(hidden_states, temb)
+
+    # 3. Define transformer execution
+    def run_transformer_blocks():
+        h = hidden_states
+        for block in module.transformer_blocks:
+            h = block(h, temb=temb)
+        return (h,)
+
+    # 4. Define postprocessing
+    def postprocess(h):
+        return module.proj_out(module.norm_out(h, temb))
+
+    return CacheContext(
+        modulated_input=modulated_input,
+        hidden_states=hidden_states,
+        encoder_hidden_states=None,
+        temb=temb,
+        run_transformer_blocks=run_transformer_blocks,
+        postprocess=postprocess,
+    )
+```
+
+Register it in `EXTRACTOR_REGISTRY`:
+```python
+EXTRACTOR_REGISTRY = {
+    ...
+    "YourTransformer2DModel": extract_your_model_context,
+}
+```
+
+### 5.2 Add Adapter for Coefficient Estimation
+
+Add an adapter in `vllm_omni/diffusion/cache/teacache/coefficient_estimator.py`:
+
+```python
+class YourModelAdapter:
+    @staticmethod
+    def load_pipeline(model_path: str, device: str, dtype: torch.dtype) -> Any:
+        # Load your pipeline
+        ...
+
+    @staticmethod
+    def get_transformer(pipeline: Any) -> tuple[Any, str]:
+        return pipeline.transformer, "YourTransformer2DModel"
+
+    @staticmethod
+    def install_hook(transformer: Any, hook: DataCollectionHook) -> None:
+        registry = HookRegistry.get_or_create(transformer)
+        registry.register_hook(hook._HOOK_NAME, hook)
+
+_MODEL_ADAPTERS["YourModel"] = YourModelAdapter
+```
+
+### 5.3 Run Coefficient Estimation
+
+Use the provided script to estimate coefficients:
+
+```python
+from vllm_omni.diffusion.cache.teacache.coefficient_estimator import (
+    TeaCacheCoefficientEstimator,
+)
+from datasets import load_dataset
+from tqdm import tqdm
+
+# Load model
+estimator = TeaCacheCoefficientEstimator(
+    model_path="/path/to/model",
+    model_type="Bagel",  # Your model type
+    device="cuda",
+)
+
+# Load prompts (paper suggests ~70 prompts)
+dataset = load_dataset("nateraw/parti-prompts", split="train")
+prompts = dataset["Prompt"][:70]
+
+# Collect data
+for prompt in tqdm(prompts):
+    estimator.collect_from_prompt(prompt, num_inference_steps=50)
+
+# Estimate coefficients
+coeffs = estimator.estimate(poly_order=4)
+print(f"Coefficients: {coeffs}")
+```
+
+### 5.4 Interpreting Coefficient Estimation Results
+
+The estimator outputs statistics and polynomial coefficients. Here's how to interpret them:
+
+**Example Output:**
+```
+Data statistics:
+Count: 48
+Input Diffs (x): min=1.1089e-02, max=5.2555e-02, mean=2.8435e-02
+Output Diffs (y): min=2.8242e-02, max=2.9792e-01, mean=7.0312e-02
+Coefficients: [1333131.29, -168644.23, 7950.51, -163.75, 1.26]
+```
+
+**What to Check:**
+- **Count**: Number of timestep pairs analyzed. Should be at least 30-50 for reliable estimation. Low count suggests insufficient prompts or inference steps.
+- **Input/Output Ranges**: Verify output differences correlate with input differences. If ranges seem unusual, check your prompt diversity.
+- **Coefficient Magnitude**: Extremely large values (>1e8) may indicate numerical instability - try collecting more diverse data.
+
+**Troubleshooting:**
+- If results seem unreliable, try:
+  - Increasing number of prompts (100+ recommended)
+  - Using more diverse prompts from multiple datasets
+  - Adjusting `num_inference_steps` (try 20, 50, 100)
+
+### 5.5 Add Coefficients to Config
+
+Add the estimated coefficients to `vllm_omni/diffusion/cache/teacache/config.py`:
+
+```python
+_MODEL_COEFFICIENTS = {
+    ...
+    "YourTransformer2DModel": [
+        1.04730573e+06,  # a4
+        -1.34150749e+05, # a3
+        6.51517806e+03,  # a2
+        -1.41209108e+02, # a1
+        1.17241808e+00,  # a0
+    ],
+}
+```
+## Step 6: Open a Pull Request
+
+When submitting a pull request to add support for a new model, please include the following information in the PR description:
+
++ Output verification: provide generation outputs to verify correctness and model behavior.
+
++ Inference speed: provide a comparison with the corresponding implementation in Diffusers.
+
++ Parallelism support: specify the supported parallel sizes and any relevant limitations.
+
++ Cache acceleration: check whether the model can be accelerated using Cache-Dit or not.
+
+
+Providing these details helps reviewers evaluate correctness, performance improvements, and parallel scalability of the new model integration.
+
 # Testing
 For comprehensive testing guidelines, please refer to the [Test File Structure and Style Guide](../ci/tests_style.md).
 

docs/contributing/profiling.md

@@ -1,48 +1,30 @@
 # Profiling vLLM-Omni
 
-Performance Profiling Guidelines Profiling capabilities in vLLM-Omni are reserved for development and maintenance tasks aimed at temporal analysis of the codebase. Production use is **strongly discouraged**; enabling the profiler incurs a substantial overhead that negatively impacts inference latency.
+> **Warning:** Profiling incurs significant overhead. Use only for development and debugging, never in production.
 
-**Mechanism**: vLLM-Omni implements cross-stage profiling via the PyTorch Profiler. To accommodate the architecture—where stages operate as distinct engine instances in separate processes—the profiling interface supports both holistic capturing (all stages) and targeted capturing (specific stages).
+vLLM-Omni uses the PyTorch Profiler to analyze performance across both **multi-stage omni-modality models** and **diffusion models**.
 
-**1. Enabling the Profiler**
+### 1. Set the Output Directory
+Before running any script, set this environment variable. The system detects this and automatically saves traces here.
 
-Before running your script, you must set the ```VLLM_TORCH_PROFILER_DIR``` environment
-variable.
-
-```Bash
-export VLLM_TORCH_PROFILER_DIR=/path/to/save/traces
+```bash
+export VLLM_TORCH_PROFILER_DIR=./profiles
 ```
 
-**Highly Recommended: Limit Profiling to a Single Iteration**  
-For most use cases (especially when profiling audio stages), you should limit the profiler to just **one iteration** to keep trace files small and readable.
+### 2. Profiling Omni-Modality Models
 
+It is best to limit profiling to one iteration to keep trace files manageable.
 
 ```bash
 export VLLM_PROFILER_MAX_ITERS=1
 ```
 
-**2. Offline Inference**
-
-For offlinie processing using ```OmniLLM```, you can wrap your ```generate``` calls with ```start_profile``` and ```stop_profile()```.
-
-Basic Usage(All Stages)
-```Python
-from vllm_omni import OmniLLM
-
-omni_llm = OmniLLM.from_engine_args(engine_args)
-
-# Start profiling all active stages
-omni_llm.start_profile()
-
-outputs = omni_llm.generate(prompts, sampling_params)
-
-# Stop profiling and save traces
-omni_llm.stop_profile()
-```
-
 **Selective Stage Profiling**
 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()
+
 # Only profile Stage 1
 omni_llm.start_profile(stages=[1])
 ```
@@ -52,48 +34,116 @@ omni_llm.start_profile(stages=[1])
 omni_llm.start_profile(stages=[0, 2])
 ```
 
-**Examples**:
-
-1. **Qwen-omni 2.5**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
-
-2. **Qwen-omni 3.0**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
-
-**3. Online Inference(Async)**
-
-For online serving using AsyncOmni, the methods are asynchronous. This allows you to toggle profiling dynamically without restarting the server.
+**Python Usage**: Wrap your generation logic with `start_profile()` and `stop_profile()`.
 
 ```python
-from vllm_omni import AsyncOmni
+from vllm_omni import omni_llm
 
-# Inside an async function:
-async_omni = AsyncOmni.from_engine_args(engine_args)
+profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
 
-await async_omni.start_profile()
+# 1. Start profiling if enabled
+if profiler_enabled:
+    omni_llm.start_profile(stages=[0])
 
-async for output in async_omni.generate(prompt, sampling_params, request_id):
-    # Process outputs...
-    pass
+# Initialize generator
+omni_generator = omni_llm.generate(prompts, sampling_params_list, py_generator=args.py_generator)
 
-await async_omni.stop_profile()
+total_requests = len(prompts)
+processed_count = 0
+
+# Main Processing Loop
+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)
+
+    # 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()
+
+        # 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()
 ```
 
-**4. Analyzing Omni Traces**
 
-After ``stop_profile()`` completes (and the file write wait time has passed), the directory specified in ```VLLM_TORCH_PROFILER_DIR``` will contain the trace files.
+**Examples**:
+
+1. **Qwen2.5-Omni**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen2_5_omni/end2end.py)
+
+2. **Qwen3-Omni**:   [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/qwen3_omni/end2end.py)
+
+
+### 3. Profiling diffusion models
+
+Diffusion profiling is End-to-End, capturing encoding, denoising loops, and decoding.
+
+**CLI Usage:**
+```python
+
+python image_to_video.py \
+    --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
+    --image qwen-bear.png \
+    --prompt "A cat playing with yarn, smooth motion" \
+    \
+    # Minimize Spatial Dimensions (Optional but helpful):
+    #    Drastically reduces memory usage so the profiler doesn't
+    #    crash due to overhead, though for accurate performance
+    #    tuning you often want target resolutions.
+    --height 48 \
+    --width 64 \
+    \
+    # Minimize Temporal Dimension (Frames):
+    #    Video models process 3D tensors (Time, Height, Width).
+    #    Reducing frames to the absolute minimum (2) keeps the
+    #    tensor size small, ensuring the trace file doesn't become
+    #    multi-gigabytes in size.
+    --num_frames 2 \
+    \
+    # Minimize Iteration Loop (Steps):
+    #    This is the most critical setting for profiling.
+    #    Diffusion models run the same loop X times.
+    #    Profiling 2 steps gives you the exact same performance
+    #    data as 50 steps, but saves minutes of runtime and
+    #    prevents the trace viewer from freezing.
+    --num_inference_steps 2 \
+    \
+    --guidance_scale 5.0 \
+    --guidance_scale_high 6.0 \
+    --boundary_ratio 0.875 \
+    --flow_shift 12.0 \
+    --fps 16 \
+    --output i2v_output.mp4
 
 ```
-Output/
-│── ...rank-0.pt.trace.json.gz   # GPU 0 trace (TP=2 Example)
-│── ...rank-1.pt.trace.json.gz   # GPU 1 trace (TP=2 Example)
-│       # Load these into Perfetto to visualize synchronization
-│
-│── profiler_out_.txt            # Summary tables (CPU/CUDA time %)
-```
+
+**Examples**:
+
+1. **Qwen image edit**:  [https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py](https://github.com/vllm-project/vllm-omni/blob/main/examples/offline_inference/image_to_image/image_edit.py)
+
+2. **Wan-AI/Wan2.2-I2V-A14B-Diffusers**:   [https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video](https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/image_to_video)
+
+> **Note:**
+As of now, asynchronous (online) profiling is not fully supported in vLLM-Omni. While start_profile() and stop_profile() methods exist, they are only reliable in offline inference scripts (e.g., the provided end2end.py examples). Do not use them in server-mode or streaming scenarios—traces may be incomplete or fail to flush.
+
+### 4. Analyzing Omni Traces
+
+Output files are saved to your configured ```VLLM_TORCH_PROFILER_DIR```.
+
+**Output**
+**Chrome Trace** (```.json.gz```): Visual timeline of kernels and stages. Open in Perfetto UI.
 
 **Viewing Tools:**
-     - [Perfetto](https://ui.perfetto.dev/): (Recommended): Best for handling large audio trace files.
-     - ```chrome://tracing```: Good for smaller text-only traces.
 
+- [Perfetto](https://ui.perfetto.dev/)(recommended)
+- ```chrome://tracing```(Chrome only)
 
-**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM.  
-For more advanced configuration options (memory profiling, custom activities, etc.), see the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/latest/dev/profiling.html)
+**Note**: vLLM-Omni reuses the PyTorch Profiler infrastructure from vLLM. See the official vLLM profiler documentation:  [vLLM Profiling Guide](https://docs.vllm.ai/en/stable/contributing/profiling/)

docs/design/architecture_overview.md

@@ -192,4 +192,4 @@ curl -sS -X POST http://localhost:8091/v1/chat/completions \
 }
 ```
 
-For more usages, please refer to [examples](../examples/README.md).
+For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples).

docs/design/feature/disaggregated_inference.md

@@ -1,27 +1,35 @@
 # Disaggregated Inference for Omni-Modality Models
 
-This guide explains how to configure and use distributed connectors (vllm_omni/distributed/connectors) in vllm-omni for multi-stage pipelines.
+This guide explains how to configure and use distributed connectors
+(`vllm_omni/distributed/omni_connectors`) in vllm-omni for multi-stage pipelines.
 
-## 1. Overview
+Backend-specific setup lives in separate docs:
+
+- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
+- [MooncakeConnector](omni_connectors/mooncake_connector.md)
+- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
+
+## Overview
 
 Connectors enable data transfer between pipeline stages (e.g., Thinker -> Talker).
-Currently supported connectors operate in **D2H2D (Device-to-Host-to-Device)** mode:
-1. **SharedMemoryConnector**: Uses system shared memory.
-2. **MooncakeConnector**: Uses [Mooncake](https://github.com/kvcache-ai/Mooncake).
+Current connectors operate in D2H2D (device to host to device) mode.
 
-*   **SharedMemoryConnector (Default)**: Zero-copy (on host), lowest latency. Best for **single-node** deployments. Auto-configured if no connectors are specified.
-*   **MooncakeConnector**: TCP/RDMA based. Best for **multi-node** distributed deployments. Requires a Mooncake Master service.
+## Connector Choices
 
-## 2. API Design
+| Use Case | Recommended Connector | Notes |
+| :--- | :--- | :--- |
+| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
+| Multi node (Mooncake) | MooncakeConnector | Requires Mooncake Master + metadata server. |
+| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
 
-The connector system is built around the `OmniConnectorBase` abstraction, which decouples data transport from stage logic.
+## Core API
 
-### Core Interface
+The connector system is built around `OmniConnectorBase`.
 
 ```python
 class OmniConnectorBase(ABC):
     @abstractmethod
-    def put(self, from_stage: str, to_stage: str, request_id: str, data: Any) -> tuple[bool, int, Optional[dict]]:
+    def put(self, from_stage: str, to_stage: str, put_key: str, data: Any) -> tuple[bool, int, Optional[dict]]:
         """
         Store data.
         Returns: (success, serialized_size, metadata)
@@ -29,145 +37,23 @@ class OmniConnectorBase(ABC):
         pass
 
     @abstractmethod
-    def get(self, from_stage: str, to_stage: str, request_id: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
+    def get(self, from_stage: str, to_stage: str, get_key: str, metadata: Optional[dict] = None) -> Optional[tuple[Any, int]]:
         """
         Retrieve data.
-        Args: metadata - Transport-specific handles returned by put() (e.g., SHM name).
+        Args: metadata - transport-specific handles returned by put() (e.g., SHM name).
         Returns: (object, serialized_size)
         """
         pass
 ```
 
-### Key Concept: Metadata Passing
-Unlike a pure key-value store, some connectors (like `SharedMemoryConnector`) generate transient resources (e.g., a shared memory block name) during `put()`. This `metadata` **must be passed via the control plane** (e.g., HTTP headers, queue messages) from the producer stage to the consumer stage so `get()` can locate the data.
+### Metadata Passing
 
-## 3. Backends & Use Cases
+Some connectors (e.g., SharedMemoryConnector) generate transient resources during `put()`.
+This `metadata` must be passed through the control plane so `get()` can locate the data.
 
-### 3.1 SharedMemoryConnector
-**Best for:** Single-node, high-performance IPC.
+## Configuration Model
 
-*   **Mechanism:**
-    *   **Small Payloads (< Threshold)**: Data is serialized and passed directly "inline" within the `metadata` dictionary. This avoids the overhead of creating SHM blocks for tiny messages.
-    *   **Large Payloads (>= Threshold)**: Data is written to a named System Shared Memory block. The block name is returned in `metadata`.
-*   **Configuration:**
-    *   `shm_threshold_bytes`: Size in bytes to switch from inline to SHM (default: 64KB).
-
-### 3.2 MooncakeConnector
-**Best for:** Multi-node distributed inference.
-
-*   **Mechanism:** Uses Mooncake's distributed KVCache store.
-    *   **Data Plane**: TCP or RDMA for high-bandwidth transfer.
-    *   **Control Plane**: Uses a centralized Mooncake Master and Metadata Server.
-    *   **Keying**: Deterministic keys based on `request_id/from_stage_to_stage`.
-*   **Requirements**: Requires a running Mooncake Master service.
-
-## 4. Relationship with vLLM
-
-vLLM provides specialized distributed mechanisms for specific artifacts:
-*   **KV Transfer** (`vllm.distributed.kv_transfer`): Optimized for transferring KV caches between prefill and decode instances (using NCCL, Mooncake, etc.).
-*   **EC Transfer** (`vllm.distributed.ec_transfer`): Optimized for sharing encoder embeddings.
-*   **Device Communicators** (`vllm.distributed.device_communicators`): Low-level primitives (NCCL, SHM) for Tensor/Pipeline Parallelism.
-
-`vllm-omni` complements this by introducing a **Generalized Connector Abstraction** (`OmniConnector`) for multimodal pipelines. While vLLM's connectors are artifact-specific, `vllm-omni`:
-
-1.  **Unifies Transport**: Provides a single API (`put`/`get`) to transport *any* stage artifact (Input Embeddings, Hidden States, Audio/Image Tensors, KV Cache, Final Output) between arbitrary pipeline stages (e.g., AudioEncoder -> LLM -> AudioGenerator).
-2.  **Extends Connectivity**: Enables flexible construction of complex DAGs (Directed Acyclic Graphs) where stages can run in the same process, same node, or across nodes, using the most appropriate backend (SHM, Mooncake, etc.) for each edge.
-3.  **Wraps & Adapts**: Can internally utilize vLLM's specialized `kv_transfer` for KV paths while using generic transports (SHM/Mooncake) for other data types, presenting a consistent interface to the application layer.
-
-## 5. Installation (Mooncake)
-
-If using `MooncakeConnector`, install the library first:
-
-```bash
-# For CUDA-enabled systems (Recommended)
-pip install mooncake-transfer-engine
-
-# For non-CUDA systems
-pip install mooncake-transfer-engine-non-cuda
-```
-
-## 6. Using MooncakeConnector
-
-### 6.1 Start Mooncake Master
-
-Start the master service on your primary node:
-
-```bash
-# if you use mooncake SSD storage
-mkdir -p ./mc_storage
-
-mooncake_master \
-  --rpc_port=50051 \
-  --enable_http_metadata_server=true \
-  --http_metadata_server_host=0.0.0.0 \
-  --http_metadata_server_port=8080 \
-  --metrics_port=9003 \
-  --root_fs_dir=./mc_storage/ \
-  --cluster_id=mc-local-1 &
-```
-
-### 6.2 Configuration (YAML)
-
-Edit your stage config (e.g., `qwen2_5_omni.yaml`).
-
-**Step 1: Define Connector in Global Runtime**
-
-```yaml
-runtime:
-  connectors:
-    connector_of_mooncake:
-      name: MooncakeConnector
-      extra:
-        host: "127.0.0.1"           # Local Worker IP
-        metadata_server: "http://<MASTER_IP>:8080/metadata"
-        master: "<MASTER_IP>:50051"
-        segment: 512000000          # 512MB segment
-        localbuf: 64000000          # 64MB buffer
-        proto: "tcp"                # "tcp" or "rdma"
-    ```
-
-**Mooncake Configuration Parameters:**
-
-*   **host**: The hostname or IP address of the local machine (worker). Mooncake uses this to register itself in the metadata server so other nodes can find it.
-*   **metadata_server**: The URL of the metadata server. This is used for service discovery and connection establishment (e.g., exchanging QP information for RDMA).
-*   **master**: The address of the Mooncake Master Server (e.g., `<MASTER_IP>:50051`). This is used for global state management and control plane operations.
-*   **segment**: The size of the global memory segment in bytes (default: ~512MB). This defines the shared memory region accessible by Mooncake for data transfer.
-*   **localbuf**: The size of the local buffer in bytes (default: ~64MB). Used for local data buffering during transfer operations.
-*   **proto**: The transport protocol to use. Options:
-    *   `tcp`: Standard TCP/IP (easier setup, universal compatibility).
-    *   `rdma`: Remote Direct Memory Access (higher performance, requires RDMA-capable hardware).
-
-For more details, refer to the [Mooncake Repository](https://github.com/kvcache-ai/Mooncake).
-
-    **Step 2: Reference in Stages**
-
-Explicitly link stages using `input_connectors` and `output_connectors`:
-
-```yaml
-stage_args:
-  - stage_id: 0
-    # ...
-    output_connectors:
-      to_stage_1: connector_of_mooncake
-
-  - stage_id: 1
-    # ...
-    input_connectors:
-      from_stage_0: connector_of_mooncake
-```
-
-## 7. Using SharedMemoryConnector (Auto-Mode)
-
-**Best for single-node.**
-
-The system will automatically create `SharedMemoryConnector`s for any pipeline edge that does not have an explicit connector defined. This is inferred from:
-1.  `runtime.edges` list in the config.
-2.  `engine_input_source` dependencies defined in `stage_args`.
-
-### Threshold Configuration
-By default, payloads larger than **64KB** (default threshold) are transferred via shared memory, while smaller ones use the control queue (inline).
-
-To adjust this threshold (e.g., to 1GB), add the following to your `runtime.connectors`:
+Define connectors in runtime:
 
 ```yaml
 runtime:
@@ -175,26 +61,46 @@ runtime:
     connector_of_shared_memory:
       name: SharedMemoryConnector
       extra:
-        shm_threshold_bytes: 1024 # 1KB threshold
+        shm_threshold_bytes: 65536
 ```
 
-## 8. Summary
+Wire stages to connectors:
 
-| Use Case | Recommended Connector | Configuration |
-| :--- | :--- | :--- |
-| **Single Node** | `SharedMemoryConnector` | **None** (Automatic) or Custom Threshold |
-| **Multi Node** | `MooncakeConnector` | Explicit YAML + Mooncake Master |
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_shared_memory
 
-## 9. Operational Notes (important)
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_shared_memory
+```
 
-- **Fail-fast config validation**: the loader raises if any expected edge is missing a connector. Define `input_connectors`/`output_connectors` or rely on auto-SHM filling; otherwise startup aborts.
-- **Missing payloads halt the stage**: workers expect connector payloads; if metadata or connector config is missing, the stage raises and stops. Verify connector wiring and metadata propagation before production.
+If a pipeline edge has no explicit connector, the system auto-creates a
+SharedMemoryConnector for that edge.
 
-## 10. Future Roadmap: Device-to-Device (D2D) Transport
+## Relationship with vLLM
 
-The current implementations (`SharedMemoryConnector`, `MooncakeConnector`) utilize a **D2H2D (Device-to-Host-to-Device)** data path. Tensors are moved to CPU memory (Host) for transport, which incurs PCIe overhead.
+vLLM provides specialized distributed mechanisms for specific artifacts:
 
-As outlined in the design RFC, future versions will introduce **D2D (Device-to-Device)** connectors:
+- KV Transfer (`vllm.distributed.kv_transfer`): optimized for KV caches.
+- EC Transfer (`vllm.distributed.ec_transfer`): optimized for encoder embeddings.
+- Device Communicators (`vllm.distributed.device_communicators`): low-level primitives (NCCL, SHM).
 
-*   **Goal**: Direct GPU-to-GPU transfer (via NCCL, UCX, or IPC) to minimize latency for large tensor payloads.
-*   **Mechanism**: The `OmniConnector` API allows `put()` to initiate a transfer and return a lightweight handle (metadata) via the control plane, while the heavy payload flows directly between devices.
+vllm-omni complements this with a generalized connector abstraction:
+
+1. Unifies transport via a single `put`/`get` API for any stage artifact.
+2. Enables DAG-style pipelines across processes or nodes with per-edge transports.
+3. Can wrap vLLM-specific transfers for KV paths while keeping a consistent interface.
+
+## Operational Notes
+
+- Fail-fast config validation: missing expected edges cause startup failures.
+- Missing payloads halt stages: verify connector wiring and metadata propagation.
+
+## Future Roadmap: D2D Transport
+
+Current connectors use D2H2D paths. Future versions will introduce direct
+device-to-device connectors (NCCL, UCX, IPC) to reduce latency for large
+tensor payloads.

docs/design/feature/omni_connectors/mooncake_connector.md

@@ -0,0 +1,74 @@
+# MooncakeConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Mooncake.
+
+## Installation
+
+```bash
+# For CUDA-enabled systems (recommended)
+pip install mooncake-transfer-engine
+
+# For non-CUDA systems
+pip install mooncake-transfer-engine-non-cuda
+```
+
+## Start Mooncake Master
+
+```bash
+# If you use Mooncake SSD storage
+mkdir -p ./mc_storage
+
+mooncake_master \
+  --rpc_port=50051 \
+  --enable_http_metadata_server=true \
+  --http_metadata_server_host=0.0.0.0 \
+  --http_metadata_server_port=8080 \
+  --metrics_port=9003 \
+  --root_fs_dir=./mc_storage/ \
+  --cluster_id=mc-local-1 &
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_mooncake:
+      name: MooncakeConnector
+      extra:
+        host: "127.0.0.1"
+        metadata_server: "http://<MASTER_IP>:8080/metadata"
+        master: "<MASTER_IP>:50051"
+        segment: 512000000
+        localbuf: 64000000
+        proto: "tcp"
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_mooncake
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_mooncake
+```
+
+Parameters:
+
+- host: local worker IP registered in the metadata server.
+- metadata_server: metadata server URL for discovery and setup.
+- master: Mooncake Master address.
+- segment: global memory segment size in bytes.
+- localbuf: local buffer size in bytes.
+- proto: transport protocol ("tcp" or "rdma").
+
+For more details, refer to the
+[Mooncake repository](https://github.com/kvcache-ai/Mooncake).

docs/design/feature/omni_connectors/shared_memory_connector.md

@@ -0,0 +1,27 @@
+# SharedMemoryConnector
+
+## When to Use
+
+Best for single-node deployments where stages run on the same host. It is
+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.
+
+## Configuration
+
+```yaml
+runtime:
+  connectors:
+    connector_of_shared_memory:
+      name: SharedMemoryConnector
+      extra:
+        shm_threshold_bytes: 65536
+```
+
+## Notes
+
+- Auto-mode uses SharedMemoryConnector if no connector is declared for an edge.

docs/design/feature/omni_connectors/yuanrong_connector.md

@@ -0,0 +1,101 @@
+# YuanrongConnector
+
+## When to Use
+
+Best for multi-node distributed inference using Yuanrong Datasystem.
+
+## Mechanism
+
+Uses Yuanrong Datasystem's distributed KV store (`datasystem.kv_client`).
+
+- Data Plane: TCP or RDMA for high-bandwidth transfer.
+- Control Plane: Yuanrong Datasystem workers and etcd.
+- Keying: deterministic keys based on `put_key` (often composed as `request_id:fromStage_toStage`).
+
+## Installation
+
+```bash
+pip install openyuanrong-datasystem
+```
+
+## Start etcd
+
+```bash
+# Download and install etcd (v3.5.12 or higher)
+ETCD_VERSION="v3.5.12"
+ETCD_ARCH="linux-arm64"
+wget https://github.com/etcd-io/etcd/releases/download/${ETCD_VERSION}/etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+tar -xvf etcd-${ETCD_VERSION}-${ETCD_ARCH}.tar.gz
+cd etcd-${ETCD_VERSION}-${ETCD_ARCH}
+sudo cp etcd etcdctl /usr/local/bin/
+
+# Start etcd
+etcd \
+  --name etcd-single \
+  --data-dir /tmp/etcd-data \
+  --listen-client-urls http://0.0.0.0:2379 \
+  --advertise-client-urls http://0.0.0.0:2379 \
+  --listen-peer-urls http://0.0.0.0:2380 \
+  --initial-advertise-peer-urls http://0.0.0.0:2380 \
+  --initial-cluster etcd-single=http://0.0.0.0:2380 &
+
+# Verify etcd is running
+etcdctl --endpoints "127.0.0.1:2379" put key "value"
+etcdctl --endpoints "127.0.0.1:2379" get key
+```
+
+For production environments, refer to the
+[official etcd clustering documentation](https://etcd.io/docs/current/op-guide/clustering/).
+
+## Start Datasystem Worker
+
+```bash
+# Replace ${ETCD_IP} with etcd node IP, ${WORKER_IP} with local node IP
+dscli start -w \
+  --worker_address "${WORKER_IP}:31501" \
+  --etcd_address "${ETCD_IP}:2379" \
+  --shared_memory_size_mb 20480
+```
+
+To stop the worker:
+
+```bash
+dscli stop --worker_address "${WORKER_IP}:31501"
+```
+
+## Configuration
+
+Define the connector in runtime:
+
+```yaml
+runtime:
+  connectors:
+    connector_of_yuanrong:
+      name: YuanrongConnector
+      extra:
+        host: "127.0.0.1"
+        port: 31501
+        get_sub_timeout_ms: 1000
+```
+
+Wire stages to the connector:
+
+```yaml
+stage_args:
+  - stage_id: 0
+    output_connectors:
+      to_stage_1: connector_of_yuanrong
+
+  - stage_id: 1
+    input_connectors:
+      from_stage_0: connector_of_yuanrong
+```
+
+Parameters:
+
+- host: datasystem worker host.
+- port: datasystem worker port.
+- 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).

docs/features/comfyui.md

@@ -0,0 +1,70 @@
+# vLLM-Omni ComfyUI Integration
+
+vLLM-Omni offers a ComfyUI integration on top of its online serving API.
+It can send model inference requests to either a locally running vLLM-Omni service or a remote one.
+
+## Requirement
+
+- Python 3.12 or above
+- [ComfyUI installed](https://docs.comfy.org/installation/system_requirements)
+- [vLLM-Omni installed](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/) on either the same device or another device discoverable via the internet.
+- No need to install additional packages apart from those already required by ComfyUI.
+
+!!! tip
+    If you run both ComfyUI and vLLM-Omni on the same device, you can create separate virtual environments and use different Python versions for them.
+
+
+## Installation
+
+Copy the `apps/ComfyUI-vLLM-Omni` folder to the `custom_nodes` subfolder of your ComfyUI installation. Your directory should look like `ComfyUI/custom_nodes/ComfyUI-vLLM-Omni`.
+
+If you are running ComfyUI during copying, you should restart ComfyUI to load this extension.
+
+!!! tip
+    You can use utility websites such as https://download-directory.github.io/ to download a subdirectory of a repo. Also checkout community discussions (e.g., https://stackoverflow.com/questions/7106012/download-a-single-folder-or-directory-from-a-github-repository) for more info.
+
+On the device and virtual environment you run ComfyUI, launch ComfyUI with
+```bash
+cd ComfyUI
+
+# The regular way
+python main.py
+
+# If you are mainly using this node, launch it faster with
+python main.py --cpu
+```
+
+On the device and virtual environment you run vLLM-Omni, start a model service with
+```bash
+vllm serve The_Model_ID_to_Serve --omni --port 8000
+```
+
+Check **ComfyUI's sidebar -> Node Library**. There should be a new folder named **vLLM-Omni**.
+If no, check your shell running the ComfyUI process. There may be some error messages before the line `Import times for custom nodes:` and the line `To see the GUI go to: http://127.0.0.1:8188`.
+
+## Quickstart
+
+This extension offers the following nodes based on the output modalities:
+
+- **Generate Image** for text-to-image and image-to-image tasks
+- **Multimodality Comprehension** for multimodality-to-text and multimodality-to-audio tasks
+- **TTS** and **TTS Voice Clone** for TTS tasks
+
+This extension also offers example workflows (at **ComfyUI sidebar -> Templates -> vLLM-Omni**)
+
+!!! info
+    The node UI and feature designs are intended to match vLLM-Omni online serving interfaces. It cannot offer more than what the interfaces support.
+
+To build a simple workflow yourself,
+
+- Drag a generation node onto the canvas.
+- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
+- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
+- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
+    - For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
+    - For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
+    - For multi-stage models where all stages are either autoregression or diffusion, you can also connect only a single Sampling Params node, indicating that this set of sampling parameters will be used for all stages.
+
+## Examples & Screenshots
+
+Please read the [ComfyUI integration's README](https://github.com/vllm-project/vllm-omni/tree/main/apps/ComfyUI-vLLM-Omni) for more info.

docs/getting_started/installation/gpu.md

@@ -26,6 +26,8 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
 ### Pre-built wheels
 
+Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
+
 === "NVIDIA CUDA"
 
     --8<-- "docs/getting_started/installation/gpu/cuda.inc.md:pre-built-wheels"

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

@@ -17,12 +17,10 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
 # --8<-- [start:pre-built-wheels]
 
 #### Installation of vLLM
-Note: Pre-built wheels are currently only available for vLLM-Omni 0.11.0rc1 and 0.12.0rc1. For the latest version, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
-
 
 vLLM-Omni is built based on vLLM. Please install it with command below.
 ```bash
-uv pip install vllm==0.12.0 --torch-backend=auto
+uv pip install vllm==0.15.0 --torch-backend=auto
 ```
 
 #### Installation of vLLM-Omni
@@ -36,13 +34,13 @@ uv pip install vllm-omni
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.12.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.15.0 release version of the library
 
 ```bash
-uv pip install vllm==0.12.0 --torch-backend=auto
+uv pip install vllm==0.15.0 --torch-backend=auto
 ```
 
-The release 0.12.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
+The release 0.14.0 of vLLM is based on PyTorch 2.9.0 which requires CUDA 12.9 environment.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,11 +56,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.12.0
+git checkout v0.15.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.12.0/vllm-0.12.0-cp38-abi3-manylinux_2_31_x86_64.whl
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.15.0/vllm-0.15.0-cp38-abi3-manylinux_2_31_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash
@@ -93,7 +91,7 @@ docker run --runtime nvidia --gpus 2 \
     --env "HF_TOKEN=$HF_TOKEN" \
     -p 8091:8091 \
     --ipc=host \
-    vllm/vllm-omni:v0.12.0rc1 \
+    vllm/vllm-omni:v0.14.0 \
     --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
 ```
 

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

@@ -9,10 +9,60 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 
 # --8<-- [start:pre-built-wheels]
 
+#### Installation of vLLM
+
+vLLM-Omni is built based on vLLM. Please install it with command below.
+```bash
+uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+```
+
+#### Installation of vLLM-Omni
+
+```bash
+# we need to add --no-build-isolation as the torch
+# is not obtained from pypi, we have to install using the
+# torch installed in our environment
+uv pip install vllm-omni
+
+# Optional if want to run Qwen3 TTS
+uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+uv pip install onnxruntime-rocm sox
+```
+
 # --8<-- [end:pre-built-wheels]
 
 # --8<-- [start:build-wheel-from-source]
 
+#### Installation of vLLM
+If you do not need to modify source code of vLLM, you can directly install the stable 0.14.0 release version of the library
+
+```bash
+uv pip install vllm==0.14.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.14.0/rocm700
+```
+
+The release 0.14.0 of vLLM requires ROCm 7.0 environment.
+
+#### Installation of vLLM-Omni
+Since vllm-omni is rapidly evolving, it's recommended to install it from source
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=rocm uv pip install -e .
+# OR
+uv pip install -e . --no-build-isolation
+```
+
+<details><summary>(Optional) Installation of vLLM from source</summary>
+If you want to check, modify or debug with source code of vLLM, install the library from source with the following instructions:
+
+```bash
+git clone https://github.com/vllm-project/vllm.git
+cd vllm
+git checkout v0.14.0
+python3 -m pip install -r requirements/rocm.txt
+python3 setup.py develop
+```
+
 # --8<-- [end:build-wheel-from-source]
 
 # --8<-- [start:build-docker]
@@ -23,19 +73,30 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
 DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.rocm -t vllm-omni-rocm .
 ```
 
-If you want to specify which GPU Arch to build for to cutdown build time:
-
-```bash
-DOCKER_BUILDKIT=1 docker build \
-  -f docker/Dockerfile.rocm \
-  --build-arg PYTORCH_ROCM_ARCH="gfx942;gfx950" \
-  -t vllm-omni-rocm .
-```
-
 #### Launch the docker image
 
+##### Launch with OpenAI API Server
+
 ```
-docker run -it \
+docker run --rm \
+--group-add=video \
+--ipc=host \
+--cap-add=SYS_PTRACE \
+--security-opt seccomp=unconfined \
+--device /dev/kfd \
+--device /dev/dri \
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--env "HF_TOKEN=$HF_TOKEN" \
+-p 8091:8091 \
+--ipc=host \
+vllm-omni-rocm \
+--model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
+```
+
+##### Launch with interactive session for development
+
+```
+docker run --rm -it \
 --network=host \
 --group-add=video \
 --ipc=host \
@@ -44,8 +105,9 @@ docker run -it \
 --device /dev/kfd \
 --device /dev/dri \
 -v <path/to/model>:/app/model \
-vllm-omni-rocm \
-bash
+-v ~/.cache/huggingface:/root/.cache/huggingface \
+--entrypoint bash \
+vllm-omni-rocm
 ```
 
 # --8<-- [end:build-docker]
@@ -57,8 +119,7 @@ vLLM-Omni offers an official docker image for deployment. These images are built
 #### Launch vLLM-Omni Server
 Here's an example deployment command that has been verified on 2 x MI300's:
 ```bash
-docker run -it \
-  --network=host \
+docker run --rm \
   --group-add=video \
   --ipc=host \
   --cap-add=SYS_PTRACE \
@@ -68,14 +129,15 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  vllm/vllm-omni-rocm:v0.12.0rc1 \
-  vllm serve --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
+  -p 8091:8091 \
+  vllm/vllm-omni-rocm:v0.14.0 \
+  --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
 #### Launch an interactive terminal with prebuilt docker image.
 If you want to run in dev environment you can launch the docker image as follows:
 ```bash
-docker run -it \
+docker run --rm -it \
   --network=host \
   --group-add=video \
   --ipc=host \
@@ -86,8 +148,8 @@ docker run -it \
   -v <path/to/model>:/app/model \
   -v ~/.cache/huggingface:/root/.cache/huggingface \
   --env "HF_TOKEN=$HF_TOKEN" \
-  vllm/vllm-omni-rocm:v0.12.0rc1 \
-  bash
+  --entrypoint bash \
+  vllm/vllm-omni-rocm:v0.14.0
 ```
 
 # --8<-- [end:pre-built-images]

docs/getting_started/installation/npu.md

@@ -16,7 +16,13 @@ vLLM-Omni supports NPU through the vLLM Ascend Plugin (vllm-ascend). This is a c
 
 ## Installation
 
-### Recommended
+### Set up using Docker
+
+=== "NPU"
+
+    --8<-- "docs/getting_started/installation/npu/npu.inc.md:pre-built-images"
+
+### Build wheel from source
 
 === "NPU"
 

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

@@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0
 export DEVICE1=/dev/davinci1
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1-a3
+export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0rc1
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -34,26 +34,79 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
-# Install the missing dependency of mooncake in the origin image.
-apt update
-apt install libjemalloc2
-echo "export LD_PRELOAD=/usr/lib/$(uname -m)-linux-gnu/libjemalloc.so.2:$LD_PRELOAD" >> ~/.bashrc
-source ~/.bashrc
-
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
-git clone -b v0.12.0rc1 https://github.com/vllm-project/vllm-omni.git
-cd vllm-omni
-pip install -v -e .
-export VLLM_WORKER_MULTIPROC_METHOD=spawn
+git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git
 
-# (Optional) Disable mooncake for stable capability
-mv /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake \
-   /usr/local/Ascend/ascend-toolkit/latest/python/site-packages/mooncake.disabled
+cd vllm-omni
+VLLM_OMNI_TARGET_DEVICE=npu pip install -v -e .
+# OR pip install -v -e . --no-build-isolation
+export VLLM_WORKER_MULTIPROC_METHOD=spawn
 ```
 
 The default workdir is `/workspace`, with vLLM, vLLM-Ascend and vLLM-Omni code placed in `/vllm-workspace` installed in development mode.
 
 For other installation methods (pip installation, building from source, custom Docker builds), please refer to the [vllm-ascend installation guide](https://docs.vllm.ai/projects/ascend/en/latest/installation.html).
 
+We are keeping [issue #997](https://github.com/vllm-project/vllm-omni/issues/997) up to date with the aligned versions of vLLM, vLLM-Ascend, and vLLM-Omni, and also outlining the Q1 roadmap there.
+
 # --8<-- [end:installation]
+
+# --8<-- [start:pre-built-images]
+
+`vllm-ascend` offers Docker images for deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Supported images as following.
+
+| image name | Hardware | OS |
+|-|-|-|
+| image-tag | Atlas A2 | Ubuntu |
+| image-tag-a3 | Atlas A3 | Ubuntu |
+
+vLLM-Omni offers Docker images for Ascend NPU deployment. You can just pull the **prebuilt image** from the image repository [ascend/vllm-ascend](https://quay.io/repository/ascend/vllm-ascend?tab=tags) and run it with bash.
+
+Here's an example deployment command that has been verified on 6 x NPUs:
+
+```bash
+# Atlas A2:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+# Atlas A3:
+# export IMAGE=quay.io/ascend/vllm-omni:v0.14.0-a3
+export IMAGE=quay.io/ascend/vllm-omni:v0.14.0
+docker run --rm \
+    --name vllm-omni-npu \
+    --shm-size=1g \
+    --device /dev/davinci0 \
+    --device /dev/davinci1 \
+    --device /dev/davinci2 \
+    --device /dev/davinci3 \
+    --device /dev/davinci4 \
+    --device /dev/davinci5 \
+    --device /dev/davinci_manager \
+    --device /dev/devmm_svm \
+    --device /dev/hisi_hdc \
+    -v /usr/local/dcmi:/usr/local/dcmi \
+    -v /usr/local/bin/npu-smi:/usr/local/bin/npu-smi \
+    -v /usr/local/Ascend/driver/lib64/:/usr/local/Ascend/driver/lib64/ \
+    -v /usr/local/Ascend/driver/version.info:/usr/local/Ascend/driver/version.info \
+    -v /etc/ascend_install.info:/etc/ascend_install.info \
+    -v /root/.cache:/root/.cache \
+    -p 8000:8000 \
+    -it $IMAGE bash
+```
+
+!!! tip
+    You can use this docker image to serve models the same way you would with in vLLM! To do so, make sure you overwrite the default entrypoint (`vllm serve --omni`) which works only for models supported in the vLLM-Omni project.
+
+Or build IMAGE from **source code**:
+
+```bash
+git clone https://github.com/vllm-project/vllm-omni.git
+cd vllm-omni
+# A2
+docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu .
+# A3
+# docker build -t vllm-omni-dev-image:latest -f ./docker/Dockerfile.npu.a3 .
+```
+
+# --8<-- [end:pre-built-images]

docs/getting_started/quickstart.md

@@ -17,7 +17,13 @@ For installation on GPU from source:
 ```bash
 uv venv --python 3.12 --seed
 source .venv/bin/activate
-uv pip install vllm==0.12.0 --torch-backend=auto
+
+# On CUDA
+uv pip install vllm==0.15.0 --torch-backend=auto
+
+# On ROCm
+uv pip install vllm==0.15.0 --extra-index-url https://wheels.vllm.ai/rocm/0.15.0/rocm700
+
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni
 uv pip install -e .
@@ -40,6 +46,46 @@ if __name__ == "__main__":
     images[0].save("coffee.png")
 ```
 
+You can pass a list of prompts and wait for them to process altogether, shown below.
+
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(
+        model="Tongyi-MAI/Z-Image-Turbo",
+        # stage_configs_path="./stage-config.yaml",  # See below
+    )
+    prompts = [
+        "a cup of coffee on a table",
+        "a toy dinosaur on a sandy beach",
+        "a fox waking up in bed and yawning",
+    ]
+    omni_outputs = omni.generate(prompts)
+    for i_prompt, prompt_output in enumerate(omni_outputs):
+        this_request_output = prompt_output.request_output[0]
+        this_images = this_request_output.images
+        for i_image, image in enumerate(this_images):
+            image.save(f"p{i_prompt}-img{i_image}.jpg")
+            print("saved to", f"p{i_prompt}-img{i_image}.jpg")
+            # saved to p0-img0.jpg
+            # saved to p1-img0.jpg
+            # saved to p2-img0.jpg
+```
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
 For more usages, please refer to [offline inference](../user_guide/examples/offline_inference/qwen2_5_omni.md)
 
 ## Online Serving with OpenAI-Completions API

docs/mkdocs/hooks/generate_api_readme.py

@@ -34,10 +34,10 @@ CATEGORIES = {
         "name": "Core",
         "description": "Core scheduling and caching components.",
     },
-    "model_executor": {
-        "name": "Model Executor",
-        "description": "Model execution components.",
-    },
+    # "model_executor": {
+    #     "name": "Model Executor",
+    #     "description": "Model execution components.",
+    # },
     "config": {
         "name": "Configuration",
         "description": "Configuration classes.",

docs/models/supported_models.md

@@ -25,6 +25,7 @@ th {
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
+| `GlmImagePipeline` | GLM-Image | `zai-org/GLM-Image` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
 | `WanPipeline` | Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` |
@@ -33,7 +34,11 @@ th {
 |`LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` |
 |`StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` |
 |`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` |
 |`StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |
 
 
 ## List of Supported Models for NPU
@@ -47,6 +52,7 @@ th {
 
 | Architecture | Models | Example HF Models |
 |--------------|--------|-------------------|
+| `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` |
 | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B`|
 | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` |
 | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` |
@@ -55,3 +61,8 @@ th {
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2511 | `Qwen/Qwen-Image-Edit-2511` |
 |`ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` |
+|`LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` |
+|`Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` |
+|`Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` |

docs/serving/image_edit_api.md

@@ -0,0 +1,205 @@
+# Image Edit API
+
+vLLM-Omni provides an OpenAI DALL-E compatible API for image edit using diffusion models.
+
+Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
+
+## Quick Start
+
+### Start the Server
+
+For example...
+
+```bash
+# Qwen-Image
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni --port 8000
+
+
+### Generate Images
+
+**Using curl:**
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@./xx.png" \
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+**Using OpenAI SDK:**
+
+```python
+import base64
+from openai import OpenAI
+from pathlib import Path
+client = OpenAI(
+    api_key="None",
+    base_url="http://localhost:8000/v1"
+)
+
+input_image_url = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png"
+
+result = client.images.edit(
+    image=[],
+    model="Qwen-Image-Edit-2511",
+    prompt="Change the bears in the two input images into walking together.",
+    size='512x512',
+    stream=False,
+    output_format='jpeg',
+    # url格式
+    extra_body={
+        "url": [input_image_url1,input_image_url],
+        "num_inference_steps": 50,
+        "guidance_scale": 1,
+        "seed": 777,
+    }
+)
+
+image_base64 = result.data[0].b64_json
+image_bytes = base64.b64decode(image_base64)
+
+# Save the image to a file
+with open("edit_out_http.jpeg", "wb") as f:
+    f.write(image_bytes)
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/images/edits
+Content-Type: multipart/form-data
+```
+
+### Request Parameters
+
+#### OpenAI Standard Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `prompt` | string | **required** | A text description of the desired image |
+| `model` | string | server's model | Model to use (optional, should match server if specified) |
+| `image` | string or array | **required** | The image(s) to edit. |
+| `n` | integer | 1 | Number of images to generate (1-10) |
+| `size` | string | "auto" | Image dimensions in WxH format (e.g., "1024x1024", "512x512"), when set to auto, it decide size from first input image. |
+| `response_format` | string | "b64_json" | Response format (only "b64_json" supported) |
+| `user` | string | null | User identifier for tracking |
+| `output_format` | string | "png" | The format in which the generated images are returned. Must be one of "png", "jpg", "jpeg", "webp". |
+| `output_compression` | integer | 100 | The compression level (0-100%) for the generated images. |
+| `background` | string or null | "auto" | Allows to set transparency for the background of the generated image(s).
+
+#### vllm-omni Extension Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `url` | string or array | None | The image(s) to edit. |
+| `negative_prompt` | string | null | Text describing what to avoid in the image |
+| `num_inference_steps` | integer | model defaults | Number of diffusion steps |
+| `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) |
+| `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) |
+| `seed` | integer | null | Random seed for reproducibility |
+
+### Response Format
+
+```json
+{
+  "created": 1701234567,
+  "data": [
+    {
+      "b64_json": "<base64-encoded PNG>",
+      "url": null,
+      "revised_prompt": null
+    }
+  ],
+  "output_format": null,
+  "size": null,
+}
+```
+
+## Examples
+
+### Multiple Images input
+
+```bash
+curl -s -D >(grep -i x-request-id >&2) \
+  -o >(jq -r '.data[0].b64_json' | base64 --decode > gift-basket.png) \
+  -X POST "http://localhost:8000/v1/images/edits" \
+  -F "model=xxx" \
+  -F "image=@xx.png" \
+  -F "image=@xx.png"
+  -F "prompt='this bear is wearing sportwear. holding a basketball, and bending one leg.'" \
+  -F "size=1024x1024" \
+  -F "output_format=png"
+```
+
+
+## Parameter Handling
+
+The API passes parameters directly to the diffusion pipeline without model-specific transformation:
+
+- **Default values**: When parameters are not specified, the underlying model uses its own defaults
+- **Pass-through design**: User-provided values are forwarded directly to the diffusion engine
+- **Minimal validation**: Only basic type checking and range validation at the API level
+
+### Parameter Compatibility
+
+The API passes parameters directly to the diffusion pipeline without model-specific validation.
+
+- Unsupported parameters may be silently ignored by the model
+- Incompatible values will result in errors from the underlying pipeline
+- Recommended values vary by model - consult model documentation
+
+**Best Practice:** Start with the model's recommended parameters, then adjust based on your needs.
+
+## Error Responses
+
+### 400 Bad Request
+
+Invalid parameters (e.g., model mismatch):
+
+```json
+{
+  "detail": "Invalid size format: '1024x'. Expected format: 'WIDTHxHEIGHT' (e.g., '1024x1024')."
+}
+```
+
+### 422 Unprocessable Entity
+
+Validation errors (missing required fields):
+
+```json
+{
+  "detail": "Field 'image' or 'url' is required"
+}
+```
+
+## Troubleshooting
+
+### Server Not Running
+
+```bash
+# Check if server is responding
+curl -X http://localhost:8000/v1/images/edit \
+  -F "prompt='test'"
+```
+
+### Out of Memory
+
+If you encounter OOM errors:
+1. Reduce image size: `"size": "512x512"`
+2. Reduce inference steps: `"num_inference_steps": 25`
+
+## Development
+
+Enable debug logging to see prompts and generation details:
+
+```bash
+vllm serve Qwen/Qwen-Image-Edit-2511 --omni \
+  --uvicorn-log-level debug
+```

docs/usage/faq.md

@@ -22,7 +22,7 @@ A: At first, you can check current [issues](https://github.com/vllm-project/vllm
 
 > Q: Does vLLM-Omni support AWQ or any other quantization?
 
-A: vLLM-Omni partitions model into several stages. For AR stages, it will reuse main logic of LLMEngine in vLLM. So current quantization supported in vLLM should be also supported in vLLM-Omni for them. But systematic verification is ongoing. For quantization for DiffusionEngine, we are working on it. Please stay tuned and welcome contribution!
+A: We plan to introduce GGUF FP8 prequantized models and online FP8 quantization in version 0.16.0. Support for other quantization types will follow in future releases. For details, please see our [Q1 quantization roadmap](https://github.com/vllm-project/vllm-omni/issues/1057).
 
 > Q: Does vLLM-Omni support multimodal streaming input and output?
 

docs/user_guide/diffusion/cache_dit_acceleration.md

@@ -18,6 +18,7 @@ Enable cache-dit acceleration by simply setting `cache_backend="cache_dit"`. Cac
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 # Simplest way: just enable cache-dit with default parameters
 omni = Omni(
@@ -27,7 +28,7 @@ omni = Omni(
 
 images = omni.generate(
     "a beautiful landscape",
-    num_inference_steps=50,
+    OmniDiffusionSamplingParams(num_inference_steps=50),
 )
 ```
 
@@ -69,7 +70,7 @@ vllm serve Qwen/Qwen-Image --omni --port 8091 \
 
 ## Acceleration Methods
 
-For comprehensive illustration, please view cache-dit [User_Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
+For comprehensive illustration, please view cache-dit [User_Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
 
 ### 1. DBCache (Dual Block Cache)
 
@@ -222,6 +223,6 @@ To test another model, you can modify `--model` with the target model identifier
 
 ## Additional Resources
 
-- [Cache-DiT User Guide](https://github.com/vipshop/cache-dit/blob/main/docs/User_Guide.md)
-- [Cache-DiT Benchmark](https://github.com/vipshop/cache-dit/tree/main/bench)
-- [DBCache Technical Details](https://github.com/vipshop/cache-dit/blob/main/docs/DBCache.md)
+- [Cache-DiT User Guide](https://cache-dit.readthedocs.io/en/latest/user_guide/OVERVIEWS/)
+- [Cache-DiT Benchmark](https://cache-dit.readthedocs.io/en/latest/benchmark/HYBRID_CACHE/)
+- [DBCache Technical Details](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/)

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -1,19 +1,32 @@
 # CPU Offloading for Diffusion Model
 
 ## Overview
+
+vLLM-Omni provides two offloading strategies to reduce GPU memory usage for diffusion models, allowing you to run larger models on GPUs with limited VRAM:
+
+1. **Model-level (Component) Offloading**: Swaps entire model components (DiT transformer, VAE, encoders) between GPU and CPU.
+2. **Layerwise (Blockwise) Offloading**: Keeps only a single or a few transformer blocks on GPU at a time, with compute - memory copy overlap.
+
+Both approaches use pinned memory for faster CPU-GPU transfers. For now, the two offloading strategies could not be used at the same time.
+
+
+## Model-level CPU Offloading
+
+### Implementation
+
 CPU offload lets the diffusion worker move large model components between GPU and CPU memory on demand. It keeps the DiT transformer resident on GPU only while it is actively running, and swaps it out when encoders modules need the device. This reduces peak VRAM usage so bigger checkpoints run on smaller GPUs, or multiple requests can share the same GPU.
 
-## Execution Model
+**Execution Flow**:
 1. Text encoders run on GPU while the DiT transformer is offloaded to CPU.
 2. Before denoising, weights are prefetched back to GPU, honoring pinned-memory copies for speed.
 3. After the diffusion step, the transformer returns to CPU and the process repeats as needed.
 
 Transfers use pinned host buffers, and the worker coordinates swaps via mutex-style hooks so components never compete for memory.
 
-## Configuration
+### Configuration
 You can enable CPU offload in two ways:
 
-- **Python API**: set `enable_cpu_offload=True`.
+1. **Python API**: set `enable_cpu_offload=True`.
 
 ```python
 from vllm_omni import Omni
@@ -23,7 +36,66 @@ if __name__ == "__main__":
     m = Omni(model="Qwen/Qwen-Image",enable_cpu_offload=True)
 ```
 
-- **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
+2. **CLI**: pass `--enable-cpu-offload` to the diffusion service entrypoint.
 
-## Known Limitations
+### Limitations
 - Cold start latency increases for over one minute for some models(e.g., Qwen-Image)
+
+
+## Layerwise (Blockwise) Offloading
+
+### Implementation
+Layerwise offload operates at transformer block granularity, keeping a single transformer block, or a specified number of blocks, on GPU while others stay in CPU memory.
+
+Unlike full model-wise CPU offload which swaps entire components like DiT and encoders, layerwise offloading applies a sliding window way of loading and offloading weights between gpu and cpu: while block `i` computes, block `i+1` get prefetched asynchronously via pinned memory. In this way, only partial blocks(s) reside on GPU at any moment during inference, so that greatly decrease the memory occupancy.
+
+**Execution Flow**:
+
+1. During model initialization, all components are loaded to CPU first. Then components other than DiT model(s) in the pipeline, such as VAE and encoders, are moved to GPU. The weights of target transformer blocks are collected as contiguous tensors per layer on CPU with pinned memory; and non-block modules (embeddings, norms, etc) in the DiT model are moved to and stay on GPU.
+2. The first block(s) are transferred to GPU during initialization of `LayerwiseOffloader`, before the first denoising step of the very first request.
+3. As each block executes, the next block prefetches on a separate CUDA stream for compute - memory copy overlap. After execution, the current block is immediately freed from GPU memory.
+4. When the last block completes, the first block prefetches for the next denoising step.
+
+
+Example of hook executions of a DiT model with n layers, by default keep a single layer on GPU:
+| Layer (block) idx | forward pre-hook               | forward          | forward post-hook         |
+|-------------------|--------------------------------|------------------|---------------------------|
+| layer-0           | prefetch layer 1 (copy stream) | compute layer 0  | free layer-0 gpu weights  |
+| layer-1           | prefetch layer 2 (copy stream) | compute layer 1  | free layer-1 gpu weights  |
+| layer-2           | prefetch layer 3 (copy stream) | compute layer 2  | free layer-2 gpu weights  |
+| ...               | ...                            | ...              | ...                       |
+| layer-(n-1)       | **prefetch layer 0 (copy stream)** | compute layer (n-1) | free layer (n-1) gpu weights  |
+
+
+### Configuration
+
+1. **Python API**: set `enable_layerwise_offload=True` and optionally `layerwise_num_gpu_layers`.
+
+```python
+from vllm_omni import Omni
+
+if __name__ == "__main__":
+    m = Omni(
+        model="Wan-AI/Wan2.2-T2V-A14B-Diffusers",
+        enable_layerwise_offload=True,
+        ...
+    )
+```
+
+2. **CLI**: pass `--enable-layerwise-offload` and `--layerwise-num-gpu-layers` to the diffusion service entrypoint.
+
+### Supported Models
+
+| Architecture | Models | Example HF Models | DiT Model Cls | Blocks Attr Name |
+|--------------|--------|-------------------|----------|----------|
+| `QwenImagePipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | "transformer_blocks" |
+| `Wan22Pipeline` | Wan2.2 | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | "blocks" |
+
+NOTE: Models must define `_layerwise_offload_blocks_attr` class attribute so that the layerwise offloader finds the target transformer blocks.
+
+### Limitations
+- Cold start latency increases because of
+    1) components are loaded to CPU first at the very first during initialization,  
+    2) weight consolidation and pinning
+- Performance depends on CPU <-> GPU interconnection (e.g., PCIe bandwidth).
+- Support single GPU only for now

docs/user_guide/diffusion/parallelism_acceleration.md

@@ -18,17 +18,28 @@ The following table shows which models are currently supported by parallelism me
 
 ### ImageGen
 
-| Model | Model Identifier | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
-|-------|------------------|------------|---------|--------------|--------------------------|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ✅ | ✅ | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ✅ | ✅ | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ❌ |
-| **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ✅ | ✅ | ✅ | ❌ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ❌ | ❌ | ✅ (TP=2 only) |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ❌ | ❌ | ❌ |
+| Model                    | Model Identifier                     | Ulysses-SP | Ring-SP | CFG-Parallel | Tensor-Parallel |
+|--------------------------|--------------------------------------|:----------:|:-------:|:------------:|:---------------:|
+| **LongCat-Image**        | `meituan-longcat/LongCat-Image`      |     ✅      |    ✅    |      ❌       |        ✅        |
+| **LongCat-Image-Edit**   | `meituan-longcat/LongCat-Image-Edit` |     ✅      |    ✅    |      ❌       |        ✅        |
+| **Ovis-Image**           | `OvisAI/Ovis-Image`                  |     ❌      |    ❌    |      ❌       |        ❌        |
+| **Qwen-Image**           | `Qwen/Qwen-Image`                    |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Edit**      | `Qwen/Qwen-Image-Edit`               |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509`          |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Qwen-Image-Layered**   | `Qwen/Qwen-Image-Layered`            |     ✅      |    ✅    |      ✅       |        ✅        |
+| **Z-Image**              | `Tongyi-MAI/Z-Image-Turbo`           |     ✅      |    ✅    |      ❌       |  ✅ (TP=2 only)  |
+| **Stable-Diffusion3.5**  | `stabilityai/stable-diffusion-3.5`   |     ❌      |    ❌    |      ❌       |        ❌        |
+| **FLUX.2-klein**         | `black-forest-labs/FLUX.2-klein-4B`  |     ❌      |    ❌    |      ❌       |        ✅        |
+| **FLUX.1-dev**           | `black-forest-labs/FLUX.1-dev`       |     ❌      |    ❌    |      ❌       |        ✅        |
+
+!!! 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.
+
+    - Good news: The text_encoder typically has minimal impact on overall inference performance.
+    - Bad news: When TP is enabled, every TP process retains a full copy of the text_encoder weights, leading to significant GPU memory waste.
+
+    We are actively refactoring this design to address this. For details and progress, please refer to [Issue #771](https://github.com/vllm-project/vllm-omni/issues/771).
+
 
 !!! note "Why Z-Image is TP=2 only"
     Z-Image Turbo is currently limited to `tensor_parallel_size` of **1 or 2** due to model shape divisibility constraints.
@@ -38,7 +49,7 @@ The following table shows which models are currently supported by parallelism me
 
 | Model | Model Identifier | Ulysses-SP | Ring-SP | Tensor-Parallel |
 |-------|------------------|------------|---------|--------------------------|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ❌ | ❌ |
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ✅ | ✅ | ✅ |
 
 ### Tensor Parallelism
 
@@ -56,10 +67,12 @@ omni = Omni(
 )
 
 outputs = omni.generate(
-    prompt="a cat reading a book",
-    num_inference_steps=9,
-    width=512,
-    height=512,
+    "a cat reading a book",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=9,
+        width=512,
+        height=512,
+    ),
 )
 ```
 
@@ -72,6 +85,7 @@ outputs = omni.generate(
 An example of offline inference script using [Ulysses-SP](https://arxiv.org/pdf/2309.14509) is shown below:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2
 
@@ -80,7 +94,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -122,6 +139,7 @@ Ring-Attention ([arxiv paper](https://arxiv.org/abs/2310.01889)) splits the inpu
 An example of offline inference script using Ring-Attention is shown below:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ring_degree = 2
 
@@ -130,7 +148,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ring_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 See `examples/offline_inference/text_to_image/text_to_image.py` for a complete working example.
@@ -172,6 +193,7 @@ You can combine both Ulysses-SP and Ring-Attention for larger scale parallelism.
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 
 # Hybrid: 2 Ulysses × 2 Ring = 4 GPUs total
@@ -180,7 +202,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 ##### Online Serving
@@ -210,107 +235,146 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 ##### How to parallelize a new model
 
-If a diffusion model has been deployed in vLLM-Omni and supports single-card inference, you can refer to the following instructions to parallelize it with [Ulysses-SP](https://arxiv.org/pdf/2309.14509).
+NOTE: "Terminology: SP vs CP"
+    Our "Sequence Parallelism" (SP) corresponds to "Context Parallelism" (CP) in the [diffusers library](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py).
+    We use "Sequence Parallelism" to align with vLLM-Omni's terminology.
 
+---
 
-This section uses **Qwen-Image** (`QwenImageTransformer2DModel`) as the reference implementation. Qwen-Image is a **dual-stream** transformer (text + image) that performs **joint attention** across the concatenated sequences. Because of that, when enabling sequence parallel you typically:
+###### Non-intrusive `_sp_plan` (Recommended)
 
-- Chunk **image tokens** (`hidden_states`) across SP ranks along the **sequence dimension**.
-- Keep **text embeddings** (`encoder_hidden_states`) **replicated** on all SP ranks for correctness (unless you implement full joint-SP semantics and explicitly split text too).
-- Chunk **image RoPE freqs** to match the chunked image tokens.
-- `all_gather` the final output back along the sequence dimension.
+The `_sp_plan` mechanism allows SP without modifying `forward()` logic. The framework automatically registers hooks to shard inputs and gather outputs at module boundaries.
 
-First, add the sequence-parallel helpers and (for Qwen-Image) the forward-context flag. Then, in the transformer's `forward()`:
+**Requirements for `forward()` function:**
 
-- Chunk `hidden_states` by SP world size.
-- Set `get_forward_context().split_text_embed_in_sp = False` because Qwen-Image uses joint attention and we would use a full text embedding in each rank.
-- After `pos_embed`, chunk `img_freqs` on `dim=0` (token axis) to match the chunked image tokens; only chunk `txt_freqs` if you explicitly split text embeddings in SP.
+- Tensor operations that need sharding/gathering must happen at **`nn.Module` boundaries** (not inline Python operations)
+- If your `forward()` contains inline tensor operations (e.g., `torch.cat`, `pad_sequence`) that need sharding, **extract them into a submodule**
 
-Taking `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` as an example:
+**When to create a submodule:**
 
-```diff
-from vllm_omni.diffusion.distributed.parallel_state import (
-    get_sequence_parallel_rank,
-    get_sequence_parallel_world_size,
-    get_sp_group,
+```python
+# ❌ BAD: Inline operations - hooks cannot intercept
+def forward(self, x, cap_feats):
+    unified = torch.cat([x, cap_feats], dim=1)  # Cannot be sharded via _sp_plan
+    ...
+
+# ✅ GOOD: Extract into a submodule
+class UnifiedPrepare(nn.Module):
+    def forward(self, x, cap_feats):
+        return torch.cat([x, cap_feats], dim=1)  # Now can be sharded via _sp_plan
+
+class MyModel(nn.Module):
+    def __init__(self):
+        self.unified_prepare = UnifiedPrepare()  # Submodule
+
+    def forward(self, x, cap_feats):
+        unified = self.unified_prepare(x, cap_feats)  # Hook can intercept here
+```
+
+---
+
+###### Defining `_sp_plan`
+
+**Type definitions** (see [diffusers `_modeling_parallel.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/_modeling_parallel.py) for reference):
+
+```python
+from vllm_omni.diffusion.distributed.sp_plan import (
+    SequenceParallelInput,   # Corresponds to diffusers' ContextParallelInput
+    SequenceParallelOutput,  # Corresponds to diffusers' ContextParallelOutput
 )
-from vllm_omni.diffusion.forward_context import get_forward_context
-
-class QwenImageTransformer2DModel(...):
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor = None,
-        encoder_hidden_states_mask: torch.Tensor = None,
-        timestep: torch.LongTensor = None,
-        img_shapes: list[tuple[int, int, int]] | None = None,
-        txt_seq_lens: list[int] | None = None,
-        ...
-    ):
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           # Chunk image tokens along sequence dimension.
-+           hidden_states = torch.chunk(
-+               hidden_states,
-+               get_sequence_parallel_world_size(),
-+               dim=-2,
-+           )[get_sequence_parallel_rank()]
-+
-+           # Qwen-Image uses *dual-stream* (text + image) and runs *joint attention*.
-+           # Text embeddings should be replicated across SP ranks for correctness.
-+           get_forward_context().split_text_embed_in_sp = False
-
-        hidden_states = self.img_in(hidden_states)
-
-        ...
-        image_rotary_emb = self.pos_embed(img_shapes, txt_seq_lens, device=hidden_states.device)
-
-+       def get_rotary_emb_chunk(freqs):
-+           return torch.chunk(freqs, get_sequence_parallel_world_size(), dim=0)[
-+               get_sequence_parallel_rank()
-+           ]
-+
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           img_freqs, txt_freqs = image_rotary_emb
-+           img_freqs = get_rotary_emb_chunk(img_freqs)
-+           if get_forward_context().split_text_embed_in_sp:
-+               txt_freqs = get_rotary_emb_chunk(txt_freqs)
-+           image_rotary_emb = (img_freqs, txt_freqs)
 ```
 
-Next, at the end of the `forward()` function, call `get_sp_group().all_gather` to gather the chunked outputs across devices and concatenate them along the sequence dimension (matching the earlier `dim=-2` chunking):
+| Parameter | Description |
+|-----------|-------------|
+| `split_dim` | Dimension to split/gather (usually `1` for sequence) |
+| `expected_dims` | Expected tensor rank for validation (optional) |
+| `split_output` | `False`: shard **input** parameters; `True`: shard **output** tensors |
+| `auto_pad` | Auto-pad if sequence not divisible by world_size (Ulysses only) |
 
-```diff
-class QwenImageTransformer2DModel(...):
-    def forward(...):
-        ...
-        hidden_states = self.norm_out(hidden_states, temb)
-        output = self.proj_out(hidden_states)
+**Key naming convention:**
 
-+       if self.parallel_config.sequence_parallel_size > 1:
-+           output = get_sp_group().all_gather(output, dim=-2)
-        return Transformer2DModelOutput(sample=output)
+| Key | Meaning | Python equivalent |
+|-----|---------|-------------------|
+| `""` | Root model | `model` |
+| `"blocks.0"` | First element of ModuleList | `model.blocks[0]` |
+| `"blocks.*"` | All elements of ModuleList | `for b in model.blocks` |
+| `"outputs.main"` | ModuleDict entry | `model.outputs["main"]` |
+
+**Dictionary key types:**
+
+| Key type | `split_output` | Description |
+|----------|----------------|-------------|
+| `"param_name"` (str) | `False` | Shard **input parameter** by name |
+| `0`, `1` (int) | `True` | Shard **output tuple** by index |
+
+**Example** (similar to [diffusers `transformer_wan.py`](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/transformers/transformer_wan.py)):
+
+```python
+class MyTransformer(nn.Module):
+    _sp_plan = {
+        # Shard rope module OUTPUTS (returns tuple)
+        "rope": {
+            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # cos
+            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),  # sin
+        },
+        # Shard transformer block INPUT parameter
+        "blocks.0": {
+            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
+        },
+        # Gather at final projection
+        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
+    }
 ```
 
-Finally, you can set the parallel configuration and pass it to `Omni` and start parallel inference with:
-```diff
-from vllm_omni import Omni
-+from vllm_omni.diffusion.data import DiffusionParallelConfig
-ulysses_degree = 2
+---
 
-omni = Omni(
-    model="Qwen/Qwen-Image",
-+    parallel_config=DiffusionParallelConfig(ulysses_degree=2)
-)
+###### Hook flow
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
 ```
+Input → [SequenceParallelSplitHook: pre_forward] → Module.forward() → [post_forward] → ...
+                                                                              ↓
+... → [SequenceParallelGatherHook: post_forward] → Output
+```
+
+1. **SplitHook** shards tensors before/after the target module
+2. **Attention layers** handle Ulysses/Ring communication internally
+3. **GatherHook** collects sharded outputs
+
+The framework automatically applies these hooks when `sequence_parallel_size > 1`.
+
+---
+
+###### Method 2: Intrusive modification (For complex cases)
+
+For models with dynamic sharding logic that cannot be expressed via `_sp_plan`:
+
+```python
+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:
+        hidden_states = sp_shard(hidden_states, dim=1)
+        # ... computation ...
+        output = sp_gather(output, dim=1)
+    return output
+```
+
+---
+
+###### Choosing the right approach
+
+| Scenario | Approach |
+|----------|----------|
+| Standard transformer | `_sp_plan` |
+| Inline tensor ops need sharding | Extract to submodule + `_sp_plan` |
+| Dynamic/conditional sharding | Intrusive modification |
 
 
 ### CFG-Parallel
 
-##### Offline Inference
+#### Offline Inference
 
-CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=...)`. The recommended configuration is `cfg_parallel_size=2` (one rank for the positive branch and one rank for the negative branch).
+CFG-Parallel is enabled through `DiffusionParallelConfig(cfg_parallel_size=2)`, which runs one rank for the positive branch and one rank for the negative branch.
 
 An example of offline inference using CFG-Parallel (image-to-image) is shown below:
 
@@ -318,23 +382,50 @@ An example of offline inference using CFG-Parallel (image-to-image) is shown bel
 from vllm_omni import Omni
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 
+image_path = "path_to_image.png"
 omni = Omni(
     model="Qwen/Qwen-Image-Edit",
     parallel_config=DiffusionParallelConfig(cfg_parallel_size=2),
 )
+input_image = Image.open(image_path).convert("RGB")
 
 outputs = omni.generate(
-    prompt="turn this cat to a dog",
-    negative_prompt="low quality, blurry",
-    true_cfg_scale=4.0,
-    pil_image=input_image,
-    num_inference_steps=50,
+    {
+        "prompt": "turn this cat to a dog",
+        "negative_prompt": "low quality, blurry",
+        "multi_modal_data": {"image": input_image},
+    },
+    OmniDiffusionSamplingParams(
+        true_cfg_scale=4.0,
+        num_inference_steps=50,
+    ),
 )
 ```
 
 Notes:
 
-- CFG-Parallel is only effective when **true CFG** is enabled (i.e., `true_cfg_scale > 1` and a `negative_prompt` is provided).
+- CFG-Parallel is only effective when a `negative_prompt` is provided AND a guidance scale (or `cfg_scale`) is greater than 1.
+
+See `examples/offline_inference/image_to_image/image_edit.py` for a complete working example.
+```bash
+cd examples/offline_inference/image_to_image/
+python image_edit.py \
+  --model "Qwen/Qwen-Image-Edit" \
+  --image "qwen_image_output.png" \
+  --prompt "turn this cat to a dog" \
+  --negative_prompt "low quality, blurry" \
+  --cfg_scale 4.0 \
+  --output "edited_image.png" \
+  --cfg_parallel_size 2
+```
+
+#### Online Serving
+
+You can enable CFG-Parallel in online serving for diffusion models via `--cfg-parallel-size`:
+
+```bash
+vllm serve Qwen/Qwen-Image-Edit --omni --port 8091 --cfg-parallel-size 2
+```
 
 #### How to parallelize a pipeline
 
@@ -347,58 +438,130 @@ In `QwenImagePipeline`, each diffusion step runs two denoiser forward passes seq
 
 CFG-Parallel assigns these two branches to different ranks in the **CFG group** and synchronizes the results.
 
-Below is an example of CFG-Parallel implementation:
+vLLM-omni provides `CFGParallelMixin` base class that encapsulates the CFG parallel logic. By inheriting from this mixin and calling its methods, pipelines can easily implement CFG parallel without writing repetitive code.
+
+**Key Methods in CFGParallelMixin:**
+- `predict_noise_maybe_with_cfg()`: Automatically handles CFG parallel noise prediction
+- `scheduler_step_maybe_with_cfg()`: Scheduler step with automatic CFG rank synchronization
+
+**Example Implementation:**
 
 ```python
-def diffuse(
+class QwenImageCFGParallelMixin(CFGParallelMixin):
+    """
+    Base Mixin class for Qwen Image pipelines providing shared CFG methods.
+    """
+
+    def diffuse(
         self,
-        ...
-        ):
-    # Enable CFG-parallel: rank0 computes positive, rank1 computes negative.
-    cfg_parallel_ready = do_true_cfg and get_classifier_free_guidance_world_size() > 1
+        prompt_embeds: torch.Tensor,
+        prompt_embeds_mask: torch.Tensor,
+        negative_prompt_embeds: torch.Tensor,
+        negative_prompt_embeds_mask: torch.Tensor,
+        latents: torch.Tensor,
+        img_shapes: torch.Tensor,
+        txt_seq_lens: torch.Tensor,
+        negative_txt_seq_lens: torch.Tensor,
+        timesteps: torch.Tensor,
+        do_true_cfg: bool,
+        guidance: torch.Tensor,
+        true_cfg_scale: float,
+        image_latents: torch.Tensor | None = None,
+        cfg_normalize: bool = True,
+        additional_transformer_kwargs: dict[str, Any] | None = None,
+    ) -> torch.Tensor:
+        self.transformer.do_true_cfg = do_true_cfg
 
-    self.transformer.do_true_cfg = do_true_cfg
+        for i, t in enumerate(timesteps):
+            timestep = t.expand(latents.shape[0]).to(device=latents.device, dtype=latents.dtype)
 
-    if cfg_parallel_ready:
-        cfg_group = get_cfg_group()
-        cfg_rank = get_classifier_free_guidance_rank()
+            # Prepare kwargs for positive (conditional) prediction
+            positive_kwargs = {
+                "hidden_states": latents,
+                "timestep": timestep / 1000,
+                "guidance": guidance,
+                "encoder_hidden_states_mask": prompt_embeds_mask,
+                "encoder_hidden_states": prompt_embeds,
+                "img_shapes": img_shapes,
+                "txt_seq_lens": txt_seq_lens,
+            }
 
-        if cfg_rank == 0:
-            local_pred = self.transformer(
-                hidden_states=latents,
-                timestep=timestep / 1000,
-                guidance=guidance,
-                encoder_hidden_states_mask=prompt_embeds_mask,
-                encoder_hidden_states=prompt_embeds,
-                img_shapes=img_shapes,
-                txt_seq_lens=txt_seq_lens,
-                attention_kwargs=self.attention_kwargs,
-                return_dict=False,
-            )[0]
-        else:
-            local_pred = self.transformer(
-                hidden_states=latents,
-                timestep=timestep / 1000,
-                guidance=guidance,
-                encoder_hidden_states_mask=negative_prompt_embeds_mask,
-                encoder_hidden_states=negative_prompt_embeds,
-                img_shapes=img_shapes,
-                txt_seq_lens=negative_txt_seq_lens,
-                attention_kwargs=self.attention_kwargs,
-                return_dict=False,
-            )[0]
+            # Prepare kwargs for negative (unconditional) prediction
+            if do_true_cfg:
+                negative_kwargs = {
+                    "hidden_states": latents,
+                    "timestep": timestep / 1000,
+                    "guidance": guidance,
+                    "encoder_hidden_states_mask": negative_prompt_embeds_mask,
+                    "encoder_hidden_states": negative_prompt_embeds,
+                    "img_shapes": img_shapes,
+                    "txt_seq_lens": negative_txt_seq_lens,
+                }
+            else:
+                negative_kwargs = None
 
-        gathered = cfg_group.all_gather(local_pred, separate_tensors=True)
-        if cfg_rank == 0:
-            noise_pred = gathered[0]
-            neg_noise_pred = gathered[1]
-            comb_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
-            cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
-            noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
-            noise_pred = comb_pred * (cond_norm / noise_norm)
-            latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
-        cfg_group.broadcast(latents, src=0)
-    else:
-        # fallback: run positive then negative sequentially on one rank
-        ...
+            # Predict noise with automatic CFG parallel handling
+            # - In CFG parallel mode: rank0 computes positive, rank1 computes negative
+            # - Automatically gathers results and combines them on rank0
+            noise_pred = self.predict_noise_maybe_with_cfg(
+                do_true_cfg=do_true_cfg,
+                true_cfg_scale=true_cfg_scale,
+                positive_kwargs=positive_kwargs,
+                negative_kwargs=negative_kwargs,
+                cfg_normalize=cfg_normalize,
+            )
+
+            # Step scheduler with automatic CFG synchronization
+            # - Only rank0 computes the scheduler step
+            # - Automatically broadcasts updated latents to all ranks
+            latents = self.scheduler_step_maybe_with_cfg(
+                noise_pred, t, latents, do_true_cfg
+            )
+
+        return latents
+```
+
+**How it works:**
+1. Prepare separate `positive_kwargs` and `negative_kwargs` for conditional and unconditional predictions
+2. Call `predict_noise_maybe_with_cfg()` which:
+   - Detects if CFG parallel is enabled (`get_classifier_free_guidance_world_size() > 1`)
+   - Distributes computation: rank0 processes positive, rank1 processes negative
+   - Gathers predictions and combines them using `combine_cfg_noise()` on rank0
+   - Returns combined noise prediction (only valid on rank0)
+3. Call `scheduler_step_maybe_with_cfg()` which:
+   - Only rank0 computes the scheduler step
+   - Broadcasts the updated latents to all ranks for synchronization
+
+**How to customize**
+
+Some pipelines may need to customize the following functions in `CFGParallelMixin`:
+1. You may need to edit `predict_noise` function for custom behaviors.
+```python
+def predict_noise(self, *args, **kwargs):
+    """
+    Forward pass through transformer to predict noise.
+
+    Subclasses should override this if they need custom behavior,
+    but the default implementation calls self.transformer.
+    """
+    return self.transformer(*args, **kwargs)[0]
+
+```
+2. The default normalization function after combining the noise predictions from both branches is as follows. You may need to customize it.
+```python
+def cfg_normalize_function(self, noise_pred, comb_pred):
+    """
+    Normalize the combined noise prediction.
+
+    Args:
+        noise_pred: positive noise prediction
+        comb_pred: combined noise prediction after CFG
+
+    Returns:
+        Normalized noise prediction tensor
+    """
+    cond_norm = torch.norm(noise_pred, dim=-1, keepdim=True)
+    noise_norm = torch.norm(comb_pred, dim=-1, keepdim=True)
+    noise_pred = comb_pred * (cond_norm / noise_norm)
+    return noise_pred
 ```

docs/user_guide/diffusion/teacache.md

@@ -8,6 +8,7 @@ Enable TeaCache by setting `cache_backend` to `"tea_cache"`:
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 # Simple configuration - model_type is automatically extracted from pipeline.__class__.__name__
 omni = Omni(
@@ -17,7 +18,12 @@ omni = Omni(
         "rel_l1_thresh": 0.2  # Optional, defaults to 0.2
     }
 )
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ### Using Environment Variable
@@ -68,13 +74,19 @@ Controls the balance between speed and quality. Lower values prioritize quality,
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 omni = Omni(
     model="Qwen/Qwen-Image",
     cache_backend="tea_cache",
     cache_config={"rel_l1_thresh": 0.2}
 )
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ## Performance Tuning
@@ -112,6 +124,7 @@ th {
 | `QwenImageEditPipeline` | Qwen-Image-Edit | `Qwen/Qwen-Image-Edit` |
 | `QwenImageEditPlusPipeline` | Qwen-Image-Edit-2509 | `Qwen/Qwen-Image-Edit-2509` |
 | `QwenImageLayeredPipeline` | Qwen-Image-Layered | `Qwen/Qwen-Image-Layered` |
+| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` |
 
 ### VideoGen
 

docs/user_guide/diffusion_acceleration.md

@@ -39,23 +39,24 @@ The following table shows which models are currently supported by each accelerat
 
 | Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention | CFG-Parallel |
 |-------|------------------|:----------:|:-----------:|:-----------:|:----------------:|:----------------:|
-| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image** | `meituan-longcat/LongCat-Image` | ❌ | ✅ | ❌ | ❌ | ✅ |
+| **LongCat-Image-Edit** | `meituan-longcat/LongCat-Image-Edit` | ❌ | ✅ | ❌ | ❌ | ✅ |
+| **Ovis-Image** | `OvisAI/Ovis-Image` | ❌ | ✅ | ❌ | ❌ | ✅ |
 | **Qwen-Image** | `Qwen/Qwen-Image` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-2512** | `Qwen/Qwen-Image-2512` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit** | `Qwen/Qwen-Image-Edit` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Edit-2509** | `Qwen/Qwen-Image-Edit-2509` | ✅ | ✅ | ✅ | ✅ | ✅ |
 | **Qwen-Image-Layered** | `Qwen/Qwen-Image-Layered` | ❌ | ✅ | ✅ | ✅ | ✅ |
-| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ❌ | ✅ |❌ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ❌ |
-| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Z-Image** | `Tongyi-MAI/Z-Image-Turbo` | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **Stable-Diffusion3.5** | `stabilityai/stable-diffusion-3.5` | ❌ | ✅ | ❌ | ❌ | ✅ |
+| **Bagel** | `ByteDance-Seed/BAGEL-7B-MoT` | ✅ | ✅ | ❌ | ❌ | ❌ |
+| **FLUX.1-dev** | `black-forest-labs/FLUX.1-dev` | ❌ | ✅ | ❌ | ❌ | ❌ |
 
 ### VideoGen
 
-| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |
-|-------|------------------|:--------:|:---------:|:----------:|:--------------:|
-| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ❌ | ❌ |
+| Model | Model Identifier | TeaCache | Cache-DiT | Ulysses-SP | Ring-Attention |CFG-Parallel |
+|-------|------------------|:--------:|:---------:|:----------:|:--------------:|:----------------:|
+| **Wan2.2** | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | ❌ | ✅ | ✅ | ✅ | ✅ |
 
 
 ## Performance Benchmarks
@@ -96,6 +97,7 @@ To measure the parallelism methods, we run benchmarks with **Qwen/Qwen-Image** m
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 omni = Omni(
     model="Qwen/Qwen-Image",
@@ -103,13 +105,19 @@ omni = Omni(
     cache_config={"rel_l1_thresh": 0.2}  # Optional, defaults to 0.2
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ### Using Cache-DiT
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 
 omni = Omni(
     model="Qwen/Qwen-Image",
@@ -123,7 +131,12 @@ omni = Omni(
     }
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(
+        num_inference_steps=50,
+    ),
+)
 ```
 
 ### Using Ulysses-SP
@@ -131,6 +144,7 @@ outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_st
 Run text-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2
 
@@ -139,13 +153,17 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 
 Run image-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ulysses_degree = 2
 
@@ -154,8 +172,13 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ulysses_degree=ulysses_degree)
 )
 
-outputs = omni.generate(prompt="turn this cat to a dog",
-        pil_image=input_image, num_inference_steps=50)
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "multi_modal_data": {"image": input_image}
+    },
+    OmniDiffusionSamplingParams(num_inference_steps=50),
+)
 ```
 
 ### Using Ring-Attention
@@ -163,6 +186,7 @@ outputs = omni.generate(prompt="turn this cat to a dog",
 Run text-to-image:
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 ring_degree = 2
 
@@ -171,7 +195,10 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(ring_degree=2)
 )
 
-outputs = omni.generate(prompt="A cat sitting on a windowsill", num_inference_steps=50, width=2048, height=2048)
+outputs = omni.generate(
+    "A cat sitting on a windowsill",
+    OmniDiffusionSamplingParams(num_inference_steps=50, width=2048, height=2048),
+)
 ```
 
 ### Using CFG-Parallel
@@ -182,6 +209,7 @@ CFG-Parallel splits the CFG positive/negative branches across GPUs. Use it when
 
 ```python
 from vllm_omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.diffusion.data import DiffusionParallelConfig
 cfg_parallel_size = 2
 
@@ -190,8 +218,13 @@ omni = Omni(
     parallel_config=DiffusionParallelConfig(cfg_parallel_size=cfg_parallel_size)
 )
 
-outputs = omni.generate(prompt="turn this cat to a dog",
-        pil_image=input_image, num_inference_steps=50, true_cfg_scale=4.0)
+outputs = omni.generate(
+    {
+        "prompt": "turn this cat to a dog",
+        "multi_modal_data": {"image": input_image}
+    },
+    OmniDiffusionSamplingParams(num_inference_steps=50, true_cfg_scale=4.0),
+)
 ```
 
 ## Documentation

docs/user_guide/examples/offline_inference/bagel.md

@@ -0,0 +1,187 @@
+# BAGEL-7B-MoT
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/bagel>.
+
+
+## Set up
+
+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.
+
+## Run examples
+
+**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
+
+Get into the bagel folder
+
+```bash
+cd examples/offline_inference/bagel
+```
+
+### Modality Control
+
+BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
+
+#### Text to Image (text2img)
+
+- **Pipeline**: Text → Thinker  → DiT → VAE Decode → Image
+- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
+- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
+
+Generate images from text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat"
+```
+
+#### Image to Image (img2img)
+
+- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
+- **Stages Used**: Stage 1 (DiT) only
+- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
+
+Transform images based on text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2img \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Let the woman wear a blue dress"
+```
+
+#### Image to Text (img2text)
+
+- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
+
+Generate text descriptions from images:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image in detail"
+```
+
+#### Text to Text (text2text)
+
+- **Pipeline**: Text → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: No visual components involved, operates as pure language model
+
+Pure text generation:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?"
+
+# You can load prompts from a text file (one prompt per line):  
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --txt-prompts /path/to/prompts.txt
+```
+
+### Inference Steps
+
+Control the number of inference steps for image generation:
+
+```bash
+# You can adjust steps to 100 to improve image quality
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --steps 50 \
+                  --prompts "A cute cat"
+```
+
+### Key arguments
+
+BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
+
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
+
+#### 📌 Command Line Arguments (end2end.py)
+
+| Argument               | Type   | Default                       | Description                                                  |
+| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
+| `--model`              | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name                                           |
+| `--modality`           | choice | `text2img`                    | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
+| `--prompts`            | list   | `None`                        | Input text prompts directly                                  |
+| `--txt-prompts`        | string | `None`                        | Path to txt file with one prompt per line                    |
+| `--image-path`         | string | `None`                        | Input image path (for `img2img`/`img2text`)                  |
+| `--steps`              | int    | `50`                          | Number of inference steps                                    |
+| `--stage-configs-path` | string | `None`                        | Custom stage config file path                                |
+| `--worker-backend`     | choice | `process`                     | Worker backend: `process` or `ray`                           |
+| `--ray-address`        | string | `None`                        | Ray cluster address                                          |
+| `--enable-stats`       | flag   | `False`                       | Enable statistics logging                                    |
+| `--init-sleep-seconds` | int    | `20`                          | Initialization sleep time                                    |
+| `--batch-timeout`      | int    | `5`                           | Batch timeout                                                |
+| `--init-timeout`       | int    | `300`                         | Initialization timeout                                       |
+
+------
+
+#### ⚙️ Stage Configuration Parameters (bagel.yaml)
+
+ **Stage 0 - Thinker (LLM Stage)**
+
+| Parameter                        | Value                           | Description              |
+| :------------------------------- | :------------------------------ | :----------------------- |
+| `stage_type`                     | `llm`                           | Stage type               |
+| `devices`                        | `"0"`                           | GPU device ID            |
+| `max_batch_size`                 | `1`                             | Maximum batch size       |
+| `model_stage`                    | `thinker`                       | Model stage identifier   |
+| `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
+| `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
+| `tensor_parallel_size`           | `1`                             | Tensor parallel size     |
+| `max_num_batched_tokens`         | `32768`                         | Maximum batched tokens   |
+| `omni_kv_config.need_send_cache` | `true`                          | Whether to send KV cache |
+
+------
+
+**Stage 1 - DiT (Diffusion Stage)**
+
+| Parameter                        | Value       | Description                 |
+| :------------------------------- | :---------- | :-------------------------- |
+| `stage_type`                     | `diffusion` | Stage type                  |
+| `devices`                        | `"0"`       | GPU device ID               |
+| `max_batch_size`                 | `1`         | Maximum batch size          |
+| `model_stage`                    | `dit`       | Model stage identifier      |
+| `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
+| `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |
+| `engine_input_source`            | `[0]`       | Input source from Stage 0   |
+
+------
+
+#### 🔗 Runtime Configuration
+
+| Parameter             | Value   | Description                      |
+| :-------------------- | :------ | :------------------------------- |
+| `window_size`         | `-1`    | Window size (-1 means unlimited) |
+| `max_inflight`        | `1`     | Maximum inflight requests        |
+| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
+
+## FAQ
+
+- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
+
+```bash
+sudo apt update
+sudo apt install ffmpeg
+```
+
+- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
+
+| Stage               | VRAM                         |
+| :------------------ | :--------------------------- |
+| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
+| Stage-1 (DiT)       | **26.50 GiB**                |
+| Total               | **~42 GiB + KV Cache**       |
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/bagel/end2end.py"
+    ``````

docs/user_guide/examples/offline_inference/image_to_image.md

@@ -50,6 +50,12 @@ Key arguments:
 - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/offline_inference/image_to_video.md

@@ -7,12 +7,19 @@ This example demonstrates how to generate videos from images using Wan2.2 Image-
 
 ## Local CLI Usage
 
+Download the example image:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+```
+
 ### Wan2.2-I2V-A14B-Diffusers (MoE)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
@@ -27,11 +34,12 @@ python image_to_video.py \
 ```
 
 ### Wan2.2-TI2V-5B-Diffusers (Unified)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
@@ -57,6 +65,12 @@ Key arguments:
 - `--num_inference_steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
+- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/offline_inference/lora_inference.md

@@ -0,0 +1,108 @@
+# LoRA Inference(Diffusion)
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/lora_inference>.
+
+
+This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
+
+## Overview
+
+Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
+
+- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
+- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
+
+Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
+
+## Usage
+
+### Pre-loaded LoRA (via --lora-path)
+
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_preloaded.png
+```
+
+**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
+
+### Per-request LoRA (via --lora-request-path)
+
+Load a LoRA adapter on-demand for each request:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-request-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_per_request.png
+```
+
+### No LoRA
+
+If no LoRA request is provided, we will use the base model without any LoRA adapters:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_no_lora.png
+```
+
+## Parameters
+
+### LoRA Parameters
+
+- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
+- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
+- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### Standard Parameters
+
+- `--prompt`: Text prompt for image generation (required)
+- `--seed`: Random seed for reproducibility (default: 42)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--output`: Output file path (default: `lora_output.png`)
+
+## How LoRA Works
+
+All LoRA adapters are handled uniformly:
+
+1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
+2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
+3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
+
+The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
+
+## LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
+## Example materials
+
+??? abstract "lora_inference.py"
+    ``````py
+    --8<-- "examples/offline_inference/lora_inference/lora_inference.py"
+    ``````

docs/user_guide/examples/offline_inference/qwen3_tts.md

@@ -0,0 +1,103 @@
+# Qwen3-TTS
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/qwen3_tts>.
+
+
+This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.
+
+## Model Overview
+
+Qwen3 TTS provides multiple task variants for speech generation:
+
+- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
+- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
+- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
+
+## 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.
+
+### ROCm Dependencies
+
+You will need to install these two dependencies `onnxruntime-rocm` and `sox`.
+
+```
+pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+pip install onnxruntime-rocm sox
+```
+
+## Quick Start
+
+Run a single sample for a task:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Generated audio files are saved to `output_audio/` by default.
+
+## Task Usage
+
+### CustomVoice
+
+Single sample:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Batch sample (multiple prompts in one run):
+
+```
+python end2end.py --query-type CustomVoice --use-batch-sample
+```
+
+### VoiceDesign
+
+Single sample:
+
+```
+python end2end.py --query-type VoiceDesign
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type VoiceDesign --use-batch-sample
+```
+
+### Base (Voice Clone)
+
+Single sample:
+
+```
+python end2end.py --query-type Base
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type Base --use-batch-sample
+```
+
+Mode selection for Base:
+
+- `--mode-tag icl` (default): standard mode
+- `--mode-tag xvec_only`: enable `x_vector_only_mode` in the request
+
+Examples:
+
+```
+python end2end.py --query-type Base --mode-tag icl
+```
+
+## Notes
+
+- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
+- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.
+
+## Example materials
+
+??? abstract "end2end.py"
+    ``````py
+    --8<-- "examples/offline_inference/qwen3_tts/end2end.py"
+    ``````

docs/user_guide/examples/offline_inference/text_to_audio.md

@@ -0,0 +1,47 @@
+# Text-To-Audio
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/text_to_audio>.
+
+
+The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
+
+## Prerequisites
+
+If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you have access:
+
+1. **Accept Model License**: Visit the model page on Hugging Face (e.g., [stabilityai/stable-audio-open-1.0]) and accept the user agreement.
+2. **Authenticate**: Log in to Hugging Face locally to access the gated model.
+   ```bash
+   huggingface-cli login
+   ```
+
+## Local CLI Usage
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative_prompt "Low quality" \
+  --seed 42 \
+  --guidance_scale 7.0 \
+  --audio_length 10.0 \
+  --num_inference_steps 100 \
+  --output stable_audio_output.wav
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--seed`: integer seed for deterministic generation.
+- `--guidance_scale`: classifier-free guidance scale.
+- `--audio_length`: audio duration in seconds.
+- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--output`: path to save the generated WAV file.
+
+## Example materials
+
+??? abstract "text_to_audio.py"
+    ``````py
+    --8<-- "examples/offline_inference/text_to_audio/text_to_audio.py"
+    ``````

docs/user_guide/examples/offline_inference/text_to_image.md

@@ -23,7 +23,7 @@ if __name__ == "__main__":
     images[0].save("coffee.png")
 ```
 
-Or put more than one prompt in a request, processing them sequentially.
+Or put more than one prompt in a request.
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
@@ -40,6 +40,42 @@ if __name__ == "__main__":
       image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md
+    ) to let the model accept a longer batch of prompts.
+
+Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
+This is useful for models that support negative prompts.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    outputs = omni.generate([
+      {
+        "prompt": "a cup of coffee on a table"，
+        "negative_prompt": "low resolution"
+      },
+      {
+        "prompt": "a toy dinosaur on a sandy beach"，
+        "negative_prompt": "cinematic, realistic"
+      }
+    ])
+    for i, output in enumerate(outputs):
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
+```
+
 ## Local CLI Usage
 
 ```bash
@@ -64,6 +100,12 @@ Key arguments:
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
 

docs/user_guide/examples/offline_inference/text_to_video.md

@@ -12,10 +12,11 @@ python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
-  --width 640 \
-  --num_frames 32 \
+  --width 832 \
+  --num_frames 33 \
   --guidance_scale 4.0 \
   --guidance_scale_high 3.0 \
+  --flow_shift 12.0 \
   --num_inference_steps 40 \
   --fps 16 \
   --output t2v_out.mp4
@@ -24,13 +25,19 @@ python text_to_video.py \
 Key arguments:
 
 - `--prompt`: text description (string).
-- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
 - `--num_frames`: Number of frames (Wan default is 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
+- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
 - `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
-- `--boundary_ratio`: Boundary split ratio for low/high DiT.
+- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](https://github.com/vllm-project/vllm-omni/tree/main/docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 ## Example materials
 

docs/user_guide/examples/online_serving/bagel.md

@@ -0,0 +1,244 @@
+# BAGEL-7B-MoT
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/bagel>.
+
+
+## 🛠️ Installation
+
+Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/README.md)
+
+## Run examples (BAGEL-7B-MoT)
+
+**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
+
+### Launch the Server
+
+```bash
+# Use default configuration
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
+```
+
+Or use the convenience script:
+
+```bash
+cd /workspace/vllm-omni/examples/online_serving/bagel
+bash run_server.sh
+```
+
+If you have a custom stage configs file, launch the server with the command below:
+
+```bash
+vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --stage-configs-path /path/to/stage_configs_file
+```
+
+### Send Multi-modal Request
+
+Get into the bagel folder:
+
+```bash
+cd examples/online_serving/bagel
+```
+
+Send request via Python
+
+```bash
+python openai_chat_client.py --prompt "A cute cat" --modality text2img
+```
+
+The Python client supports the following command-line arguments:
+
+- `--prompt` (or `-p`): Text prompt for generation (default: `A cute cat`)
+- `--output` (or `-o`): Output file path for image results (default: `bagel_output.png`)
+- `--server` (or `-s`): Server URL (default: `http://localhost:8091`)
+- `--image-url` (or `-i`): Input image URL or local file path (for img2img/img2text modes)
+- `--modality` (or `-m`): Task modality (default: `text2img`). Options: `text2img`, `img2img`, `img2text`, `text2text`
+- `--height`: Image height in pixels (default: 512)
+- `--width`: Image width in pixels (default: 512)
+- `--steps`: Number of inference steps (default: 25)
+- `--seed`: Random seed (default: 42)
+- `--negative`: Negative prompt for image generation
+
+Example with custom parameters:
+
+```bash
+python openai_chat_client.py \
+    --prompt "A futuristic city" \
+    --modality text2img \
+    --height 768 \
+    --width 768 \
+    --steps 50 \
+    --seed 42 \
+    --negative "blurry, low quality"
+```
+
+## Modality Control
+
+BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
+
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](https://github.com/vllm-project/vllm-omni/tree/main/vllm_omni/model_executor/stage_configs/bagel.yaml)
+
+| Modality    | Input        | Output | Description                            |
+| ----------- | ------------ | ------ | -------------------------------------- |
+| `text2img`  | Text         | Image  | Generate images from text prompts      |
+| `img2img`   | Image + Text | Image  | Transform images using text guidance   |
+| `img2text`  | Image + Text | Text   | Generate text descriptions from images |
+| `text2text` | Text         | Text   | Pure text generation                   |
+
+### Text to Image (text2img)
+
+Generate images from text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "A beautiful sunset over mountains" \
+    --modality text2img \
+    --output sunset.png \
+    --steps 50
+```
+
+**Using curl**
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
+    "modalities": ["image"],
+    "height": 512,
+    "width": 512,
+    "num_inference_steps": 50,
+    "seed": 42
+  }'
+```
+
+
+### Image to Image (img2img)
+
+Transform images based on text prompts:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Make the cat stand up" \
+    --modality img2img \
+    --image-url /path/to/input.jpg \
+    --output transformed.png
+```
+
+**Using curl**
+
+```bash
+IMAGE_BASE64=$(base64 -w 0 cat.jpg)
+
+cat <<EOF > payload.json
+{
+    "messages": [{
+      "role": "user",
+      "content": [
+        {"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
+        {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
+      ]
+    }],
+    "modalities": ["image"],
+    "height": 512,
+    "width": 512,
+    "num_inference_steps": 50,
+    "seed": 42
+}
+EOF
+
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @payload.json
+
+```
+
+### Image to Text (img2text)
+
+Generate text descriptions from images:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "Describe this image in detail" \
+    --modality img2text \
+    --image-url /path/to/image.jpg
+```
+
+**Using curl**
+
+```bash
+IMAGE_BASE64=$(base64 -w 0 cat.jpg)
+
+cat <<EOF > payload.json
+{
+  "messages": [{
+    "role": "user",
+    "content": [
+      {"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
+      {"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
+    ]
+  }],
+  "modalities": ["text"]
+}
+EOF
+
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d @payload.json
+```
+
+### Text to Text (text2text)
+
+Pure text generation:
+
+**Using Python client**
+
+```bash
+python openai_chat_client.py \
+    --prompt "What is the capital of France?" \
+    --modality text2text
+```
+
+**Using curl**
+
+```bash
+curl http://localhost:8091/v1/chat/completions \
+  -H "Content-Type: application/json" \
+  -d '{
+    "messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}]
+    "modalities": ["text"]
+  }'
+```
+
+## FAQ
+
+- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
+
+```bash
+sudo apt update
+sudo apt install ffmpeg
+```
+
+- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
+
+| Stage               | VRAM                         |
+| :------------------ | :--------------------------- |
+| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
+| Stage-1 (DiT)       | **26.50 GiB**                |
+| Total               | **~42 GiB + KV Cache**       |
+
+## Example materials
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/bagel/openai_chat_client.py"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/bagel/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/lora_inference.md

@@ -0,0 +1,72 @@
+# LoRA Inference(Diffusion)
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/lora_inference>.
+
+
+This example shows how to use **per-request LoRA** with vLLM-Omni diffusion models via the OpenAI-compatible Chat Completions API.
+
+> Note: The LoRA adapter path must be readable on the **server** machine (usually a local path or a mounted directory).
+> Note: This example uses `/v1/chat/completions`. LoRA payloads for other OpenAI endpoints are not implemented here.
+
+## Start Server
+
+```bash
+# Pick a diffusion model (examples)
+# export MODEL=stabilityai/stable-diffusion-3.5-medium
+# export MODEL=Qwen/Qwen-Image
+
+bash run_server.sh
+```
+
+## Call API (curl)
+
+```bash
+# Required: local LoRA folder on the server
+export LORA_PATH=/path/to/lora_adapter
+
+# Optional
+export SERVER=http://localhost:8091
+export PROMPT="A piece of cheesecake"
+export LORA_NAME=my_lora
+export LORA_SCALE=1.0
+# Optional: if omitted, the server derives a stable id from LORA_PATH.
+# export LORA_INT_ID=123
+
+bash run_curl_lora_inference.sh
+```
+
+## Call API (Python)
+
+```bash
+python openai_chat_client.py \
+  --prompt "A piece of cheesecake" \
+  --lora-path /path/to/lora_adapter \
+  --lora-name my_lora \
+  --lora-scale 1.0 \
+  --output output.png
+```
+
+## LoRA Format
+
+LoRA adapters should be in PEFT format, for example:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```
+
+## Example materials
+
+??? abstract "openai_chat_client.py"
+    ``````py
+    --8<-- "examples/online_serving/lora_inference/openai_chat_client.py"
+    ``````
+??? abstract "run_curl_lora_inference.sh"
+    ``````sh
+    --8<-- "examples/online_serving/lora_inference/run_curl_lora_inference.sh"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/lora_inference/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/qwen3_omni.md

@@ -15,6 +15,12 @@ Please refer to [README.md](https://github.com/vllm-project/vllm-omni/tree/main/
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
 ```
 
+If you want to open async chunking for qwen3-omni, launch the server with command below
+
+```bash
+vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /vllm_omni/model_executor/stage_configs/qwen3_omni_moe_async_chunk.yaml
+```
+
 If you have custom stage configs file, launch the server with command below
 ```bash
 vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -0,0 +1,189 @@
+# Qwen3-TTS Online Serving
+
+Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/qwen3_tts>.
+
+
+This directory contains examples for running Qwen3-TTS models with vLLM-Omni's online serving API.
+
+## Supported Models
+
+| Model | Task Type | Description |
+|-------|-----------|-------------|
+| `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | CustomVoice | Predefined speaker voices with optional style control |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | VoiceDesign | Natural language voice style description |
+| `Qwen/Qwen3-TTS-12Hz-1.7B-Base` | Base | Voice cloning from reference audio |
+
+## Quick Start
+
+### 1. Start the Server
+
+```bash
+# CustomVoice model (default)
+./run_server.sh
+
+# Or specify task type
+./run_server.sh CustomVoice
+./run_server.sh VoiceDesign
+./run_server.sh Base
+```
+
+### 2. Run the Client
+
+```bash
+# CustomVoice: Use predefined speaker
+python openai_speech_client.py \
+    --text "你好，我是通义千问" \
+    --voice Vivian \
+    --language Chinese
+
+# CustomVoice with style instruction
+python openai_speech_client.py \
+    --text "今天天气真好" \
+    --voice Ryan \
+    --instructions "用开心的语气说"
+
+# VoiceDesign: Describe the voice style
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign \
+    --task-type VoiceDesign \
+    --text "哥哥，你回来啦" \
+    --instructions "体现撒娇稚嫩的萝莉女声，音调偏高"
+
+# Base: Voice cloning
+python openai_speech_client.py \
+    --model Qwen/Qwen3-TTS-12Hz-1.7B-Base \
+    --task-type Base \
+    --text "Hello, this is a cloned voice" \
+    --ref-audio /path/to/reference.wav \
+    --ref-text "Original transcript of the reference audio"
+```
+
+### 3. Using curl
+
+```bash
+# Simple TTS request
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "Hello, how are you?",
+        "voice": "Vivian",
+        "language": "English"
+    }' --output output.wav
+
+# With style instruction
+curl -X POST http://localhost:8000/v1/audio/speech \
+    -H "Content-Type: application/json" \
+    -d '{
+        "input": "I am so excited!",
+        "voice": "Vivian",
+        "instructions": "Speak with great enthusiasm"
+    }' --output excited.wav
+
+# List available voices in CustomVoice models
+curl http://localhost:8000/v1/audio/voices
+```
+
+## API Reference
+
+### Endpoint
+
+```
+POST /v1/audio/speech
+```
+
+This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/docs/api-reference/audio/createSpeech) format with additional Qwen3-TTS parameters.
+
+### Request Body
+
+```json
+{
+    "input": "Text to synthesize",
+    "voice": "Vivian",
+    "response_format": "wav",
+    "task_type": "CustomVoice",
+    "language": "Auto",
+    "instructions": "Optional style instructions",
+    "ref_audio": "URL or base64 for voice cloning",
+    "ref_text": "Reference audio transcript",
+    "x_vector_only_mode": false,
+    "max_new_tokens": 2048
+}
+```
+
+> **Note:** The `model` field is optional when serving a single model, as the server already knows which model is loaded.
+
+### Response
+
+Returns audio data in the requested format (default: WAV).
+
+## Parameters
+
+### Standard OpenAI Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `input` | string | required | Text to synthesize |
+| `voice` | string | "Vivian" | Speaker/voice name |
+| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
+| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
+| `model` | string | optional | Model name (optional when serving single model) |
+
+### Qwen3-TTS Parameters
+
+| Parameter | Type | Default | Description |
+|-----------|------|---------|-------------|
+| `task_type` | string | "CustomVoice" | Task: CustomVoice, VoiceDesign, or Base |
+| `language` | string | "Auto" | Language: Auto, Chinese, English, Japanese, Korean |
+| `instructions` | string | "" | Voice style/emotion instructions |
+| `max_new_tokens` | int | 2048 | Maximum tokens to generate |
+
+### Voice Clone Parameters (Base task)
+
+| Parameter | Type | Required | Description |
+|-----------|------|----------|-------------|
+| `ref_audio` | string | Yes* | Reference audio (file path, URL, or base64) |
+| `ref_text` | string | No | Transcript of reference audio (for ICL mode) |
+| `x_vector_only_mode` | bool | false | Use speaker embedding only (no ICL) |
+
+## Python Usage
+
+```python
+import httpx
+
+# Simple request
+response = httpx.post(
+    "http://localhost:8000/v1/audio/speech",
+    json={
+        "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+        "input": "Hello world",
+        "voice": "Vivian",
+    },
+    timeout=300.0,
+)
+
+with open("output.wav", "wb") as f:
+    f.write(response.content)
+```
+
+## Limitations
+
+- **No streaming**: Audio is generated completely before being returned. Streaming will be supported after the pipeline is disaggregated (see RFC #938).
+- **Single request**: Batch processing is not yet optimized for online serving.
+
+## Troubleshooting
+
+1. **Connection refused**: Make sure the server is running on the correct port
+2. **Out of memory**: Reduce `--gpu-memory-utilization` in run_server.sh
+3. **Unsupported speaker**: Check supported speakers via model documentation
+4. **Voice clone fails**: Ensure you're using the Base model variant for voice cloning
+
+## Example materials
+
+??? abstract "openai_speech_client.py"
+    ``````py
+    --8<-- "examples/online_serving/qwen3_tts/openai_speech_client.py"
+    ``````
+??? abstract "run_server.sh"
+    ``````sh
+    --8<-- "examples/online_serving/qwen3_tts/run_server.sh"
+    ``````

docs/user_guide/examples/online_serving/text_to_image.md

@@ -119,6 +119,7 @@ Use `extra_body` to pass generation parameters:
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |
+| `--cfg-parallel-size`.   | int   | 1       | Number of GPUs for CFG parallelism |
 
 ## Response Format
 

examples/offline_inference/bagel/README.md

@@ -0,0 +1,177 @@
+# BAGEL-7B-MoT
+
+## Set up
+
+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.
+
+## Run examples
+
+**Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, please modify the stage configuration to distribute the model across devices.
+
+Get into the bagel folder
+
+```bash
+cd examples/offline_inference/bagel
+```
+
+### Modality Control
+
+BAGEL-7B-MoT supports multiple modality modes. You can control the mode using the `--modality` argument:
+
+#### Text to Image (text2img)
+
+- **Pipeline**: Text → Thinker  → DiT → VAE Decode → Image
+- **Stages Used**: Stage 0 (Thinker) + Stage 1 (DiT)
+- **KV Transfer**: Thinker sends KV cache to DiT for conditioned generation
+
+Generate images from text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --prompts "A cute cat"
+```
+
+#### Image to Image (img2img)
+
+- **Pipeline**: Image → VAE Encode → DiT → VAE Decode → New Image
+- **Stages Used**: Stage 1 (DiT) only
+- **Special**: Bypasses the Thinker stage, direct image-to-image transformation
+
+Transform images based on text prompts:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2img \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Let the woman wear a blue dress"
+```
+
+#### Image to Text (img2text)
+
+- **Pipeline**: Image → ViT + VAE Encode → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: Uses both VAE latent encoding AND ViT semantic encoding for comprehensive image understanding
+
+Generate text descriptions from images:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality img2text \
+                  --image-path /path/to/image.jpg \
+                  --prompts "Describe this image in detail"
+```
+
+#### Text to Text (text2text)
+
+- **Pipeline**: Text → Thinker → Text Output
+- **Stages Used**: Stage 0 (Thinker) only
+- **Special**: No visual components involved, operates as pure language model
+
+Pure text generation:
+
+```bash
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --prompts "What is the capital of France?"
+
+# You can load prompts from a text file (one prompt per line):  
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2text \
+                  --txt-prompts /path/to/prompts.txt
+```
+
+### Inference Steps
+
+Control the number of inference steps for image generation:
+
+```bash
+# You can adjust steps to 100 to improve image quality
+python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
+                  --modality text2img \
+                  --steps 50 \
+                  --prompts "A cute cat"
+```
+
+### Key arguments
+
+BAGEL-7B-MoT supports **multiple modality modes** for different use cases.
+
+The default yaml configuration deploys Thinker and DiT on the same GPU. You can use the default configuration file: [`bagel.yaml`](../../../vllm_omni/model_executor/stage_configs/bagel.yaml)
+
+#### 📌 Command Line Arguments (end2end.py)
+
+| Argument               | Type   | Default                       | Description                                                  |
+| :--------------------- | :----- | :---------------------------- | :----------------------------------------------------------- |
+| `--model`              | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or name                                           |
+| `--modality`           | choice | `text2img`                    | Modality mode: `text2img`, `img2img`, `img2text`, `text2text` |
+| `--prompts`            | list   | `None`                        | Input text prompts directly                                  |
+| `--txt-prompts`        | string | `None`                        | Path to txt file with one prompt per line                    |
+| `--image-path`         | string | `None`                        | Input image path (for `img2img`/`img2text`)                  |
+| `--steps`              | int    | `50`                          | Number of inference steps                                    |
+| `--stage-configs-path` | string | `None`                        | Custom stage config file path                                |
+| `--worker-backend`     | choice | `process`                     | Worker backend: `process` or `ray`                           |
+| `--ray-address`        | string | `None`                        | Ray cluster address                                          |
+| `--enable-stats`       | flag   | `False`                       | Enable statistics logging                                    |
+| `--init-sleep-seconds` | int    | `20`                          | Initialization sleep time                                    |
+| `--batch-timeout`      | int    | `5`                           | Batch timeout                                                |
+| `--init-timeout`       | int    | `300`                         | Initialization timeout                                       |
+
+------
+
+#### ⚙️ Stage Configuration Parameters (bagel.yaml)
+
+ **Stage 0 - Thinker (LLM Stage)**
+
+| Parameter                        | Value                           | Description              |
+| :------------------------------- | :------------------------------ | :----------------------- |
+| `stage_type`                     | `llm`                           | Stage type               |
+| `devices`                        | `"0"`                           | GPU device ID            |
+| `max_batch_size`                 | `1`                             | Maximum batch size       |
+| `model_stage`                    | `thinker`                       | Model stage identifier   |
+| `model_arch`                     | `BagelForConditionalGeneration` | Model architecture       |
+| `gpu_memory_utilization`         | `0.4`                           | GPU memory utilization   |
+| `tensor_parallel_size`           | `1`                             | Tensor parallel size     |
+| `max_num_batched_tokens`         | `32768`                         | Maximum batched tokens   |
+| `omni_kv_config.need_send_cache` | `true`                          | Whether to send KV cache |
+
+------
+
+**Stage 1 - DiT (Diffusion Stage)**
+
+| Parameter                        | Value       | Description                 |
+| :------------------------------- | :---------- | :-------------------------- |
+| `stage_type`                     | `diffusion` | Stage type                  |
+| `devices`                        | `"0"`       | GPU device ID               |
+| `max_batch_size`                 | `1`         | Maximum batch size          |
+| `model_stage`                    | `dit`       | Model stage identifier      |
+| `gpu_memory_utilization`         | `0.4`       | GPU memory utilization      |
+| `omni_kv_config.need_recv_cache` | `true`      | Whether to receive KV cache |
+| `engine_input_source`            | `[0]`       | Input source from Stage 0   |
+
+------
+
+#### 🔗 Runtime Configuration
+
+| Parameter             | Value   | Description                      |
+| :-------------------- | :------ | :------------------------------- |
+| `window_size`         | `-1`    | Window size (-1 means unlimited) |
+| `max_inflight`        | `1`     | Maximum inflight requests        |
+| `shm_threshold_bytes` | `65536` | Shared memory threshold (64KB)   |
+
+## FAQ
+
+- If you encounter an error about the backend of librosa, try to install ffmpeg with the command below.
+
+```bash
+sudo apt update
+sudo apt install ffmpeg
+```
+
+- If you don’t know how much VRAM is needed for the model or encounter the OOM error, you can try to decrease the max_model_len.
+
+| Stage               | VRAM                         |
+| :------------------ | :--------------------------- |
+| Stage-0 (Thinker)   | **15.04 GiB** **+ KV Cache** |
+| Stage-1 (DiT)       | **26.50 GiB**                |
+| Total               | **~42 GiB + KV Cache**       |

examples/offline_inference/bagel/end2end.py

@@ -0,0 +1,189 @@
+import argparse
+import os
+from typing import cast
+
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniPromptType
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model",
+        default="ByteDance-Seed/BAGEL-7B-MoT",
+        help="Path to merged model directory.",
+    )
+    parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    parser.add_argument("--prompt_type", default="text", choices=["text"])
+
+    parser.add_argument(
+        "--modality",
+        default="text2img",
+        choices=["text2img", "img2img", "img2text", "text2text"],
+        help="Modality mode to control stage execution.",
+    )
+
+    parser.add_argument(
+        "--image-path",
+        type=str,
+        default=None,
+        help="Path to input image for img2img.",
+    )
+
+    # OmniLLM init args
+    parser.add_argument("--enable-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)
+    parser.add_argument("--init-timeout", type=int, default=300)
+    parser.add_argument("--shm-threshold-bytes", type=int, default=65536)
+    parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"])
+    parser.add_argument("--ray-address", type=str, default=None)
+    parser.add_argument("--stage-configs-path", type=str, default=None)
+    parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
+
+    args = parser.parse_args()
+    return args
+
+
+def main():
+    args = parse_args()
+    model_name = args.model
+    prompts: list[OmniPromptType] = []
+    try:
+        # Preferred: load from txt file (one prompt per line)
+        if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
+            with open(args.txt_prompts, encoding="utf-8") as f:
+                lines = [ln.strip() for ln in f.readlines()]
+            prompts = [ln for ln in lines if ln != ""]
+            print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
+        else:
+            prompts = args.prompts
+    except Exception as e:
+        print(f"[Error] Failed to load prompts: {e}")
+        raise
+
+    if not prompts:
+        # Default prompt for text2img test if none provided
+        prompts = ["<|im_start|>A cute cat<|im_end|>"]
+        print(f"[Info] No prompts provided, using default: {prompts}")
+    omni_outputs = []
+
+    from PIL import Image
+
+    if args.modality == "img2img":
+        from PIL import Image
+
+        from vllm_omni.entrypoints.omni_diffusion import OmniDiffusion
+
+        print("[Info] Running in img2img mode (Stage 1 only)")
+        client = OmniDiffusion(model=model_name)
+
+        if args.image_path:
+            if os.path.exists(args.image_path):
+                loaded_image = Image.open(args.image_path).convert("RGB")
+                prompts = [
+                    {
+                        "prompt": cast(str, p),
+                        "multi_modal_data": {"image": loaded_image},
+                    }
+                    for p in prompts
+                ]
+            else:
+                print(f"[Warning] Image path {args.image_path} does not exist.")
+
+        result = client.generate(
+            prompts,
+            OmniDiffusionSamplingParams(
+                seed=52,
+                need_kv_receive=False,
+                num_inference_steps=args.steps,
+            ),
+        )
+
+        # Ensure result is a list for iteration
+        if not isinstance(result, list):
+            omni_outputs = [result]
+        else:
+            omni_outputs = result
+
+    else:
+        from vllm_omni.entrypoints.omni import Omni
+
+        omni_kwargs = {}
+        if args.stage_configs_path:
+            omni_kwargs["stage_configs_path"] = args.stage_configs_path
+
+        omni_kwargs.update(
+            {
+                "log_stats": args.enable_stats,
+                "init_sleep_seconds": args.init_sleep_seconds,
+                "batch_timeout": args.batch_timeout,
+                "init_timeout": args.init_timeout,
+                "shm_threshold_bytes": args.shm_threshold_bytes,
+                "worker_backend": args.worker_backend,
+                "ray_address": args.ray_address,
+            }
+        )
+
+        omni = Omni(model=model_name, **omni_kwargs)
+
+        formatted_prompts = []
+        for p in args.prompts:
+            if args.modality == "img2text":
+                if args.image_path:
+                    loaded_image = Image.open(args.image_path).convert("RGB")
+                    final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                    prompt_dict = {
+                        "prompt": final_prompt_text,
+                        "multi_modal_data": {"image": loaded_image},
+                        "modalities": ["text"],
+                    }
+                    formatted_prompts.append(prompt_dict)
+            elif args.modality == "text2text":
+                final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
+                prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
+                formatted_prompts.append(prompt_dict)
+            else:
+                # text2img
+                final_prompt_text = f"<|im_start|>{p}<|im_end|>"
+                prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
+                formatted_prompts.append(prompt_dict)
+
+        params_list = omni.default_sampling_params_list
+        if args.modality == "text2img":
+            params_list[0].max_tokens = 1  # type: ignore # The first stage is a SamplingParam (vllm)
+            if len(params_list) > 1:
+                params_list[1].num_inference_steps = args.steps  # type: ignore # The second stage is an OmniDiffusionSamplingParam
+
+        omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
+
+    for i, req_output in enumerate(omni_outputs):
+        images = getattr(req_output, "images", None)
+        if not images and hasattr(req_output, "output"):
+            if isinstance(req_output.output, list):
+                images = req_output.output
+            else:
+                images = [req_output.output]
+
+        if 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}")
+
+    print(omni_outputs)
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/image_to_image/image_edit.py

@@ -77,10 +77,11 @@ from pathlib import Path
 import torch
 from PIL import Image
 
-from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -110,7 +111,7 @@ def parse_args() -> argparse.Namespace:
     parser.add_argument(
         "--negative_prompt",
         type=str,
-        default=" ",
+        default=None,
         required=False,
     )
     parser.add_argument(
@@ -181,6 +182,12 @@ def parse_args() -> argparse.Namespace:
         default=1,
         help="Number of GPUs used for ring sequence parallelism.",
     )
+    parser.add_argument(
+        "--tensor_parallel_size",
+        type=int,
+        default=1,
+        help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
+    )
     parser.add_argument("--layers", type=int, default=4, help="Number of layers to decompose the input image into.")
     parser.add_argument(
         "--resolution",
@@ -273,6 +280,32 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        help="Enable CPU offloading for diffusion models.",
+    )
+    parser.add_argument(
+        "--enable-layerwise-offload",
+        action="store_true",
+        help="Enable layerwise (blockwise) offloading on DiT modules.",
+    )
+    parser.add_argument(
+        "--layerwise-num-gpu-layers",
+        type=int,
+        default=1,
+        help="Number of ready layers (blocks) to keep on GPU during generation.",
+    )
     return parser.parse_args()
 
 
@@ -294,14 +327,13 @@ def main():
     else:
         input_image = input_images
 
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
     parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree, ring_degree=args.ring_degree, cfg_parallel_size=args.cfg_parallel_size
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        cfg_parallel_size=args.cfg_parallel_size,
+        tensor_parallel_size=args.tensor_parallel_size,
     )
 
     # Configure cache based on backend type
@@ -329,15 +361,21 @@ def main():
     # Initialize Omni with appropriate pipeline
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        enable_layerwise_offload=args.enable_layerwise_offload,
+        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        enable_cpu_offload=args.enable_cpu_offload,
     )
     print("Pipeline loaded")
 
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
     # Time profiling for generation
     print(f"\n{'=' * 60}")
     print("Generation Configuration:")
@@ -351,35 +389,58 @@ def main():
     else:
         print(f"  Input image size: {input_image.size}")
     print(
-        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}"
+        f"  Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}"
     )
     print(f"{'=' * 60}\n")
 
     generation_start = time.perf_counter()
-    # Generate edited image
-    generate_kwargs = {
-        "prompt": args.prompt,
-        "pil_image": input_image,
-        "negative_prompt": args.negative_prompt,
-        "generator": generator,
-        "true_cfg_scale": args.cfg_scale,
-        "guidance_scale": args.guidance_scale,
-        "num_inference_steps": args.num_inference_steps,
-        "num_outputs_per_prompt": args.num_outputs_per_prompt,
-        "layers": args.layers,
-        "resolution": args.resolution,
-    }
 
-    outputs = omni.generate(**generate_kwargs)
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
+    # Generate edited image
+    outputs = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": input_image},
+        },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_outputs_per_prompt,
+            layers=args.layers,
+            resolution=args.resolution,
+        ),
+    )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
 
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
     if not outputs:
         raise ValueError("No output generated from omni.generate()")
-    logger.info("Outputs: %s", outputs)
 
     # Extract images from OmniRequestOutput
     # omni.generate() returns list[OmniRequestOutput], extract images from request_output[0].images

examples/offline_inference/image_to_image/image_to_image.md

@@ -47,3 +47,9 @@ Key arguments:
 - `--guidance_scale`: guidance scale for guidance-distilled models (default: 1.0, disabled). Unlike classifier-free guidance (--cfg_scale), guidance-distilled models take the guidance scale directly as an input parameter. Enabled when guidance_scale > 1. Ignored when not using guidance-distilled models.
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

examples/offline_inference/image_to_video/README.md

@@ -4,12 +4,19 @@ This example demonstrates how to generate videos from images using Wan2.2 Image-
 
 ## Local CLI Usage
 
+Download the example image:
+
+```bash
+wget https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg
+```
+
 ### Wan2.2-I2V-A14B-Diffusers (MoE)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-I2V-A14B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
@@ -24,11 +31,12 @@ python image_to_video.py \
 ```
 
 ### Wan2.2-TI2V-5B-Diffusers (Unified)
+
 ```bash
 python image_to_video.py \
   --model Wan-AI/Wan2.2-TI2V-5B-Diffusers \
-  --image input.png \
-  --prompt "A cat playing with yarn, smooth motion" \
+  --image cherry_blossom.jpg \
+  --prompt "Cherry blossoms swaying gently in the breeze, petals falling, smooth motion" \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
   --width 832 \
@@ -54,3 +62,9 @@ Key arguments:
 - `--num_inference_steps`: Number of denoising steps (default 50).
 - `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: Path to save the generated video.
+- `--vae_use_slicing`: Enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: Enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -19,15 +19,18 @@ Usage:
 """
 
 import argparse
+import os
 from pathlib import Path
 
 import numpy as np
 import PIL.Image
 import torch
 
+from vllm_omni.diffusion.data import DiffusionParallelConfig
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -57,6 +60,44 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--output", type=str, default="i2v_output.mp4", help="Path to save the video (mp4).")
     parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        help="Enable CPU offloading for diffusion models.",
+    )
+    parser.add_argument(
+        "--enable-layerwise-offload",
+        action="store_true",
+        help="Enable layerwise (blockwise) offloading on DiT modules.",
+    )
+    parser.add_argument(
+        "--layerwise-num-gpu-layers",
+        type=int,
+        default=1,
+        help="Number of ready layers (blocks) to keep on GPU during 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(
+        "--enforce_eager",
+        action="store_true",
+        help="Disable torch.compile and force eager execution.",
+    )
     return parser.parse_args()
 
 
@@ -73,8 +114,7 @@ def calculate_dimensions(image: PIL.Image.Image, max_area: int = 480 * 832) -> t
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Load input image
     image = PIL.Image.open(args.image).convert("RGB")
@@ -91,30 +131,54 @@ def main():
     # Resize image to target dimensions
     image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
 
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
-
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+    parallel_config = DiffusionParallelConfig(
+        cfg_parallel_size=args.cfg_parallel_size,
+    )
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        enable_layerwise_offload=args.enable_layerwise_offload,
+        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         boundary_ratio=args.boundary_ratio,
         flow_shift=args.flow_shift,
+        enable_cpu_offload=args.enable_cpu_offload,
+        parallel_config=parallel_config,
+        enforce_eager=args.enforce_eager,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
+    # Print generation configuration
+    print(f"\n{'=' * 60}")
+    print("Generation Configuration:")
+    print(f"  Model: {args.model}")
+    print(f"  Inference steps: {args.num_inference_steps}")
+    print(f"  Frames: {args.num_frames}")
+    print(f"  Parallel configuration: cfg_parallel_size={args.cfg_parallel_size}")
+    print(f"  Video size: {args.width}x{args.height}")
+    print(f"{'=' * 60}\n")
+
     # omni.generate() returns Generator[OmniRequestOutput, None, None]
     frames = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        pil_image=image,
-        height=height,
-        width=width,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        guidance_scale_2=args.guidance_scale_high,
-        num_inference_steps=args.num_inference_steps,
-        num_frames=args.num_frames,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+            "multi_modal_data": {"image": image},
+        },
+        OmniDiffusionSamplingParams(
+            height=height,
+            width=width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            guidance_scale_2=args.guidance_scale_high,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
     )
 
     # Extract video frames from OmniRequestOutput
@@ -178,6 +242,23 @@ def main():
     export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
 
 if __name__ == "__main__":
     main()

examples/offline_inference/lora_inference/README.md

@@ -0,0 +1,98 @@
+# LoRA Inference(Diffusion)
+
+This directory contains examples for using LoRA (Low-Rank Adaptation) adapters with vLLM-omni diffusion models for offline inference.
+The example uses the  `stabilityai/stable-diffusion-3.5-medium` as the default model, but you can replace it with other models in vLLM-omni.
+
+## Overview
+
+Similar to vLLM, vLLM-omni uses a unified LoRA handling mechanism:
+
+- **Pre-loaded LoRA**: Loaded at initialization via `--lora-path` (pre-loaded into cache)
+- **Per-request LoRA**: Loaded on-demand. In the example, the LoRA is loaded via `--lora-request-path` in each request
+
+Both approaches use the same underlying mechanism - all LoRA adapters are handled uniformly through `set_active_adapter()`. If no LoRA request is provided in a request, all adapters are deactivated.
+
+## Usage
+
+### Pre-loaded LoRA (via --lora-path)
+
+Load a LoRA adapter at initialization. This adapter is pre-loaded into the cache and can be activated by requests:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_preloaded.png
+```
+
+**Note**: When using `--lora-path`, the adapter is loaded at init time with a stable ID derived from the adapter path. This example activates it automatically for the request.
+
+### Per-request LoRA (via --lora-request-path)
+
+Load a LoRA adapter on-demand for each request:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --lora-request-path /path/to/lora/ \
+    --lora-scale 1.0 \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_per_request.png
+```
+
+### No LoRA
+
+If no LoRA request is provided, we will use the base model without any LoRA adapters:
+
+```bash
+python -m examples.offline_inference.lora_inference.lora_inference \
+    --prompt "A piece of cheesecake" \
+    --num_inference_steps 50 \
+    --height 1024 \
+    --width 1024 \
+    --output output_no_lora.png
+```
+
+## Parameters
+
+### LoRA Parameters
+
+- `--lora-path`: Path to LoRA adapter folder to pre-load at initialization (loads into cache with a stable ID derived from the path)
+- `--lora-request-path`: Path to LoRA adapter folder for per-request loading
+- `--lora-request-id`: Integer ID for the LoRA adapter (optional). If not provided and `--lora-request-path` is set, will derive a stable ID from the path.
+- `--lora-scale`: Scale factor for LoRA weights (default: 1.0). Higher values increase the influence of the LoRA adapter.
+
+### Standard Parameters
+
+- `--prompt`: Text prompt for image generation (required)
+- `--seed`: Random seed for reproducibility (default: 42)
+- `--height`: Image height in pixels (default: 1024)
+- `--width`: Image width in pixels (default: 1024)
+- `--num_inference_steps`: Number of denoising steps (default: 50)
+- `--output`: Output file path (default: `lora_output.png`)
+
+## How LoRA Works
+
+All LoRA adapters are handled uniformly:
+
+1. **Initialization**: If `--lora-path` is provided, the adapter is loaded into cache with a stable ID derived from the adapter path
+2. **Per-request**: If `--lora-request-path` is provided, the adapter is loaded/activated for that request
+3. **No LoRA**: If no LoRA request is provided (`req.lora_request` is None), all adapters are deactivated
+
+The system uses LRU cache management - adapters are cached and evicted when the cache is full (unless pinned).
+
+## LoRA Adapter Format
+
+LoRA adapters must be in PEFT (Parameter-Efficient Fine-Tuning) format. A typical LoRA adapter directory structure:
+
+```
+lora_adapter/
+├── adapter_config.json
+└── adapter_model.safetensors
+```

examples/offline_inference/lora_inference/lora_inference.py

@@ -0,0 +1,147 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import argparse
+from pathlib import Path
+
+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
+
+
+def parse_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(description="Generate images with LoRA adapters.")
+    parser.add_argument("--model", default="stabilityai/stable-diffusion-3.5-medium", help="Model name or path.")
+    parser.add_argument("--prompt", required=True, help="Text prompt for image generation.")
+    parser.add_argument("--seed", type=int, default=42, help="Random seed for deterministic results.")
+    parser.add_argument("--height", type=int, default=1024, help="Height of generated image.")
+    parser.add_argument("--width", type=int, default=1024, help="Width of generated image.")
+    parser.add_argument(
+        "--num_inference_steps",
+        type=int,
+        default=50,
+        help="Number of denoising steps for the diffusion sampler.",
+    )
+    parser.add_argument(
+        "--output",
+        type=str,
+        default="lora_output.png",
+        help="Path to save the generated image (PNG).",
+    )
+    parser.add_argument(
+        "--lora-path",
+        type=str,
+        default=None,
+        help="Path to LoRA adapter folder to pre-load at initialization (PEFT format). "
+        "Note: pre-loading populates the cache; you still need to pass a lora_request to activate it.",
+    )
+    parser.add_argument(
+        "--lora-request-path",
+        type=str,
+        default=None,
+        help="Path to LoRA adapter folder for per-request activation (dynamic LoRA). "
+        "If --lora-request-id is not provided, a stable ID will be derived from this path.",
+    )
+    parser.add_argument(
+        "--lora-request-id",
+        type=int,
+        default=None,
+        help="Integer ID for the LoRA adapter (for dynamic LoRA). "
+        "If not provided and --lora-request-path is set, will derive a stable ID from the path.",
+    )
+    parser.add_argument(
+        "--lora-scale",
+        type=float,
+        default=1.0,
+        help="Scale factor for LoRA weights (default: 1.0).",
+    )
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    model = args.model
+
+    omni_kwargs = {}
+
+    if args.lora_path:
+        omni_kwargs["lora_path"] = args.lora_path
+        print(f"Using static LoRA from: {args.lora_path}")
+
+    omni = Omni(model=model, **omni_kwargs)
+
+    lora_request = None
+    if args.lora_request_path:
+        if args.lora_request_id is None:
+            lora_request_id = stable_lora_int_id(args.lora_request_path)
+        else:
+            lora_request_id = args.lora_request_id
+
+        lora_name = Path(args.lora_request_path).stem
+        lora_request = LoRARequest(
+            lora_name=lora_name,
+            lora_int_id=lora_request_id,
+            lora_path=args.lora_request_path,
+        )
+        print(f"Using per-request LoRA: name={lora_name}, id={lora_request_id}, scale={args.lora_scale}")
+    elif args.lora_path:
+        # pre-loaded LoRA
+        lora_request_id = stable_lora_int_id(args.lora_path)
+        lora_request = LoRARequest(
+            lora_name="preloaded",
+            lora_int_id=lora_request_id,
+            lora_path=args.lora_path,
+        )
+        print(f"Activating pre-loaded LoRA: id={lora_request_id}, scale={args.lora_scale}")
+
+    sampling_params = OmniDiffusionSamplingParams(
+        height=args.height,
+        width=args.width,
+        num_inference_steps=args.num_inference_steps,
+    )
+
+    if lora_request:
+        sampling_params.lora_request = lora_request
+        sampling_params.lora_scale = args.lora_scale
+
+    outputs = omni.generate(args.prompt, sampling_params)
+
+    if not outputs or len(outputs) == 0:
+        raise ValueError("No output generated from omni.generate()")
+
+    if isinstance(outputs, list):
+        first_output = outputs[0]
+    else:
+        first_output = outputs
+
+    images = None
+    if hasattr(first_output, "images") and first_output.images:
+        images = first_output.images
+    elif hasattr(first_output, "request_output") and first_output.request_output:
+        req_out = first_output.request_output
+        if isinstance(req_out, list) and len(req_out) > 0:
+            req_out = req_out[0]
+        if hasattr(req_out, "images") and req_out.images:
+            images = req_out.images
+
+    if not images:
+        raise ValueError("No images found in request_output")
+
+    output_path = Path(args.output)
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    suffix = output_path.suffix or ".png"
+    stem = output_path.stem or "lora_output"
+    if len(images) <= 1:
+        images[0].save(output_path)
+        print(f"Saved generated image to {output_path}")
+    else:
+        for idx, img in enumerate(images):
+            save_path = output_path.parent / f"{stem}_{idx}{suffix}"
+            img.save(save_path)
+            print(f"Saved generated image to {save_path}")
+
+
+if __name__ == "__main__":
+    main()

examples/offline_inference/qwen2_5_omni/end2end.py

@@ -278,9 +278,9 @@ def get_audio_query(question: str = None, audio_path: str | None = None, samplin
 
 
 query_map = {
-    "mixed_modalities": get_mixed_modalities_query,
+    "use_mixed_modalities": get_mixed_modalities_query,
     "use_audio_in_video": get_use_audio_in_video_query,
-    "multi_audios": get_multi_audios_query,
+    "use_multi_audios": get_multi_audios_query,
     "use_image": get_image_query,
     "use_video": get_video_query,
     "use_audio": get_audio_query,
@@ -410,7 +410,7 @@ def main(args):
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
                 request_id = output.request_id
-                audio_tensor = output.multimodal_output["audio"]
+                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}")
@@ -434,7 +434,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="mixed_modalities",
+        default="use_mixed_modalities",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/offline_inference/qwen2_5_omni/run_single_prompt.sh

@@ -1,2 +1,2 @@
 python end2end.py --output-wav output_audio \
-                  --query-type mixed_modalities
+                  --query-type use_mixed_modalities

examples/offline_inference/qwen3_omni/end2end.py

@@ -12,6 +12,7 @@ from typing import NamedTuple
 import librosa
 import numpy as np
 import soundfile as sf
+import vllm
 from PIL import Image
 from vllm import SamplingParams
 from vllm.assets.audio import AudioAsset
@@ -225,18 +226,76 @@ def get_multi_audios_query() -> QueryResult:
     )
 
 
+# def get_use_audio_in_video_query(video_path: str | None = None) -> QueryResult:
+# question = (
+#     "Describe the content of the video in details, then convert what the "
+#     "baby say into text."
+# )
+# prompt = (
+#     f"<|im_start|>system\n{default_system}<|im_end|>\n"
+#     "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+#     f"{question}<|im_end|>\n"
+#     f"<|im_start|>assistant\n"
+# )
+# if video_path:
+#     if not os.path.exists(video_path):
+#         raise FileNotFoundError(f"Video file not found: {video_path}")
+#     video_frames = video_to_ndarrays(video_path, num_frames=16)
+# else:
+#     video_frames = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays
+# audio = extract_video_audio(video_path, sampling_rate=16000)
+# return QueryResult(
+#     inputs={
+#         "prompt": prompt,
+#         "multi_modal_data": {
+#             "video": video_frames,
+#             "audio": audio,
+#         },
+#         "mm_processor_kwargs": {
+#             "use_audio_in_video": True,
+#         },
+#     },
+#     limit_mm_per_prompt={"audio": 1, "video": 1},
+# )
+def get_use_audio_in_video_query() -> QueryResult:
+    question = "Describe the content of the video in details, then convert what the baby say into text."
+    prompt = (
+        f"<|im_start|>system\n{default_system}<|im_end|>\n"
+        "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>"
+        f"{question}<|im_end|>\n"
+        f"<|im_start|>assistant\n"
+    )
+    asset = VideoAsset(name="baby_reading", num_frames=16)
+    audio = asset.get_audio(sampling_rate=16000)
+    return QueryResult(
+        inputs={
+            "prompt": prompt,
+            "multi_modal_data": {
+                "video": asset.np_ndarrays,
+                "audio": audio,
+            },
+            "mm_processor_kwargs": {
+                "use_audio_in_video": True,
+            },
+        },
+        limit_mm_per_prompt={"audio": 1, "video": 1},
+    )
+
+
 query_map = {
     "text": get_text_query,
     "use_audio": get_audio_query,
     "use_image": get_image_query,
     "use_video": get_video_query,
-    "multi_audios": get_multi_audios_query,
-    "mixed_modalities": get_mixed_modalities_query,
+    "use_multi_audios": get_multi_audios_query,
+    "use_mixed_modalities": get_mixed_modalities_query,
+    "use_audio_in_video": get_use_audio_in_video_query,
 }
 
 
 def main(args):
     model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+    print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
 
     # Get paths from args
     video_path = getattr(args, "video_path", None)
@@ -259,6 +318,10 @@ def main(args):
             num_frames=getattr(args, "num_frames", 16),
             sampling_rate=getattr(args, "sampling_rate", 16000),
         )
+    elif args.query_type == "multi_audios":
+        query_result = query_func()
+    elif args.query_type == "use_audio_in_video":
+        query_result = query_func()
     else:
         query_result = query_func()
 
@@ -270,7 +333,7 @@ def main(args):
     )
 
     thinker_sampling_params = SamplingParams(
-        temperature=0.4,
+        temperature=0.9,
         top_p=0.9,
         top_k=-1,
         max_tokens=1200,
@@ -331,6 +394,8 @@ def main(args):
     total_requests = len(prompts)
     processed_count = 0
 
+    print(f"query type: {args.query_type}")
+
     for stage_outputs in omni_generator:
         if stage_outputs.final_output_type == "text":
             for output in stage_outputs.request_output:
@@ -353,7 +418,7 @@ def main(args):
         elif stage_outputs.final_output_type == "audio":
             for output in stage_outputs.request_output:
                 request_id = output.request_id
-                audio_tensor = output.multimodal_output["audio"]
+                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
@@ -385,7 +450,7 @@ def parse_args():
         "--query-type",
         "-q",
         type=str,
-        default="mixed_modalities",
+        default="use_mixed_modalities",
         choices=query_map.keys(),
         help="Query type.",
     )

examples/offline_inference/qwen3_tts/README.md

@@ -0,0 +1,93 @@
+# Qwen3-TTS
+
+This directory contains an offline demo for running Qwen3 TTS models with vLLM Omni. It builds task-specific inputs and generates WAV files locally.
+
+## Model Overview
+
+Qwen3 TTS provides multiple task variants for speech generation:
+
+- **CustomVoice**: Generate speech with a known speaker identity (speaker ID) and optional instruction.
+- **VoiceDesign**: Generate speech from text plus a descriptive instruction that designs a new voice.
+- **Base**: Voice cloning using a reference audio + reference transcript, with optional mode selection.
+
+## 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.
+
+### ROCm Dependencies
+
+You will need to install these two dependencies `onnxruntime-rocm` and `sox`.
+
+```
+pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm
+pip install onnxruntime-rocm sox
+```
+
+## Quick Start
+
+Run a single sample for a task:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Generated audio files are saved to `output_audio/` by default.
+
+## Task Usage
+
+### CustomVoice
+
+Single sample:
+
+```
+python end2end.py --query-type CustomVoice
+```
+
+Batch sample (multiple prompts in one run):
+
+```
+python end2end.py --query-type CustomVoice --use-batch-sample
+```
+
+### VoiceDesign
+
+Single sample:
+
+```
+python end2end.py --query-type VoiceDesign
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type VoiceDesign --use-batch-sample
+```
+
+### Base (Voice Clone)
+
+Single sample:
+
+```
+python end2end.py --query-type Base
+```
+
+Batch sample:
+
+```
+python end2end.py --query-type Base --use-batch-sample
+```
+
+Mode selection for Base:
+
+- `--mode-tag icl` (default): standard mode
+- `--mode-tag xvec_only`: enable `x_vector_only_mode` in the request
+
+Examples:
+
+```
+python end2end.py --query-type Base --mode-tag icl
+```
+
+## Notes
+
+- The script uses the model paths embedded in `end2end.py`. Update them if your local cache path differs.
+- Use `--output-dir` (preferred) or `--output-wav` to change the output folder.

examples/offline_inference/qwen3_tts/end2end.py

@@ -0,0 +1,381 @@
+"""Offline inference demo for Qwen3 TTS via vLLM Omni.
+
+Provides single and batch sample inputs for CustomVoice, VoiceDesign, and Base
+tasks, then runs Omni generation and saves output wav files.
+"""
+
+import os
+from typing import NamedTuple
+
+import soundfile as sf
+
+os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
+
+from vllm import SamplingParams
+from vllm.utils.argparse_utils import FlexibleArgumentParser
+
+from vllm_omni import Omni
+
+
+class QueryResult(NamedTuple):
+    """Container for a prepared Omni request."""
+
+    inputs: dict
+    model_name: str
+
+
+def get_custom_voice_query(use_batch_sample: bool = False) -> QueryResult:
+    """Build CustomVoice sample inputs.
+
+    Args:
+        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
+
+    Returns:
+        QueryResult with Omni inputs and the CustomVoice model path.
+    """
+    task_type = "CustomVoice"
+    if use_batch_sample:
+        texts = ["其实我真的有发现，我是一个特别善于观察别人情绪的人。", "She said she would be here by noon."]
+        instructs = ["", "Very happy."]
+        languages = ["Chinese", "English"]
+        speakers = ["Vivian", "Ryan"]
+        inputs = []
+        for text, instruct, language, speaker in zip(texts, instructs, languages, speakers):
+            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            inputs.append(
+                {
+                    "prompt": prompt,
+                    "additional_information": {
+                        "task_type": [task_type],
+                        "text": [text],
+                        "instruct": [instruct],
+                        "language": [language],
+                        "speaker": [speaker],
+                        "max_new_tokens": [2048],
+                    },
+                }
+            )
+    else:
+        text = "其实我真的有发现，我是一个特别善于观察别人情绪的人。"
+        language = "Chinese"
+        speaker = "Vivian"
+        instruct = "用特别愤怒的语气说"
+        prompts = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        inputs = {
+            "prompt": prompts,
+            "additional_information": {
+                "task_type": [task_type],
+                "text": [text],
+                "language": [language],
+                "speaker": [speaker],
+                "instruct": [instruct],
+                "max_new_tokens": [2048],
+            },
+        }
+    return QueryResult(
+        inputs=inputs,
+        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
+    )
+
+
+def get_voice_design_query(use_batch_sample: bool = False) -> QueryResult:
+    """Build VoiceDesign sample inputs.
+
+    Args:
+        use_batch_sample: When True, return a batch of prompts; otherwise a single prompt.
+
+    Returns:
+        QueryResult with Omni inputs and the VoiceDesign model path.
+    """
+    task_type = "VoiceDesign"
+    if use_batch_sample:
+        texts = [
+            "哥哥，你回来啦，人家等了你好久好久了，要抱抱！",
+            "It's in the top drawer... wait, it's empty? No way, that's impossible! I'm sure I put it there!",
+        ]
+        instructs = [
+            "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。",
+            "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice.",
+        ]
+        languages = ["Chinese", "English"]
+        inputs = []
+        for text, instruct, language in zip(texts, instructs, languages):
+            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            inputs.append(
+                {
+                    "prompt": prompt,
+                    "additional_information": {
+                        "task_type": [task_type],
+                        "text": [text],
+                        "language": [language],
+                        "instruct": [instruct],
+                        "max_new_tokens": [2048],
+                        "non_streaming_mode": [True],
+                    },
+                }
+            )
+    else:
+        text = "哥哥，你回来啦，人家等了你好久好久了，要抱抱！"
+        instruct = "体现撒娇稚嫩的萝莉女声，音调偏高且起伏明显，营造出黏人、做作又刻意卖萌的听觉效果。"
+        language = "Chinese"
+        prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+        inputs = {
+            "prompt": prompt,
+            "additional_information": {
+                "task_type": [task_type],
+                "text": [text],
+                "language": [language],
+                "instruct": [instruct],
+                "max_new_tokens": [2048],
+                "non_streaming_mode": [True],
+            },
+        }
+    return QueryResult(
+        inputs=inputs,
+        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign",
+    )
+
+
+def get_base_query(use_batch_sample: bool = False, mode_tag: str = "icl") -> QueryResult:
+    """Build Base (voice clone) sample inputs.
+
+    Args:
+        use_batch_sample: When True, return a batch of prompts (Case 2).
+        mode_tag: "icl" or "xvec_only" to control x_vector_only_mode behavior.
+
+    Returns:
+        QueryResult with Omni inputs and the Base model path.
+    """
+    task_type = "Base"
+    ref_audio_path_1 = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
+    ref_audio_single = ref_audio_path_1
+    ref_text_single = (
+        "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
+    )
+    syn_text_single = "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye."
+    syn_lang_single = "Auto"
+    x_vector_only_mode = mode_tag == "xvec_only"
+    if use_batch_sample:
+        syn_text_batch = [
+            "Good one. Okay, fine, I'm just gonna leave this sock monkey here. Goodbye.",
+            "其实我真的有发现，我是一个特别善于观察别人情绪的人。",
+        ]
+        syn_lang_batch = ["Chinese", "English"]
+        inputs = []
+        for text, language in zip(syn_text_batch, syn_lang_batch):
+            prompt = f"<|im_start|>assistant\n{text}<|im_end|>\n<|im_start|>assistant\n"
+            inputs.append(
+                {
+                    "prompt": prompt,
+                    "additional_information": {
+                        "task_type": [task_type],
+                        "ref_audio": [ref_audio_single],
+                        "ref_text": [ref_text_single],
+                        "text": [text],
+                        "language": [language],
+                        "x_vector_only_mode": [x_vector_only_mode],
+                        "max_new_tokens": [2048],
+                    },
+                }
+            )
+    else:
+        prompt = f"<|im_start|>assistant\n{syn_text_single}<|im_end|>\n<|im_start|>assistant\n"
+        inputs = {
+            "prompt": prompt,
+            "additional_information": {
+                "task_type": [task_type],
+                "ref_audio": [ref_audio_single],
+                "ref_text": [ref_text_single],
+                "text": [syn_text_single],
+                "language": [syn_lang_single],
+                "x_vector_only_mode": [x_vector_only_mode],
+                "max_new_tokens": [2048],
+            },
+        }
+    return QueryResult(
+        inputs=inputs,
+        model_name="Qwen/Qwen3-TTS-12Hz-1.7B-Base",
+    )
+
+
+def main(args):
+    """Run offline inference with Omni using prepared sample inputs.
+
+    Args:
+        args: Parsed CLI args from parse_args().
+    """
+    query_func = query_map[args.query_type]
+    if args.query_type in {"CustomVoice", "VoiceDesign"}:
+        query_result = query_func(use_batch_sample=args.use_batch_sample)
+    elif args.query_type == "Base":
+        query_result = query_func(
+            use_batch_sample=args.use_batch_sample,
+            mode_tag=args.mode_tag,
+        )
+    else:
+        query_result = query_func()
+
+    model_name = query_result.model_name
+    omni = Omni(
+        model=model_name,
+        stage_configs_path=args.stage_configs_path,
+        log_stats=args.enable_stats,
+        stage_init_timeout=args.stage_init_timeout,
+    )
+
+    sampling_params = SamplingParams(
+        temperature=0.9,
+        top_p=1.0,
+        top_k=50,
+        max_tokens=2048,
+        seed=42,
+        detokenize=False,
+        repetition_penalty=1.05,
+    )
+
+    sampling_params_list = [
+        sampling_params,
+    ]
+
+    output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
+    os.makedirs(output_dir, exist_ok=True)
+
+    omni_generator = omni.generate(query_result.inputs, sampling_params_list)
+    for stage_outputs in omni_generator:
+        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")
+            audio_samplerate = output.outputs[0].multimodal_output["sr"].item()
+            # 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()
+
+            # Save audio file with explicit WAV format
+            sf.write(output_wav, audio_numpy, samplerate=audio_samplerate, format="WAV")
+            print(f"Request ID: {request_id}, Saved audio to {output_wav}")
+
+
+def parse_args():
+    """Parse CLI arguments for offline TTS inference.
+
+    Returns:
+        argparse.Namespace with CLI options.
+    """
+    parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models")
+    parser.add_argument(
+        "--query-type",
+        "-q",
+        type=str,
+        default="CustomVoice",
+        choices=query_map.keys(),
+        help="Query type.",
+    )
+    parser.add_argument(
+        "--enable-stats",
+        action="store_true",
+        default=False,
+        help="Enable writing detailed statistics (default: disabled)",
+    )
+    parser.add_argument(
+        "--stage-init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing a single stage in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--batch-timeout",
+        type=int,
+        default=5,
+        help="Timeout for batching in seconds (default: 5)",
+    )
+    parser.add_argument(
+        "--init-timeout",
+        type=int,
+        default=300,
+        help="Timeout for initializing stages in seconds (default: 300)",
+    )
+    parser.add_argument(
+        "--shm-threshold-bytes",
+        type=int,
+        default=65536,
+        help="Threshold for using shared memory in bytes (default: 65536)",
+    )
+    parser.add_argument(
+        "--output-wav",
+        default="output_audio",
+        help="[Deprecated] Output wav directory (use --output-dir).",
+    )
+    parser.add_argument(
+        "--num-prompts",
+        type=int,
+        default=1,
+        help="Number of prompts to generate.",
+    )
+    parser.add_argument(
+        "--txt-prompts",
+        type=str,
+        default=None,
+        help="Path to a .txt file with one prompt per line (preferred).",
+    )
+    parser.add_argument(
+        "--stage-configs-path",
+        type=str,
+        default=None,
+        help="Path to a stage configs file.",
+    )
+    parser.add_argument(
+        "--audio-path",
+        "-a",
+        type=str,
+        default=None,
+        help="Path to local audio file. If not provided, uses default audio asset.",
+    )
+    parser.add_argument(
+        "--sampling-rate",
+        type=int,
+        default=16000,
+        help="Sampling rate for audio loading (default: 16000).",
+    )
+    parser.add_argument(
+        "--log-dir",
+        type=str,
+        default="logs",
+        help="Log directory (default: logs).",
+    )
+    parser.add_argument(
+        "--py-generator",
+        action="store_true",
+        default=False,
+        help="Use py_generator mode. The returned type of Omni.generate() is a Python Generator object.",
+    )
+    parser.add_argument(
+        "--use-batch-sample",
+        action="store_true",
+        default=False,
+        help="Use batch input sample for CustomVoice/VoiceDesign/Base query.",
+    )
+    parser.add_argument(
+        "--mode-tag",
+        type=str,
+        default="icl",
+        choices=["icl", "xvec_only"],
+        help="Mode tag for Base query x_vector_only_mode (default: icl).",
+    )
+
+    return parser.parse_args()
+
+
+query_map = {
+    "CustomVoice": get_custom_voice_query,
+    "VoiceDesign": get_voice_design_query,
+    "Base": get_base_query,
+}
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    main(args)

examples/offline_inference/text_to_audio/README.md

@@ -0,0 +1,37 @@
+# Text-To-Audio
+
+The `stabilityai/stable-audio-open-1.0` pipeline generates audio from text prompts.
+
+## Prerequisites
+
+If you use a gated model (e.g., `stabilityai/stable-audio-open-1.0`), ensure you have access:
+
+1. **Accept Model License**: Visit the model page on Hugging Face (e.g., [stabilityai/stable-audio-open-1.0]) and accept the user agreement.
+2. **Authenticate**: Log in to Hugging Face locally to access the gated model.
+   ```bash
+   huggingface-cli login
+   ```
+
+## Local CLI Usage
+
+```bash
+python text_to_audio.py \
+  --model stabilityai/stable-audio-open-1.0 \
+  --prompt "The sound of a hammer hitting a wooden surface" \
+  --negative_prompt "Low quality" \
+  --seed 42 \
+  --guidance_scale 7.0 \
+  --audio_length 10.0 \
+  --num_inference_steps 100 \
+  --output stable_audio_output.wav
+```
+
+Key arguments:
+
+- `--prompt`: text description (string).
+- `--negative_prompt`: negative prompt for classifier-free guidance.
+- `--seed`: integer seed for deterministic generation.
+- `--guidance_scale`: classifier-free guidance scale.
+- `--audio_length`: audio duration in seconds.
+- `--num_inference_steps`: diffusion sampling steps.(more steps = higher quality, slower).
+- `--output`: path to save the generated WAV file.

examples/offline_inference/text_to_audio/text_to_audio.py

@@ -21,7 +21,8 @@ import numpy as np
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.utils.platform_utils import detect_device_type
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -117,8 +118,7 @@ def save_audio(audio_data: np.ndarray, output_path: str, sample_rate: int = 4410
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     print(f"\n{'=' * 60}")
     print("Stable Audio Open - Text-to-Audio Generation")
@@ -142,17 +142,21 @@ def main():
     generation_start = time.perf_counter()
 
     # Generate audio
-    audio = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        generator=generator,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_waveforms,
-        extra={
-            "audio_start_in_s": args.audio_start,
-            "audio_end_in_s": audio_end_in_s,
+    outputs = omni.generate(
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
         },
+        OmniDiffusionSamplingParams(
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_waveforms,
+            extra_args={
+                "audio_start_in_s": args.audio_start,
+                "audio_end_in_s": audio_end_in_s,
+            },
+        ),
     )
 
     generation_end = time.perf_counter()
@@ -166,6 +170,21 @@ def main():
     suffix = output_path.suffix or ".wav"
     stem = output_path.stem or "stable_audio_output"
 
+    # Extract audio from omni.generate() outputs
+    if not outputs:
+        raise ValueError("No output generated from omni.generate()")
+
+    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]
+    if not hasattr(request_output, "multimodal_output"):
+        raise ValueError("No multimodal_output found in request_output")
+
+    audio = request_output.multimodal_output.get("audio")
+    if audio is None:
+        raise ValueError("No audio output found in request_output")
+
     # Handle different output formats
     if isinstance(audio, torch.Tensor):
         audio = audio.cpu().float().numpy()

examples/offline_inference/text_to_image/README.md

@@ -20,7 +20,7 @@ if __name__ == "__main__":
     images[0].save("coffee.png")
 ```
 
-Or put more than one prompt in a request, processing them sequentially.
+Or put more than one prompt in a request.
 
 ```python
 from vllm_omni.entrypoints.omni import Omni
@@ -37,6 +37,41 @@ if __name__ == "__main__":
       image = output.request_output[0].images[0].save(f"{i}.jpg")
 ```
 
+!!! info
+
+    However, it is not currently recommended to do so
+    because not all models support batch inference,
+    and batch requesting mostly does not provide significant performance improvement (despite the impression that it does).
+    This feature is primarily for the sake of interface compatibility with vLLM and to allow for future improvements.
+
+!!! info
+
+    For diffusion pipelines, the stage config field `stage_args.[].runtime.max_batch_size` is 1 by default, and the input
+    list is sliced into single-item requests before feeding into the diffusion pipeline. For models that do internally support
+    batched inputs, you can [modify this configuration](../../../configuration/stage_configs.md) to let the model accept a longer batch of prompts.
+
+Apart from string prompt, vLLM-Omni also supports dictionary prompts in the same style as vLLM.
+This is useful for models that support negative prompts.
+
+```python
+from vllm_omni.entrypoints.omni import Omni
+
+if __name__ == "__main__":
+    omni = Omni(model="Qwen/Qwen-Image")
+    outputs = omni.generate([
+      {
+        "prompt": "a cup of coffee on a table"，
+        "negative_prompt": "low resolution"
+      },
+      {
+        "prompt": "a toy dinosaur on a sandy beach"，
+        "negative_prompt": "cinematic, realistic"
+      }
+    ])
+    for i, output in enumerate(outputs):
+      image = output.request_output[0].images[0].save(f"{i}.jpg")
+```
+
 ## Local CLI Usage
 
 ```bash
@@ -61,6 +96,12 @@ Key arguments:
 - `--num_inference_steps`: diffusion sampling steps (more steps = higher quality, slower).
 - `--height/--width`: output resolution (defaults 1024x1024).
 - `--output`: path to save the generated PNG.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.
 
 > ℹ️ Qwen-Image currently publishes best-effort presets at `1328x1328`, `1664x928`, `928x1664`, `1472x1140`, `1140x1472`, `1584x1056`, and `1056x1584`. Adjust `--height/--width` accordingly for the most reliable outcomes.
 

examples/offline_inference/text_to_image/gradio_demo.py

@@ -5,8 +5,9 @@ import gradio as gr
 import torch
 
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 ASPECT_RATIOS: dict[str, tuple[int, int]] = {
     "1:1": (1328, 1328),
@@ -61,8 +62,8 @@ def parse_args() -> argparse.Namespace:
 @lru_cache(maxsize=1)
 def get_omni(model_name: str) -> Omni:
     # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    vae_use_slicing = current_omni_platform.is_npu()
+    vae_use_tiling = current_omni_platform.is_npu()
     return Omni(
         model=model_name,
         vae_use_slicing=vae_use_slicing,
@@ -71,7 +72,6 @@ def get_omni(model_name: str) -> Omni:
 
 
 def build_demo(args: argparse.Namespace) -> gr.Blocks:
-    device = detect_device_type()
     omni = get_omni(args.model)
 
     def run_inference(
@@ -98,15 +98,17 @@ def build_demo(args: argparse.Namespace) -> gr.Blocks:
             raise gr.Error("Inference steps must be a positive integer.")
         if num_images not in {1, 2, 3, 4}:
             raise gr.Error("Number of images must be 1, 2, 3, or 4.")
-        generator = torch.Generator(device=device).manual_seed(seed)
+        generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
         outputs = omni.generate(
             prompt.strip(),
-            height=height,
-            width=width,
-            generator=generator,
-            true_cfg_scale=float(cfg_scale_value),
-            num_inference_steps=num_steps,
-            num_outputs_per_prompt=num_images,
+            OmniDiffusionSamplingParams(
+                height=height,
+                width=width,
+                generator=generator,
+                true_cfg_scale=float(cfg_scale_value),
+                num_inference_steps=num_steps,
+                num_outputs_per_prompt=num_images,
+            ),
         )
         images_outputs = []
         for output in outputs:

examples/offline_inference/text_to_image/text_to_image.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 
 import argparse
+import os
 import time
 from pathlib import Path
 
@@ -9,8 +10,9 @@ import torch
 
 from vllm_omni.diffusion.data import DiffusionParallelConfig, logger
 from vllm_omni.entrypoints.omni import Omni
+from vllm_omni.inputs.data import OmniDiffusionSamplingParams
 from vllm_omni.outputs import OmniRequestOutput
-from vllm_omni.utils.platform_utils import detect_device_type, is_npu
+from vllm_omni.platforms import current_omni_platform
 
 
 def parse_args() -> argparse.Namespace:
@@ -23,7 +25,9 @@ def parse_args() -> argparse.Namespace:
     )
     parser.add_argument("--prompt", default="a cup of coffee on the table", help="Text prompt for image generation.")
     parser.add_argument(
-        "--negative_prompt", default="", help="negative prompt for classifier-free conditional guidance."
+        "--negative_prompt",
+        default=None,
+        help="negative prompt for classifier-free conditional guidance.",
     )
     parser.add_argument("--seed", type=int, default=142, help="Random seed for deterministic results.")
     parser.add_argument(
@@ -69,6 +73,11 @@ def parse_args() -> argparse.Namespace:
             "Default: None (no cache acceleration)."
         ),
     )
+    parser.add_argument(
+        "--enable-cache-dit-summary",
+        action="store_true",
+        help="Enable cache-dit summary logging after diffusion forward passes.",
+    )
     parser.add_argument(
         "--ulysses_degree",
         type=int,
@@ -93,23 +102,44 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Disable torch.compile and force eager execution.",
     )
+    parser.add_argument(
+        "--enable-cpu-offload",
+        action="store_true",
+        help="Enable CPU offloading for diffusion models.",
+    )
+    parser.add_argument(
+        "--enable-layerwise-offload",
+        action="store_true",
+        help="Enable layerwise (blockwise) offloading on DiT modules.",
+    )
+    parser.add_argument(
+        "--layerwise-num-gpu-layers",
+        type=int,
+        default=1,
+        help="Number of ready layers (blocks) to keep on GPU during generation.",
+    )
     parser.add_argument(
         "--tensor_parallel_size",
         type=int,
         default=1,
         help="Number of GPUs used for tensor parallelism (TP) inside the DiT.",
     )
+    parser.add_argument(
+        "--vae_use_slicing",
+        action="store_true",
+        help="Enable VAE slicing for memory optimization.",
+    )
+    parser.add_argument(
+        "--vae_use_tiling",
+        action="store_true",
+        help="Enable VAE tiling for memory optimization.",
+    )
     return parser.parse_args()
 
 
 def main():
     args = parse_args()
-    device = detect_device_type()
-    generator = torch.Generator(device=device).manual_seed(args.seed)
-
-    # Enable VAE memory optimizations on NPU
-    vae_use_slicing = is_npu()
-    vae_use_tiling = is_npu()
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
 
     # Configure cache based on backend type
     cache_config = None
@@ -148,16 +178,27 @@ def main():
         tensor_parallel_size=args.tensor_parallel_size,
     )
 
+    # Check if profiling is requested via environment variable
+    profiler_enabled = bool(os.getenv("VLLM_TORCH_PROFILER_DIR"))
+
     omni = Omni(
         model=args.model,
-        vae_use_slicing=vae_use_slicing,
-        vae_use_tiling=vae_use_tiling,
+        enable_layerwise_offload=args.enable_layerwise_offload,
+        layerwise_num_gpu_layers=args.layerwise_num_gpu_layers,
+        vae_use_slicing=args.vae_use_slicing,
+        vae_use_tiling=args.vae_use_tiling,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
+        enable_cache_dit_summary=args.enable_cache_dit_summary,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
+        enable_cpu_offload=args.enable_cpu_offload,
     )
 
+    if profiler_enabled:
+        print("[Profiler] Starting profiling...")
+        omni.start_profile()
+
     # Time profiling for generation
     print(f"\n{'=' * 60}")
     print("Generation Configuration:")
@@ -173,15 +214,19 @@ def main():
 
     generation_start = time.perf_counter()
     outputs = omni.generate(
-        args.prompt,
-        negative_prompt=args.negative_prompt,
-        height=args.height,
-        width=args.width,
-        generator=generator,
-        true_cfg_scale=args.cfg_scale,
-        guidance_scale=args.guidance_scale,
-        num_inference_steps=args.num_inference_steps,
-        num_outputs_per_prompt=args.num_images_per_prompt,
+        {
+            "prompt": args.prompt,
+            "negative_prompt": args.negative_prompt,
+        },
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            true_cfg_scale=args.cfg_scale,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_outputs_per_prompt=args.num_images_per_prompt,
+        ),
     )
     generation_end = time.perf_counter()
     generation_time = generation_end - generation_start
@@ -189,6 +234,23 @@ def main():
     # Print profiling results
     print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
 
+    if profiler_enabled:
+        print("\n[Profiler] Stopping profiler and collecting results...")
+        profile_results = omni.stop_profile()
+        if profile_results and isinstance(profile_results, dict):
+            traces = profile_results.get("traces", [])
+            print("\n" + "=" * 60)
+            print("PROFILING RESULTS:")
+            for rank, trace in enumerate(traces):
+                print(f"\nRank {rank}:")
+                if trace:
+                    print(f"  • Trace: {trace}")
+            if not traces:
+                print("  No traces collected.")
+            print("=" * 60)
+        else:
+            print("[Profiler] No valid profiling data returned.")
+
     # Extract images from OmniRequestOutput
     # omni.generate() returns list[OmniRequestOutput], extract images from the first output
     if not outputs or len(outputs) == 0:

examples/offline_inference/text_to_video/text_to_video.md

@@ -9,10 +9,11 @@ python text_to_video.py \
   --prompt "Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
   --negative_prompt "<optional quality filter>" \
   --height 480 \
-  --width 640 \
-  --num_frames 32 \
+  --width 832 \
+  --num_frames 33 \
   --guidance_scale 4.0 \
   --guidance_scale_high 3.0 \
+  --flow_shift 12.0 \
   --num_inference_steps 40 \
   --fps 16 \
   --output t2v_out.mp4
@@ -21,10 +22,16 @@ python text_to_video.py \
 Key arguments:
 
 - `--prompt`: text description (string).
-- `--height/--width`: output resolution (defaults 720x1280). Dimensions should align with Wan VAE downsampling (multiples of 8).
+- `--height/--width`: output resolution (defaults 480x832, i.e. 480P). Dimensions should align with Wan VAE downsampling (multiples of 8).
 - `--num_frames`: Number of frames (Wan default is 81).
-- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high)..
+- `--guidance_scale` and `--guidance_scale_high`: CFG scale (applied to low/high).
 - `--negative_prompt`: optional list of artifacts to suppress (the PR demo used a long Chinese string).
-- `--boundary_ratio`: Boundary split ratio for low/high DiT.
+- `--boundary_ratio`: Boundary split ratio for low/high DiT. Default `0.875` uses both transformers for best quality. Set to `1.0` to load only the low-noise transformer (saves noticeable memory with good quality, recommended if memory is limited). Set to `0.0` loads only the high-noise transformer (not recommended, lower quality).
 - `--fps`: frames per second for the saved MP4 (requires `diffusers` export_to_video).
 - `--output`: path to save the generated video.
+- `--vae_use_slicing`: enable VAE slicing for memory optimization.
+- `--vae_use_tiling`: enable VAE tiling for memory optimization.
+- `--cfg_parallel_size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
+- `--enable-cpu-offload`: enable CPU offloading for diffusion models.
+
+> ℹ️ If you encounter OOM errors, try using `--vae_use_slicing` and `--vae_use_tiling` to reduce memory usage.