cherry-pick pr3067 (#3087)

Signed-off-by: lyj-jjj <liuyingjun5@huawei.com>

.buildkite/bootstrap-amd-omni.sh

@@ -0,0 +1,246 @@
+#!/bin/bash
+# vllm-omni customized version
+# Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/bootstrap-amd.sh
+# Last synced: 2025-12-15
+# Modifications: Use local template file instead of downloading from ci-infra
+
+set -euo pipefail
+
+if [[ -z "${RUN_ALL:-}" ]]; then
+    RUN_ALL=0
+fi
+
+if [[ -z "${NIGHTLY:-}" ]]; then
+    NIGHTLY=0
+fi
+
+if [[ -z "${VLLM_CI_BRANCH:-}" ]]; then
+    VLLM_CI_BRANCH="main"
+fi
+
+if [[ -z "${AMD_MIRROR_HW:-}" ]]; then
+    AMD_MIRROR_HW="amdproduction"
+fi
+
+if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
+    DOCS_ONLY_DISABLE=0
+fi
+
+fail_fast() {
+    DISABLE_LABEL="ci-no-fail-fast"
+    # If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq
+    if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
+        PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name')
+        if [[ $PR_LABELS == *"$DISABLE_LABEL"* ]]; then
+            echo false
+        else
+            echo true
+        fi
+    else
+        echo false  # not a PR or BUILDKITE_PULL_REQUEST not set
+    fi
+}
+
+check_run_all_label() {
+    RUN_ALL_LABEL="ready-run-all-tests"
+    # If BUILDKITE_PULL_REQUEST != "false", then we check the PR labels using curl and jq
+    if [ "$BUILDKITE_PULL_REQUEST" != "false" ]; then
+        PR_LABELS=$(curl -s "https://api.github.com/repos/vllm-project/vllm-omni/pulls/$BUILDKITE_PULL_REQUEST" | jq -r '.labels[].name')
+        if [[ $PR_LABELS == *"$RUN_ALL_LABEL"* ]]; then
+            echo true
+        else
+            echo false
+        fi
+    else
+        echo false  # not a PR or BUILDKITE_PULL_REQUEST not set
+    fi
+}
+
+if [[ -z "${COV_ENABLED:-}" ]]; then
+    COV_ENABLED=0
+fi
+
+upload_pipeline() {
+    echo "Uploading pipeline..."
+    # Install minijinja
+    ls .buildkite || buildkite-agent annotate --style error 'Please merge upstream main branch for buildkite CI'
+    curl -sSfL https://github.com/mitsuhiko/minijinja/releases/download/2.3.1/minijinja-cli-installer.sh | sh
+    source /var/lib/buildkite-agent/.cargo/env
+
+    if [[ $BUILDKITE_PIPELINE_SLUG == "fastcheck" ]]; then
+        AMD_MIRROR_HW="amdtentative"
+    fi
+
+    # Use local template file for vllm-omni
+    cp .buildkite/test-template-amd-omni.j2 .buildkite/test-template.j2
+
+
+    # (WIP) Use pipeline generator instead of jinja template
+    if [ -e ".buildkite/pipeline_generator/pipeline_generator.py" ]; then
+        python -m pip install click pydantic
+        python .buildkite/pipeline_generator/pipeline_generator.py --run_all=$RUN_ALL --list_file_diff="$LIST_FILE_DIFF" --nightly="$NIGHTLY" --mirror_hw="$AMD_MIRROR_HW"
+        buildkite-agent pipeline upload .buildkite/pipeline.yaml
+        exit 0
+    fi
+    echo "List file diff: $LIST_FILE_DIFF"
+    echo "Run all: $RUN_ALL"
+    echo "Nightly: $NIGHTLY"
+    echo "AMD Mirror HW: $AMD_MIRROR_HW"
+
+    FAIL_FAST=$(fail_fast)
+
+    cd .buildkite
+
+    # Select test definition file: merge suite for main, ready suite for PRs
+    if [[ $BUILDKITE_BRANCH == "main" ]]; then
+        TEST_YAML="test-amd-merge.yml"
+    else
+        TEST_YAML="test-amd-ready.yaml"
+    fi
+
+    (
+        set -x
+        # Output pipeline.yaml with all blank lines removed
+        minijinja-cli test-template.j2 "$TEST_YAML" \
+            -D branch="$BUILDKITE_BRANCH" \
+            -D list_file_diff="$LIST_FILE_DIFF" \
+            -D run_all="$RUN_ALL" \
+            -D nightly="$NIGHTLY" \
+            -D mirror_hw="$AMD_MIRROR_HW" \
+            -D fail_fast="$FAIL_FAST" \
+            -D vllm_use_precompiled="$VLLM_USE_PRECOMPILED" \
+            -D vllm_merge_base_commit="$(git merge-base origin/main HEAD)" \
+            -D cov_enabled="$COV_ENABLED" \
+            -D vllm_ci_branch="$VLLM_CI_BRANCH" \
+            | sed '/^[[:space:]]*$/d' \
+            > pipeline.yaml
+    )
+    cat pipeline.yaml
+    buildkite-agent artifact upload pipeline.yaml
+    buildkite-agent pipeline upload pipeline.yaml
+    exit 0
+}
+
+get_diff() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
+}
+
+get_diff_main() {
+    $(git add .)
+    echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
+}
+
+file_diff=$(get_diff)
+if [[ $BUILDKITE_BRANCH == "main" ]]; then
+    file_diff=$(get_diff_main)
+fi
+
+# ----------------------------------------------------------------------
+# Early exit start: skip pipeline if conditions are met
+# ----------------------------------------------------------------------
+
+# skip pipeline if all changed files are under docs/
+if [[ "${DOCS_ONLY_DISABLE}" != "1" ]]; then
+  if [[ -n "${file_diff:-}" ]]; then
+    docs_only=1
+    # Robust iteration over newline-separated file_diff
+    while IFS= read -r f; do
+      [[ -z "$f" ]] && continue
+      # **Policy:** only skip if *every* path starts with docs/
+      if [[ "$f" != docs/* ]]; then
+        docs_only=0
+        break
+      fi
+    done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
+
+    if [[ "$docs_only" -eq 1 ]]; then
+      buildkite-agent annotate ":memo: CI skipped — docs/** only changes detected
+
+\`\`\`
+${file_diff}
+\`\`\`" --style "info" || true
+      echo "[docs-only] All changes are under docs/. Exiting before pipeline upload."
+      exit 0
+    fi
+  fi
+fi
+
+# ----------------------------------------------------------------------
+# Early exit end
+# ----------------------------------------------------------------------
+
+patterns=(
+    "docker/Dockerfile"
+    "CMakeLists.txt"
+    "requirements/common.txt"
+    "requirements/cuda.txt"
+    "requirements/build.txt"
+    "requirements/test.txt"
+    "setup.py"
+    "csrc/"
+    "cmake/"
+)
+
+ignore_patterns=(
+    "docker/Dockerfile."
+    "csrc/cpu"
+    "csrc/rocm"
+    "cmake/hipify.py"
+    "cmake/cpu_extension.cmake"
+)
+
+for file in $file_diff; do
+    # First check if file matches any pattern
+    matches_pattern=0
+    for pattern in "${patterns[@]}"; do
+        if [[ $file == $pattern* ]] || [[ $file == $pattern ]]; then
+            matches_pattern=1
+            break
+        fi
+    done
+
+    # If file matches pattern, check it's not in ignore patterns
+    if [[ $matches_pattern -eq 1 ]]; then
+        matches_ignore=0
+        for ignore in "${ignore_patterns[@]}"; do
+            if [[ $file == $ignore* ]] || [[ $file == $ignore ]]; then
+                matches_ignore=1
+                break
+            fi
+        done
+
+        if [[ $matches_ignore -eq 0 ]]; then
+            RUN_ALL=1
+            echo "Found changes: $file. Run all tests"
+            break
+        fi
+    fi
+done
+
+# Check for ready-run-all-tests label
+LABEL_RUN_ALL=$(check_run_all_label)
+if [[ $LABEL_RUN_ALL == true ]]; then
+    RUN_ALL=1
+    NIGHTLY=1
+    echo "Found 'ready-run-all-tests' label. Running all tests including optional tests."
+fi
+
+# Decide whether to use precompiled wheels
+# Relies on existing patterns array as a basis.
+if [[ -n "${VLLM_USE_PRECOMPILED:-}" ]]; then
+    echo "VLLM_USE_PRECOMPILED is already set to: $VLLM_USE_PRECOMPILED"
+elif [[ $RUN_ALL -eq 1 ]]; then
+    export VLLM_USE_PRECOMPILED=0
+    echo "Detected critical changes, building wheels from source"
+else
+    export VLLM_USE_PRECOMPILED=1
+    echo "No critical changes, using precompiled wheels"
+fi
+
+
+LIST_FILE_DIFF=$(get_diff | tr ' ' '|')
+if [[ $BUILDKITE_BRANCH == "main" ]]; then
+    LIST_FILE_DIFF=$(get_diff_main | tr ' ' '|')
+fi
+upload_pipeline

.buildkite/bootstrap-intel-omni.sh

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

.buildkite/pipeline-intel.yaml

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

.buildkite/pipeline.yml

@@ -1,6 +1,40 @@
 steps:
-  - label: "Simple Unit Test"
+  - label: ":docker: Build image"
+    key: image-build
     commands:
-      - ".buildkite/scripts/simple_test.sh"
+      - "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
+      - "docker build --progress=plain --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"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    if: build.branch != "main" && build.pull_request.labels includes "ready"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    if: build.branch == "main" && build.env("NIGHTLY") != "1"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
     agents:
       queue: "cpu_queue_premerge"

.buildkite/release-pipeline.yaml

@@ -0,0 +1,38 @@
+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 --progress=plain --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"
+
+  # L2 Test
+  - label: "Upload Ready Pipeline"
+    depends_on: image-build
+    key: upload-ready-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-ready.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L3 Test
+  - label: "Upload Merge Pipeline"
+    depends_on: image-build
+    key: upload-merge-pipeline
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-merge.yml
+    agents:
+      queue: "cpu_queue_premerge"
+
+  # L4 Test
+  - label: "Upload Nightly Pipeline"
+    depends_on: image-build
+    key: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - buildkite-agent pipeline upload .buildkite/test-nightly.yml
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/scripts/docker_login_ecr_public.sh

@@ -0,0 +1,73 @@
+#!/bin/bash
+# Helper function to safely login to ECR Public with per-job config isolation
+# Uses DOCKER_CONFIG environment variable to prevent race conditions
+#
+# This script prevents the "device or resource busy" error by giving each
+# Buildkite job its own isolated Docker config directory.
+#
+# Usage:
+#   source docker_login_ecr_public.sh && safe_docker_login_ecr_public
+
+set -euo pipefail
+
+# Configuration
+ECR_REGISTRY="public.ecr.aws"
+
+setup_isolated_docker_config() {
+    # Use BUILDKITE_JOB_ID for job-specific isolation
+    # Fallback to PID if running outside Buildkite
+    local job_id="${BUILDKITE_JOB_ID:-$$}"
+
+    # Set Docker config to job-specific directory
+    export DOCKER_CONFIG="/tmp/docker-config-${job_id}"
+
+    # Create directory if it doesn't exist
+    mkdir -p "$DOCKER_CONFIG"
+
+    echo "[docker-config] Using isolated Docker config: $DOCKER_CONFIG"
+}
+
+check_docker_auth() {
+    # Check if already authenticated to the given registry
+    # Returns 0 if authenticated, 1 if not
+    local registry="$1"
+
+    # Check if credentials exist in the isolated config
+    if [[ -f "$DOCKER_CONFIG/config.json" ]]; then
+        # Check if registry is present in config
+        if grep -q "$registry" "$DOCKER_CONFIG/config.json" 2>/dev/null; then
+            return 0
+        fi
+    fi
+
+    return 1
+}
+
+safe_docker_login_ecr_public() {
+    # Setup isolated config first
+    setup_isolated_docker_config
+
+    local registry="$ECR_REGISTRY"
+
+    # Check if already authenticated (within this job)
+    if check_docker_auth "$registry"; then
+        echo "[docker-login] Already authenticated to $registry in this job"
+        return 0
+    fi
+
+    # Perform login to isolated config directory
+    echo "[docker-login] Logging in to $ECR_REGISTRY (isolated config)..."
+    if aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin "$ECR_REGISTRY"; then
+        echo "[docker-login] Login successful (config: $DOCKER_CONFIG)"
+        return 0
+    else
+        local exit_code=$?
+        echo "[docker-login] ERROR: Login failed with exit code $exit_code" >&2
+        return $exit_code
+    fi
+}
+
+# Execute if run as script (not sourced)
+if [[ "${BASH_SOURCE[0]}" == "${0}" ]]; then
+    safe_docker_login_ecr_public
+fi

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

@@ -0,0 +1,171 @@
+#!/bin/bash
+# vllm-omni customized version
+# Based on: vllm/.buildkite/scripts/hardware_ci/run-amd-test.sh
+# Last synced: 2025-12-15
+# Modifications: docker image name for vllm-omni
+
+# This script runs test inside the corresponding ROCm docker container.
+set -o pipefail
+
+# Export Python path
+export PYTHONPATH=".."
+
+# Print ROCm version
+echo "--- Confirming Clean Initial State"
+while true; do
+        sleep 3
+        if grep -q clean /opt/amdgpu/etc/gpu_state; then
+                echo "GPUs state is \"clean\""
+                break
+        fi
+done
+
+echo "--- ROCm info"
+rocminfo
+
+# cleanup older docker images
+cleanup_docker() {
+  # Get Docker's root directory
+  docker_root=$(docker info -f '{{.DockerRootDir}}')
+  if [ -z "$docker_root" ]; then
+    echo "Failed to determine Docker root directory."
+    exit 1
+  fi
+  echo "Docker root directory: $docker_root"
+  # Check disk usage of the filesystem where Docker's root directory is located
+  disk_usage=$(df "$docker_root" | tail -1 | awk '{print $5}' | sed 's/%//')
+  # Define the threshold
+  threshold=70
+  if [ "$disk_usage" -gt "$threshold" ]; then
+    echo "Disk usage is above $threshold%. Cleaning up Docker images and volumes..."
+    # Remove dangling images (those that are not tagged and not used by any container)
+    docker image prune -f
+    # Remove unused volumes / force the system prune for old images as well.
+    docker volume prune -f && docker system prune --force --filter "until=72h" --all
+    echo "Docker images and volumes cleanup completed."
+  else
+    echo "Disk usage is below $threshold%. No cleanup needed."
+  fi
+}
+
+# Call the cleanup docker function
+cleanup_docker
+
+echo "--- Resetting GPUs"
+
+echo "reset" > /opt/amdgpu/etc/gpu_state
+
+while true; do
+        sleep 3
+        if grep -q clean /opt/amdgpu/etc/gpu_state; then
+                echo "GPUs state is \"clean\""
+                break
+        fi
+done
+
+echo "--- Pulling container"
+image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
+container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+
+# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
+sudo apt-get update && sudo apt-get install -y awscli
+# Use safe docker login helper to prevent race conditions
+source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
+safe_docker_login_ecr_public
+# Pull the container from ECR Public Gallery
+
+docker pull "${image_name}"
+
+remove_docker_container() {
+   docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
+}
+trap remove_docker_container EXIT
+
+echo "--- Running container"
+
+HF_CACHE="$(realpath ~)/huggingface"
+mkdir -p "${HF_CACHE}"
+HF_MOUNT="/root/.cache/huggingface"
+
+if [[ -n "${TEST_COMMAND:-}" ]]; then
+    commands="$TEST_COMMAND"
+else
+    commands="$@"
+fi
+echo "Commands:$commands"
+
+PARALLEL_JOB_COUNT=8
+MYPYTHONPATH=".."
+
+# Test that we're launching on the machine that has
+# proper access to GPUs
+render_gid=$(getent group render | cut -d: -f3)
+if [[ -z "$render_gid" ]]; then
+  echo "Error: 'render' group not found. This is required for GPU access." >&2
+  exit 1
+fi
+
+# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
+# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
+if [[ $commands == *"--shard-id="* ]]; then
+  # assign job count as the number of shards used
+  commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
+  for GPU in $(seq 0 $(($PARALLEL_JOB_COUNT-1))); do
+    # assign shard-id for each shard
+    commands_gpu=$(echo "$commands" | sed -E "s/--shard-id[[:blank:]]*=[[:blank:]]*[0-9]*/--shard-id=${GPU} /g" | sed 's/ \\ / /g')
+    echo "Shard ${GPU} commands:$commands_gpu"
+    echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
+    docker run \
+        --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
+        --network=host \
+        --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=0 \
+        -e HIP_VISIBLE_DEVICES="${GPU}" \
+        -e HF_TOKEN \
+        -e AWS_ACCESS_KEY_ID \
+        -e AWS_SECRET_ACCESS_KEY \
+        -v "${HF_CACHE}:${HF_MOUNT}" \
+        -e "HF_HOME=${HF_MOUNT}" \
+        -e "PYTHONPATH=${MYPYTHONPATH}" \
+        --name "${container_name}_${GPU}" \
+        "${image_name}" \
+        /bin/bash -c "${commands_gpu}" \
+        |& while read -r line; do echo ">>Shard $GPU: $line"; done &
+    PIDS+=($!)
+  done
+  #wait for all processes to finish and collect exit codes
+  for pid in "${PIDS[@]}"; do
+    wait "${pid}"
+    STATUS+=($?)
+  done
+  for st in "${STATUS[@]}"; do
+    if [[ ${st} -ne 0 ]]; then
+      echo "One of the processes failed with $st"
+      exit "${st}"
+    fi
+  done
+else
+  echo "Render devices: $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES"
+  docker run \
+          --device /dev/kfd $BUILDKITE_AGENT_META_DATA_RENDER_DEVICES \
+          --network=host \
+          --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=0 \
+          -e HF_TOKEN \
+          -e AWS_ACCESS_KEY_ID \
+          -e AWS_SECRET_ACCESS_KEY \
+          -v "${HF_CACHE}:${HF_MOUNT}" \
+          -e "HF_HOME=${HF_MOUNT}" \
+          -e "PYTHONPATH=${MYPYTHONPATH}" \
+          --name "${container_name}" \
+          "${image_name}" \
+          /bin/bash -c "${commands}"
+fi

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

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

.buildkite/scripts/hardware_ci/run_npu_test.sh

@@ -0,0 +1,145 @@
+#!/bin/bash
+
+# This script build the Ascend NPU docker image and run the offline inference inside the container.
+# It serves a sanity check for compilation and basic model usage.
+set -ex
+
+# Base ubuntu image with basic ascend development libraries and python installed
+VLLM_OMNI_REPO="https://github.com/vllm-project/vllm-omni.git"
+BASE_IMAGE_NAME="quay.nju.edu.cn/ascend/vllm-ascend:v0.11.0rc2"
+image_name="npu/vllm-omni-ci:${BUILDKITE_COMMIT}_${EPOCHSECONDS}"
+# image_name="npu/vllm-ci:${BUILDKITE_COMMIT}_${EPOCHSECONDS}"
+container_name="npu_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
+
+# BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards
+agent_idx=$(echo "${BUILDKITE_AGENT_NAME}" | awk -F'-' '{print $(NF-1)}')
+echo "agent_idx: ${agent_idx}"
+builder_name="cachebuilder${agent_idx}"
+builder_cache_dir="/mnt/docker-cache${agent_idx}"
+mkdir -p ${builder_cache_dir}
+
+# Try building the docker image
+cat <<EOF | DOCKER_BUILDKIT=1 docker build \
+    --add-host pypi-cache:${PYPI_CACHE_HOST} \
+    --builder ${builder_name} --cache-from type=local,src=${builder_cache_dir} \
+                           --cache-to type=local,dest=${builder_cache_dir},mode=max \
+    --build-arg BUILDKITE_PULL_REQUEST="${BUILDKITE_PULL_REQUEST}" \
+    --build-arg BUILDKITE_PULL_REQUEST_REPO="${BUILDKITE_PULL_REQUEST_REPO}" \
+    --progress=plain --load -t ${image_name} -f - .
+FROM ${BASE_IMAGE_NAME}
+
+# Define environments
+ENV DEBIAN_FRONTEND=noninteractive
+
+RUN pip config set global.index-url http://pypi-cache:${PYPI_CACHE_PORT}/pypi/simple && \
+    pip config set global.trusted-host pypi-cache && \
+    apt-get update -y && \
+    apt-get install -y python3-pip git vim wget net-tools gcc g++ cmake libnuma-dev && \
+    rm -rf /var/cache/apt/* && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install for pytest to make the docker build cache layer always valid
+RUN --mount=type=cache,target=/root/.cache/pip \
+    pip install pytest>=6.0  pytest-cov modelscope
+
+COPY . .
+
+# Install vllm-omni
+WORKDIR /workspace
+ARG VLLM_OMNI_REPO=https://github.com/vllm-project/vllm-omni.git
+ARG VLLM_OMNI_TAG=main
+ARG BUILDKITE_PULL_REQUEST
+ARG BUILDKITE_PULL_REQUEST_REPO
+RUN git config --global url."https://gh-proxy.test.osinfra.cn/https://github.com/".insteadOf "https://github.com/" && \
+    if [ "\$BUILDKITE_PULL_REQUEST" != "false" ] && [ -n "\$BUILDKITE_PULL_REQUEST" ]; then \
+        echo "Cloning and checking out PR #\$BUILDKITE_PULL_REQUEST..." && \
+        git clone \$VLLM_OMNI_REPO /workspace/vllm-omni && \
+        cd /workspace/vllm-omni && \
+        git fetch origin pull/\$BUILDKITE_PULL_REQUEST/head:pr-\$BUILDKITE_PULL_REQUEST && \
+        git checkout pr-\$BUILDKITE_PULL_REQUEST; \
+    else \
+        echo "Not a PR build, using main branch" && \
+        git clone --depth 1 \$VLLM_OMNI_REPO /workspace/vllm-omni; \
+    fi
+
+RUN --mount=type=cache,target=/root/.cache/pip \
+    export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=\$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /workspace/vllm-omni/
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+WORKDIR /workspace/vllm-omni
+CMD ["/bin/bash"]
+
+EOF
+
+# Setup cleanup
+remove_docker_container() {
+  docker rm -f "${container_name}" || true;
+  docker image rm -f "${image_name}" || true;
+  docker system prune -f || true;
+}
+trap remove_docker_container EXIT
+
+# Generate corresponding --device args based on BUILDKITE_AGENT_NAME
+# Ascend NPU BUILDKITE_AGENT_NAME format is {hostname}-{agent_idx}-{npu_card_num}cards, and agent_idx starts from 1.
+#   e.g. atlas-a2-001-1-2cards means this is the 1-th agent on atlas-a2-001 host, and it has 2 NPU cards.
+#   returns --device /dev/davinci0 --device /dev/davinci1
+parse_and_gen_devices() {
+    local input="$1"
+    local index cards_num
+    if [[ "$input" =~ ([0-9]+)-([0-9]+)cards$ ]]; then
+        index="${BASH_REMATCH[1]}"
+        cards_num="${BASH_REMATCH[2]}"
+    else
+        echo "parse error" >&2
+        return 1
+    fi
+
+    local devices=""
+    local i=0
+    while (( i < cards_num )); do
+        local dev_idx=$(((index - 1)*cards_num + i ))
+        devices="$devices --device /dev/davinci${dev_idx}"
+        ((i++))
+    done
+
+    # trim leading space
+    devices="${devices#"${devices%%[![:space:]]*}"}"
+    # Output devices: assigned to the caller variable
+    printf '%s' "$devices"
+}
+
+devices=$(parse_and_gen_devices "${BUILDKITE_AGENT_NAME}") || exit 1
+
+# Run the image and execute the Out-Of-Tree (OOT) platform interface test case on Ascend NPU hardware.
+# This test checks whether the OOT platform interface is functioning properly in conjunction with
+# the hardware plugin vllm-ascend.
+hf_model_cache_dir=/mnt/hf_cache${agent_idx}
+ms_model_cache_dir=/mnt/modelscope${agent_idx}
+mkdir -p ${hf_model_cache_dir}
+mkdir -p ${ms_model_cache_dir}
+docker run \
+    --init \
+    ${devices} \
+    --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 ${hf_model_cache_dir}:/root/.cache/huggingface \
+    -v ${ms_model_cache_dir}:/root/.cache/modelscope \
+    --network host \
+    --entrypoint="" \
+    --name "${container_name}" \
+    "${image_name}" \
+    bash -c '
+    set -e
+    VLLM_USE_MODELSCOPE=True pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+'

.buildkite/scripts/simple_test.sh

@@ -1,52 +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.11.0
-"${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]"
-"${VENV_PYTHON}" -m pytest -v -s tests/test_omni_llm.py

.buildkite/test-amd-merge.yml

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

.buildkite/test-amd-ready.yaml

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

.buildkite/test-amd.yaml

@@ -0,0 +1,119 @@
+steps:
+
+- label: "Diffusion Model Test"
+  timeout_in_minutes: 30
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export GPU_ARCHS=gfx942
+    - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
+
+- label: "Diffusion Images API LoRA E2E"
+  timeout_in_minutes: 30
+  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
+    - pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
+
+- label: "Diffusion Cache Backend Test"
+  timeout_in_minutes: 15
+  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 -m 'core_model and cache and diffusion and not distributed_rocm and MI325'
+
+- label: "Diffusion Sequence Parallelism Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  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_sequence_parallel.py -m core_model
+
+- 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_parallelism.py
+
+- label: "Diffusion GPU Worker Test"
+  timeout_in_minutes: 20
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - pytest -s -v tests/diffusion/test_diffusion_worker.py
+
+- label: "Omni Model Test Qwen2-5-Omni"
+  timeout_in_minutes: 15
+  agent_pool: mi325_2
+  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_qwen2_5_omni.py
+    - pytest -s -v tests/engine/test_async_omni_engine_abort.py
+
+- label: "Omni Model Test Qwen3-Omni"
+  timeout_in_minutes: 15
+  agent_pool: mi325_2
+  depends_on: amd-build
+  mirror_hardwares: [amdproduction]
+  grade: Blocking
+  commands:
+    - export VLLM_LOGGING_LEVEL=DEBUG
+    - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+    - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+    - export VLLM_ROCM_USE_AITER=0
+    - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
+    - pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
+
+
+- label: "Diffusion Image Edit Test"
+  timeout_in_minutes: 15
+  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_image_gen_edit.py

.buildkite/test-merge.yml

@@ -0,0 +1,380 @@
+steps:
+  - label: "Simple Unit Test"
+    depends_on: upload-merge-pipeline
+    commands:
+      - "pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model Test"
+    timeout_in_minutes: 30
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Images API LoRA E2E"
+    timeout_in_minutes: 30
+    depends_on: upload-merge-pipeline
+    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:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model CPU offloading Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    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:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Cache Backend Test"
+    timeout_in_minutes: 15
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Sequence Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
+      - pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Tensor Parallelism Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion GPU Worker Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - pytest -s -v tests/diffusion/test_diffusion_worker.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Engine Test"
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                pytest -s -v tests/engine/test_async_omni_engine_abort.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Omni Model Test"
+    timeout_in_minutes: 30
+    depends_on: upload-merge-pipeline
+    commands:
+      - export VLLM_LOGGING_LEVEL=DEBUG
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
+      - pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Omni Model Test with H100"
+    timeout_in_minutes: 30
+    depends_on: upload-merge-pipeline
+    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 -m "advanced_model" --run-level "advanced_model"
+      - pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Diffusion Image Edit Test with H100 (1 GPU)"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Bagel Model Test with H100 (Real Weights)"
+    timeout_in_minutes: 60
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 55m bash -c '
+          set -e
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Voxtral-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-merge-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

.buildkite/test-nightly.yml

@@ -0,0 +1,444 @@
+steps:
+  - label: ":full_moon: Omni Model Test with H100"
+    timeout_in_minutes: 90
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - |
+        set +e
+        pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
+        EXIT1=$$?
+        pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
+        EXIT2=$$?
+        pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        EXIT3=$$?
+        pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        EXIT4=$$?
+        exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Omni Model Test"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: ":full_moon: Diffusion Model Test with H100"
+    timeout_in_minutes: 120
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Diffusion Model Wan22 completed Test with H100"
+    timeout_in_minutes: 90
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Diffusion Model Test with L4"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: ":full_moon: Documentation Example Code Test with H100"
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export VLLM_TEST_CLEAN_GPU_MEMORY="1"
+      - pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Qwen3-TTS Non-Async-Chunk E2E Test"
+    timeout_in_minutes: 30
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - |
+        huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+        export VLLM_WORKER_MULTIPROC_METHOD=spawn
+        pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py::TestQwen3TTSNoAsyncChunk
+    agents:
+      queue: "gpu_4_queue"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
+    key: nightly-omni-performance
+    timeout_in_minutes: 180
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export BENCHMARK_DIR=tests/dfx/perf/results
+      - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
+      - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: GEBench Accuracy Test with H100"
+    key: nightly-gebench-accuracy
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1
+      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: GEdit-Bench Accuracy Test with H100"
+    key: nightly-gedit-bench-accuracy
+    timeout_in_minutes: 60
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1
+      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv"
+      - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE
+                    value: "120"
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: ":full_moon: Qwen-Image Diffusion Perf Test with H100"
+    key: nightly-qwen-image-performance
+    timeout_in_minutes: 300
+    depends_on: upload-nightly-pipeline
+    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
+    commands:
+      - export VLLM_WORKER_MULTIPROC_METHOD=spawn
+      - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
+      - export CACHE_DIT_VERSION=1.3.0
+      - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json
+      - buildkite-agent artifact upload "tests/dfx/perf/results/benchmark_results_*.json"
+      - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 4
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  # No need to run this step for PRs with label nightly-test
+  - label: ":email: Nightly Collection & Email"
+    key: nightly-perf-distribution
+    depends_on:
+      - nightly-omni-performance
+      - nightly-qwen-image-performance
+    if: build.env("NIGHTLY") == "1"
+    commands:
+      - pip install openpyxl
+      - export DEFAULT_INPUT_DIR=tests/dfx/perf/results
+      - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance
+      - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-omni-performance
+      - python tools/nightly/generate_nightly_perf_excel.py
+      - python tools/nightly/generate_nightly_perf_html.py
+      - python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.xlsx"
+      - buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
+    agents:
+      queue: "cpu_queue_premerge"

.buildkite/test-ready.yml

@@ -0,0 +1,530 @@
+steps:
+  - label: "Simple Unit Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Voxtral TTS CUDA Unit Test"
+    timeout_in_minutes: 10
+    depends_on: upload-ready-pipeline
+    commands:
+      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Batching Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Custom Pipeline Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
+    agents:
+      queue: "gpu_1_queue"
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Model CPU offloading Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
+    agents:
+      queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
+    plugins:
+      - docker#v5.2.0:
+          image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+          always-pull: true
+          propagate-environment: true
+          environment:
+            - "HF_HOME=/fsx/hf_cache"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Audio Generation Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Cache Backend Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 15m pytest -s -v -m 'core_model and cache and diffusion and not distributed_cuda and L4'
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion Sequence Parallelism Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Diffusion GPU Worker Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Engine Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 15m bash -c '
+                export VLLM_WORKER_MULTIPROC_METHOD=spawn
+                pytest -s -v tests/engine/test_async_omni_engine_abort.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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+
+  - label: "Omni Model Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 17m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
+        '
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  # - label: "Omni Model Test with H100"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         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 -m "core_model" --run-level "core_model"
+  #       '
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 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: "MiMo-Audio E2E Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 2
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Qwen3-TTS E2E Test"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          huggingface-cli download Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py -k "not NoAsyncChunk"
+        '
+    agents:
+      queue: "gpu_4_queue"
+    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"
+            - "HF_TOKEN"
+          volumes:
+            - "/fsx/hf_cache:/fsx/hf_cache"
+
+  - label: "Voxtral-TTS E2E Test"
+    timeout_in_minutes: 20
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 20m bash -c '
+          export VLLM_LOGGING_LEVEL=DEBUG
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  # - label: "Diffusion Image Edit Test with H100 (1 GPU)"
+  #   depends_on: upload-ready-pipeline
+  #   commands:
+  #     - |
+  #       timeout 20m bash -c '
+  #         export VLLM_WORKER_MULTIPROC_METHOD=spawn
+  #         pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
+  #       '
+  #   agents:
+  #     queue: "mithril-h100-pool"
+  #   plugins:
+  #     - kubernetes:
+  #         podSpec:
+  #           containers:
+  #             - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+  #               resources:
+  #                 limits:
+  #                   nvidia.com/gpu: 1
+  #               volumeMounts:
+  #                 - name: devshm
+  #                   mountPath: /dev/shm
+  #                 - name: hf-cache
+  #                   mountPath: /root/.cache/huggingface
+  #               env:
+  #                 - name: HF_HOME
+  #                   value: /root/.cache/huggingface
+  #           nodeSelector:
+  #             node.kubernetes.io/instance-type: gpu-h100-sxm
+  #           volumes:
+  #             - name: devshm
+  #               emptyDir:
+  #                 medium: Memory
+  #             - name: hf-cache
+  #               hostPath:
+  #                 path: /mnt/hf-cache
+  #                 type: DirectoryOrCreate
+
+  - label: "Bagel Text2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Bagel Img2Img Model Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 30m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate
+
+  - label: "Bagel Online Serving Test with H100"
+    depends_on: upload-ready-pipeline
+    commands:
+      - |
+        timeout 40m bash -c '
+          export VLLM_WORKER_MULTIPROC_METHOD=spawn
+          export VLLM_TEST_CLEAN_GPU_MEMORY=1
+          export VLLM_IMAGE_FETCH_TIMEOUT=60
+          pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
+        '
+    agents:
+      queue: "mithril-h100-pool"
+    plugins:
+      - kubernetes:
+          podSpec:
+            containers:
+              - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
+                resources:
+                  limits:
+                    nvidia.com/gpu: 1
+                volumeMounts:
+                  - name: devshm
+                    mountPath: /dev/shm
+                  - name: hf-cache
+                    mountPath: /root/.cache/huggingface
+                env:
+                  - name: HF_HOME
+                    value: /root/.cache/huggingface
+                  - name: HF_TOKEN
+                    valueFrom:
+                      secretKeyRef:
+                        name: hf-token-secret
+                        key: token
+            nodeSelector:
+              node.kubernetes.io/instance-type: gpu-h100-sxm
+            volumes:
+              - name: devshm
+                emptyDir:
+                  medium: Memory
+              - name: hf-cache
+                hostPath:
+                  path: /mnt/hf-cache
+                  type: DirectoryOrCreate

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

@@ -0,0 +1,58 @@
+{# vllm-omni customized version
+   Based on: https://github.com/vllm-project/ci-infra/blob/main/buildkite/test-template-amd.j2
+   Last synced: 2025-12-15
+   Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
+#}
+{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
+{% set default_working_dir = "/app/vllm-omni" %}
+
+  - group: "AMD Tests"
+    depends_on: ~
+    steps:
+      - label: "AMD: :docker: build image"
+        depends_on: ~
+        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 }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
+          - "docker push {{ docker_image_amd }}"
+        key: "amd-build"
+        env:
+          DOCKER_BUILDKIT: "1"
+        retry:
+          automatic:
+            - exit_status: -1  # Agent was lost
+              limit: 1
+            - exit_status: -10  # Agent was lost
+              limit: 1
+            - exit_status: 1  # Machine occasionally fail
+              limit: 1
+        agents:
+          queue: cpu_queue_premerge
+
+    {% for step in steps %}
+    {% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
+      - label: "{{ step.agent_pool }}: {{ step.label }}"
+        depends_on: amd-build
+        agents:
+          {% if step.agent_pool %}
+          queue: amd_{{ step.agent_pool }}
+          {% else %}
+          queue: amd_mi325_1
+          {% endif %}
+{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
+{% set indented_cmd = cmd_body | replace("\n", "\n            ") %}
+        command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
+        env:
+          DOCKER_BUILDKIT: "1"
+          TEST_COMMAND: |-
+            (command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
+            {{ indented_cmd | safe }}
+        priority: 100
+        {% if step.grade and step.grade == "Blocking" %}
+        soft_fail: false
+        {% else %}
+        soft_fail: true
+        {% endif%}
+    {% endif %}
+    {% endfor %}

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

@@ -24,7 +24,7 @@ body:
     description: |
       Please run the following and paste the output below.
       ```sh
-      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/vllm_omni/collect_env.py
+      wget https://raw.githubusercontent.com/vllm-project/vllm-omni/main/collect_env.py
       # For security purposes, please feel free to check the contents of collect_env.py before running it.
       python collect_env.py
       ```
@@ -40,6 +40,31 @@ body:
       </details>
   validations:
     required: true
+
+- type: textarea
+  attributes:
+    label: Your code version
+    description: |
+      Please provide a description of which commit id or version the code is.
+
+    value: |
+      <details>
+      <summary>The commit id or version of vllm</summary>
+
+      ```text
+
+      ```
+      </details>
+      <details>
+      <summary>The commit id or version of vllm-omni</summary>
+
+      ```text
+
+      ```
+      </details>
+  validations:
+    required: true
+
 - type: textarea
   attributes:
     label: 🐛 Describe the bug
@@ -49,28 +74,21 @@ body:
       If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
 
       ```python
-      from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
+      from vllm_omni.entrypoints.omni import Omni
+      from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+      from vllm import SamplingParams
 
-      # Create stage configurations
-      ar_config = create_ar_stage_config(
-          stage_id=0,
-          model_path="Qwen/Qwen3-0.6B",
-          input_modalities=["text"],
-          output_modalities=["text"]
+      omni = Omni(
+          model="Qwen/Qwen-Image",
+          stage_configs_path="/path/to/stage_configs.yaml",
       )
 
-      dit_config = create_dit_stage_config(
-          stage_id=1,
-          model_path="stabilityai/stable-diffusion-2-1",
-          input_modalities=["text"],
-          output_modalities=["image"]
-      )
-
-      # Initialize OmniLLM
-      omni = OmniLLM([ar_config, dit_config])
-
-      # Generate
-      outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
+      prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
+      sampling_params_list = [
+          SamplingParams(max_tokens=1),
+          OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
+      ]
+      outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
       ```
 
       If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -20,6 +20,8 @@ body:
     label: Proposed Change.
     description: >
       The proposed change of the RFC.
+    value: |
+      Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
   validations:
     required: true
 - type: textarea

.github/PULL_REQUEST_TEMPLATE.md

@@ -12,10 +12,10 @@ PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTT
 <summary> Essential Elements of an Effective PR Description Checklist </summary>
 
 - [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
-- [ ] The test plan, such as providing test command.
-- [ ] The test results, such as pasting the results comparison before and after, or e2e results
-- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model.
-- [ ] (Optional) Release notes update. If your change is user facing, please update the release notes draft.
+- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
+- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
+- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
+- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
 </details>
 
 **BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)

.github/scripts/pr_reviewer.py

@@ -0,0 +1,629 @@
+#!/usr/bin/env python3
+"""
+PR Reviewer using GLM API for vllm-omni project.
+"""
+
+import json
+import logging
+import os
+import sys
+import time
+from dataclasses import dataclass
+from typing import Any, TypedDict
+
+import requests
+
+
+# Type definitions for API responses
+class PRDetails(TypedDict):
+    """Type definition for GitHub PR details response."""
+
+    title: str
+    body: str
+    number: int
+    state: str
+    user: dict[str, Any]
+
+
+class GLMMessage(TypedDict):
+    """Type definition for GLM API message."""
+
+    role: str
+    content: str
+
+
+class GLMChoice(TypedDict):
+    """Type definition for GLM API choice."""
+
+    message: GLMMessage
+    finish_reason: str
+
+
+class GLMResponse(TypedDict):
+    """Type definition for GLM API response."""
+
+    choices: list[GLMChoice]
+    usage: dict[str, int] | None
+
+
+class GitHubComment(TypedDict):
+    """Type definition for GitHub comment."""
+
+    id: int
+    body: str
+    created_at: str
+    user: dict[str, Any]
+
+
+# Configuration
+TRIGGER_PHRASE: str = "@vllm-omni-reviewer"
+DEFAULT_GLM_API_URL: str = "https://open.bigmodel.cn/api/paas/v4/chat/completions"  # noqa: E501
+DEFAULT_GLM_MODEL: str = "glm-5"
+DEFAULT_COOLDOWN_MINUTES: int = 5
+DEFAULT_MAX_RETRIES: int = 3
+DEFAULT_RETRY_DELAY: float = 1.0
+MAX_DIFF_SIZE: int = 100_000  # Maximum diff size in characters
+
+
+@dataclass
+class Config:
+    """Configuration for the PR reviewer."""
+
+    glm_api_url: str
+    glm_model: str
+    cooldown_minutes: int
+    max_retries: int
+    retry_delay: float
+    max_diff_size: int
+
+
+# Setup logging
+logging.basicConfig(
+    level=logging.INFO,
+    format="[PR Reviewer] %(asctime)s - %(levelname)s - %(message)s",
+    datefmt="%Y-%m-%d %H:%M:%S",
+)
+logger: logging.Logger = logging.getLogger(__name__)
+
+
+def get_config() -> Config:
+    """Load configuration from environment variables with defaults."""
+    return Config(
+        glm_api_url=os.getenv("GLM_API_URL", DEFAULT_GLM_API_URL),
+        glm_model=os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL),
+        cooldown_minutes=int(
+            os.getenv(
+                "PR_REVIEWER_COOLDOWN_MINUTES",
+                str(DEFAULT_COOLDOWN_MINUTES),
+            )
+        ),
+        max_retries=int(
+            os.getenv(
+                "PR_REVIEWER_MAX_RETRIES",
+                str(DEFAULT_MAX_RETRIES),
+            )
+        ),
+        retry_delay=float(os.getenv("PR_REVIEWER_RETRY_DELAY", str(DEFAULT_RETRY_DELAY))),
+        max_diff_size=int(os.getenv("PR_REVIEWER_MAX_DIFF_SIZE", str(MAX_DIFF_SIZE))),  # noqa: E501
+    )
+
+
+def get_env_var(name: str) -> str:
+    """
+    Get an environment variable or raise an error.
+
+    Args:
+        name: Name of the environment variable.
+
+    Returns:
+        The value of the environment variable.
+
+    Raises:
+        SystemExit: If the environment variable is not set.
+    """
+    value = os.environ.get(name)
+    if not value:
+        logger.error(f"Environment variable {name} is not set")
+        sys.exit(1)
+    return value
+
+
+def check_trigger(comment_body: str) -> bool:
+    """
+    Check if the comment contains the trigger phrase.
+
+    Args:
+        comment_body: The body of the comment to check.
+
+    Returns:
+        True if the trigger phrase is found, False otherwise.
+    """
+    return TRIGGER_PHRASE in comment_body
+
+
+def fetch_pr_diff(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    max_size: int = MAX_DIFF_SIZE,
+) -> str | None:
+    """
+    Fetch the diff for a pull request.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        max_size: Maximum diff size in characters.
+
+    Returns:
+        The diff content as a string, or None if fetching failed.
+        Returns empty string if diff is larger than max_size.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3.diff",
+    }
+
+    logger.info(f"Fetching PR diff from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        diff: str = response.text
+        if len(diff) > max_size:
+            logger.warning(
+                f"Diff size ({len(diff)} bytes) exceeds maximum "
+                f"({max_size} bytes), truncating to first "
+                f"{max_size} bytes"
+            )
+            return diff[:max_size] + "\n\n... [Diff truncated due to size] ..."
+        logger.info(f"Successfully fetched diff ({len(diff)} bytes)")
+        return diff
+    else:
+        logger.error(f"Failed to fetch PR diff: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return None
+
+
+def fetch_pr_details(
+    repo_name: str,
+    pr_number: int,
+    token: str,
+) -> PRDetails | None:
+    """
+    Fetch PR details including title and description.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+
+    Returns:
+        A dictionary containing PR details, or None if fetching failed.
+    """
+    url: str = f"https://api.github.com/repos/{repo_name}/pulls/{pr_number}"
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Fetching PR details from {url}")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code == 200:
+        return response.json()
+    else:
+        logger.error(f"Failed to fetch PR details: {response.status_code}")
+        return None
+
+
+def build_review_prompt(pr_title: str, pr_description: str, diff: str) -> str:
+    """
+    Build the prompt for the GLM-4.7 API.
+
+    Args:
+        pr_title: The title of the pull request.
+        pr_description: The description/body of the pull request.
+        diff: The diff content of the pull request.
+
+    Returns:
+        The formatted prompt string for the API.
+    """
+    return f"""You are an expert code reviewer for the VLLM-Omni project. \
+Please review the following pull request:
+
+## Pull Request Details
+**Title:** {pr_title}
+
+**Description:**
+{pr_description if pr_description else "No description provided."}
+
+## Code Changes (Diff)
+{diff}
+
+## Review Guidelines
+
+Please provide a comprehensive code review with the following sections:
+
+### 1. Overview
+- Brief summary of the changes
+- Overall assessment (positive, neutral, or concerns)
+
+### 2. Code Quality
+- Code style and consistency
+- Potential bugs or edge cases
+- Performance considerations
+- Error handling
+
+### 3. Architecture & Design
+- Integration with existing codebase
+- Design patterns and best practices
+- Potential improvements
+
+### 4. Security & Safety
+- Security concerns (if any)
+- Resource management
+- Input validation
+
+### 5. Testing & Documentation
+- Test coverage considerations
+- Documentation completeness
+- Examples and usage clarity
+
+### 6. Specific Suggestions
+- Line-by-line specific feedback (use `file:line` format)
+- Concrete actionable suggestions
+- Code examples for improvements (if applicable)
+
+### 7. Approval Status
+- **LGTM** (Looks Good To Me) if the PR is ready to merge
+- **LGTM with suggestions** if the PR is good but has minor suggestions
+- **Changes requested** if significant changes are needed
+
+## Important Notes
+- Be constructive and helpful
+- Focus on objective technical feedback
+- Acknowledge good practices when you see them
+- Prioritize critical issues over nitpicks
+- If the diff is empty or minimal, acknowledge this and provide
+  any relevant context-specific guidance
+
+Please format your response in Markdown with clear section headers.
+"""
+
+
+def validate_glm_response(data: dict[str, Any]) -> str | None:
+    """
+    Validate and extract content from GLM API response.
+
+    Args:
+        data: The response data from GLM API.
+
+    Returns:
+        The review content string if valid, None otherwise.
+    """
+    # Check if choices exists and is a non-empty list
+    if "choices" not in data:
+        logger.error("GLM API response missing 'choices' field")
+        logger.error(f"Response structure: {json.dumps(data, indent=2)}")
+        return None
+
+    choices = data["choices"]
+    if not isinstance(choices, list):
+        logger.error(f"GLM API 'choices' is not a list: {type(choices)}")
+        return None
+
+    if len(choices) == 0:
+        logger.error("GLM API 'choices' is an empty list")
+        return None
+
+    # Check if first choice has message
+    try:
+        first_choice = choices[0]
+        if not isinstance(first_choice, dict):
+            logger.error(f"GLM API choice is not a dict: {type(first_choice)}")
+            return None
+
+        if "message" not in first_choice:
+            logger.error("GLM API choice missing 'message' field")
+            logger.error(f"Choice structure: {json.dumps(first_choice, indent=2)}")  # noqa: E501
+            return None
+
+        message = first_choice["message"]
+        if not isinstance(message, dict):
+            logger.error(f"GLM API message is not a dict: {type(message)}")
+            return None
+
+        if "content" not in message:
+            logger.error("GLM API message missing 'content' field")
+            logger.error(f"Message structure: {json.dumps(message, indent=2)}")
+            return None
+
+        content = message["content"]
+        if not isinstance(content, str):
+            logger.error(f"GLM API content is not a string: {type(content)}")
+            return None
+
+        return content
+
+    except (KeyError, IndexError, TypeError) as e:
+        logger.error(f"Failed to parse GLM API response: {e}")
+        logger.error(f"Response: {json.dumps(data, indent=2)}")
+        return None
+
+
+def call_glm_api(prompt: str, api_key: str, config: Config) -> str | None:
+    """
+    Call the GLM-4.7 API to get code review with retry logic.
+
+    Args:
+        prompt: The prompt to send to the API.
+        api_key: The GLM API key.
+        config: Configuration object.
+
+    Returns:
+        The review content as a string, or None if all retries failed.
+    """
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {api_key}",
+        "Content-Type": "application/json",
+    }
+
+    payload: dict[str, Any] = {
+        "model": config.glm_model,
+        "messages": [{"role": "user", "content": prompt}],
+        "temperature": 0.3,
+        "max_tokens": 32000,
+        "top_p": 0.9,
+    }
+
+    last_error: str | None = None
+
+    for attempt in range(config.max_retries):
+        try:
+            logger.info(f"Calling GLM API ({config.glm_model}) - Attempt {attempt + 1}/{config.max_retries}")
+            response = requests.post(
+                config.glm_api_url,
+                headers=headers,
+                json=payload,
+                timeout=120,
+            )
+
+            if response.status_code == 200:
+                data = response.json()
+                review = validate_glm_response(data)
+                if review:
+                    logger.info(f"Successfully received review ({len(review)} chars)")  # noqa: E501
+                    return review
+                else:
+                    last_error = "Failed to validate API response structure"
+                    logger.error(last_error)
+            else:
+                last_error = f"GLM API request failed: {response.status_code} - {response.text}"
+                logger.error(last_error)
+
+        except requests.exceptions.Timeout:
+            last_error = f"GLM API request timed out (attempt {attempt + 1})"
+            logger.error(last_error)
+        except requests.exceptions.RequestException as e:
+            last_error = f"GLM API request exception: {e}"
+            logger.error(last_error)
+        except json.JSONDecodeError as e:
+            last_error = f"Failed to decode GLM API response as JSON: {e}"
+            logger.error(last_error)
+
+        # Exponential backoff before retry
+        if attempt < config.max_retries - 1:
+            wait_time: float = config.retry_delay * (2**attempt)
+            logger.info(f"Waiting {wait_time}s before retry...")  # noqa: E501
+            time.sleep(wait_time)
+
+    logger.error(
+        f"All {config.max_retries} attempts failed. Last error: {last_error}"  # noqa: E501
+    )
+    return None
+
+
+def check_cooldown(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    cooldown_minutes: int,
+) -> bool:
+    """
+    Check if the PR is within the cooldown period.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        cooldown_minutes: Cooldown period in minutes.
+
+    Returns:
+        True if within cooldown period (should skip), False otherwise.
+    """
+    from datetime import datetime, timedelta
+
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    logger.info(f"Checking cooldown period ({cooldown_minutes} minutes)")
+    response = requests.get(url, headers=headers, timeout=30)
+
+    if response.status_code != 200:
+        logger.warning(f"Failed to check cooldown: {response.status_code}, proceeding with review")
+        return False
+
+    comments: list[dict[str, Any]] = response.json()
+    cutoff_time: datetime = datetime.utcnow() - timedelta(minutes=cooldown_minutes)  # noqa: E501
+
+    for comment in reversed(comments):
+        # Check if this is a bot comment
+        body: str = comment.get("body", "")
+        if "VLLM-Omni PR Review" in body or "PR Reviewer Bot" in body:
+            created_at_str: str = comment.get("created_at", "")
+            try:
+                # Parse GitHub timestamp format
+                created_at = datetime.fromisoformat(created_at_str.replace("Z", "+00:00"))
+                created_at = created_at.replace(tzinfo=None)
+                if created_at > cutoff_time:
+                    logger.info(f"PR is within cooldown period (last review: {created_at_str})")
+                    return True
+            except ValueError:
+                logger.warning(f"Failed to parse comment timestamp: {created_at_str}")  # noqa: E501
+                continue
+
+    logger.info("PR is outside cooldown period, proceeding with review")
+    return False
+
+
+def post_review_comment(  # noqa: E501
+    repo_name: str,
+    pr_number: int,
+    token: str,
+    review: str,
+) -> bool:
+    """
+    Post the review as a comment on the PR.
+
+    Args:
+        repo_name: The repository name in format "owner/repo".
+        pr_number: The pull request number.
+        token: GitHub authentication token.
+        review: The review content to post.
+
+    Returns:
+        True if posting succeeded, False otherwise.
+    """
+    url: str = (
+        f"https://api.github.com/repos/{repo_name}/issues/"
+        f"{pr_number}/comments"  # noqa: E501
+    )
+    headers: dict[str, str] = {
+        "Authorization": f"Bearer {token}",
+        "Accept": "application/vnd.github.v3+json",
+    }
+
+    # Format the review comment
+    comment_body: str = f"""## 🤖 VLLM-Omni PR Review
+
+{review}
+
+---
+*This review was generated automatically by the VLLM-Omni PR Reviewer Bot
+using {os.getenv("GLM_MODEL", DEFAULT_GLM_MODEL)}.*
+"""
+
+    payload: dict[str, str] = {"body": comment_body}
+
+    logger.info(f"Posting review comment to PR #{pr_number}")
+    response = requests.post(url, headers=headers, json=payload, timeout=30)
+
+    if response.status_code == 201:
+        logger.info("Successfully posted review comment")
+        return True
+    else:
+        logger.error(f"Failed to post comment: {response.status_code}")
+        logger.error(f"Response: {response.text}")
+        return False
+
+
+def main() -> int:
+    """
+    Main entry point for the PR reviewer bot.
+
+    Returns:
+        0 on success, 1 on error.
+    """
+    logger.info("VLLM-Omni PR Reviewer Bot starting...")
+
+    # Load configuration
+    config: Config = get_config()
+    logger.info(
+        f"Configuration: model={config.glm_model}, "
+        f"cooldown={config.cooldown_minutes}min, "
+        f"max_retries={config.max_retries}"
+    )
+
+    # Get environment variables
+    token: str = get_env_var("GITHUB_TOKEN")
+    api_key: str = get_env_var("GLM_API_KEY")
+    repo_name: str = get_env_var("REPO_NAME")
+    pr_number_str: str = get_env_var("PR_NUMBER")
+    comment_body: str = get_env_var("COMMENT_BODY")
+
+    try:
+        pr_number: int = int(pr_number_str)
+    except ValueError:
+        logger.error(f"Invalid PR number: {pr_number_str}")
+        return 1
+
+    logger.info(f"Repository: {repo_name}")
+    logger.info(f"PR Number: {pr_number}")
+
+    # Check if the comment contains the trigger phrase
+    if not check_trigger(comment_body):
+        logger.info(
+            f"Comment does not contain trigger phrase '{TRIGGER_PHRASE}', exiting"  # noqa: E501
+        )
+        return 0
+
+    logger.info("Trigger phrase detected! Starting review process...")
+
+    # Check cooldown period
+    if check_cooldown(repo_name, pr_number, token, config.cooldown_minutes):
+        logger.info("Skipping review due to cooldown period")
+        return 0
+
+    # Fetch PR details
+    logger.info("Step 1/4: Fetching PR details...")
+    pr_details: PRDetails | None = fetch_pr_details(repo_name, pr_number, token)  # noqa: E501
+    if not pr_details:
+        logger.error("Failed to fetch PR details")
+        return 1
+
+    pr_title: str = pr_details.get("title", "Unknown")
+    pr_description: str = pr_details.get("body", "")
+
+    logger.info(f"PR Title: {pr_title}")
+
+    # Fetch PR diff
+    logger.info("Step 2/4: Fetching PR diff...")
+    diff: str | None = fetch_pr_diff(repo_name, pr_number, token, config.max_diff_size)
+    if diff is None:
+        logger.error("Failed to fetch PR diff")
+        return 1
+
+    if not diff:
+        logger.warning("Warning: Empty diff - this might be a draft PR or no code changes")
+
+    # Build prompt
+    logger.info("Step 3/4: Building review prompt...")
+    prompt: str = build_review_prompt(pr_title, pr_description, diff)
+
+    # Call GLM API
+    logger.info("Step 4/4: Calling GLM API...")
+    review: str | None = call_glm_api(prompt, api_key, config)
+    if not review:
+        logger.error("Failed to get review from GLM API")
+        return 1
+
+    # Post review comment
+    logger.info("Posting review comment...")
+    if not post_review_comment(repo_name, pr_number, token, review):
+        logger.error("Failed to post review comment")
+        return 1
+
+    logger.info("PR review completed successfully!")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

.github/workflows/build_wheel.yml

@@ -0,0 +1,63 @@
+name: Build Wheel
+
+on:
+  push:
+    branches:
+      - main
+    paths-ignore:
+      - 'docs/**'
+      - 'examples/**'
+      - 'apps/**'
+      - '.github/**'
+      - '.buildkite/**'
+      - '**.md'
+      - '*.yml'
+      - '.gitignore'
+      - '.gitattributes'
+      - 'LICENSE'
+  pull_request:
+    branches:
+      - main
+    paths-ignore:
+      - 'docs/**'
+      - 'examples/**'
+      - 'apps/**'
+      - '.github/**'
+      - '.buildkite/**'
+      - '*.yml'
+      - '**.md'
+      - '.gitignore'
+      - '.gitattributes'
+      - 'LICENSE'
+  workflow_dispatch:
+
+jobs:
+  build:
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: ["3.11", "3.12"]
+
+    steps:
+      - uses: actions/checkout@v6
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v6
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Install uv
+        run: |
+          pip install uv
+
+      - name: Build Wheel
+        env:
+          UV_SYSTEM_PYTHON: 1
+        run: |
+          bash scripts/build_wheel.sh --python python
+
+      - name: Upload Wheel Artifact
+        uses: actions/upload-artifact@v7
+        with:
+          name: vllm-omni-wheel-py${{ matrix.python-version }}
+          path: dist/*.whl

.github/workflows/pr-reviewer.yml

@@ -0,0 +1,62 @@
+name: VLLM-Omni PR Reviewer
+
+on:
+  issue_comment:
+    types: [created]
+
+permissions:
+  contents: read
+  pull-requests: write
+  issues: write
+
+jobs:
+  pr-reviewer:
+    name: Review Pull Request
+    runs-on: ubuntu-latest
+    timeout-minutes: 10
+    # Only run when the comment is from a collaborator/owner/member
+    if: |
+      github.event_name == 'issue_comment' &&
+      github.event.issue.pull_request != null &&
+      contains(github.event.comment.body, '@vllm-omni-reviewer') &&
+      (github.event.comment.author_association == 'MEMBER' ||
+       github.event.comment.author_association == 'COLLABORATOR' ||
+       github.event.comment.author_association == 'OWNER')
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v6.0.2
+        with:
+          fetch-depth: 0
+
+      - name: Set up Python
+        uses: actions/setup-python@v6.2.0
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip==24.0
+          pip install requests==2.31.0 pyyaml==6.0.1
+
+      - name: Run PR Reviewer
+        id: reviewer
+        env:
+          GLM_API_KEY: ${{ secrets.GLM_API_KEY }}
+          GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
+          PR_NUMBER: ${{ github.event.issue.number || github.event.pull_request.number }}
+          COMMENT_BODY: ${{ github.event.comment.body }}
+          REPO_NAME: ${{ github.repository }}
+          PR_REVIEWER_COOLDOWN_MINUTES: 5
+          PR_REVIEWER_MAX_RETRIES: 3
+        run: |
+          python .github/scripts/pr_reviewer.py 2>&1 | tee "/tmp/pr_review_${PR_NUMBER}.log"
+
+      - name: Upload review logs
+        if: always()
+        uses: actions/upload-artifact@v7.0.0
+        with:
+          name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
+          path: /tmp/pr_review_*.log
+          retention-days: 7
+          if-no-files-found: ignore

.github/workflows/pre-commit.yml

@@ -18,8 +18,8 @@ jobs:
   pre-commit:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
-      - uses: actions/setup-python@v5
+      - uses: actions/checkout@v6
+      - uses: actions/setup-python@v6
         with:
           python-version: "3.12"
           cache: 'pip'

.gitignore

@@ -153,12 +153,17 @@ dmypy.json
 # Cython debug symbols
 cython_debug/
 
+# Claude
+CLAUDE.md
+.claude/
+
+# Codex
+AGENTS.md
+.codex/
+
 # cursor
 .cursor/
 
-# docker
-docker/
-
 # scripts
 /scripts/
 
@@ -182,6 +187,7 @@ checkpoints/
 
 # Cache directories
 cache/
+!vllm_omni/diffusion/cache/
 .cache/
 diffusion_cache/
 kv_cache/
@@ -216,6 +222,7 @@ data/
 datasets/
 *.csv
 *.json
+!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
 *.jsonl
 *.parquet
 
@@ -223,6 +230,7 @@ datasets/
 outputs/
 results/
 generated/
+output_*/
 
 # Configuration overrides
 configs/local.yaml
@@ -232,3 +240,10 @@ configs/development.yaml
 # Docker
 .dockerignore
 Dockerfile.dev
+discussion
+tmp_test
+
+# Auto-generated version file (created by setuptools_scm during build)
+vllm_omni/_version.py
+# output files
+*.wav

.pre-commit-config.yaml

@@ -3,11 +3,12 @@ default_install_hook_types:
   - commit-msg
 repos:
   - repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.6.0
+    rev: v6.0.0
     hooks:
       # list of supported hooks: https://pre-commit.com/hooks.html
       - id: check-yaml
         args: ["--unsafe"]
+        exclude: "examples/online_serving/chart-helm/templates/"
       - id: debug-statements
       - id: end-of-file-fixer
       - id: mixed-line-ending
@@ -16,19 +17,22 @@ repos:
         args: ["--markdown-linebreak-ext=md"]
 
   - repo: https://github.com/astral-sh/ruff-pre-commit
-    rev: v0.14.0
+    rev: v0.14.10
     hooks:
     - id: ruff-check
       args: [--output-format, github, --fix]
     - id: ruff-format
 
   - repo: https://github.com/crate-ci/typos
-    rev: v1.35.5
+    rev: typos-dict-v0.13.13
     hooks:
       - id: typos
         # only for staged files
 
   - repo: https://github.com/rhysd/actionlint
+    # 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
@@ -58,3 +62,10 @@ repos:
       verbose: true
       pass_filenames: false
     # Insert new entries above the `suggestion` entry
+
+    - id: check-pickle-imports
+      name: Prevent new pickle/cloudpickle imports
+      entry: python tools/pre_commit/check_pickle_imports.py
+      language: python
+      types: [python]
+      additional_dependencies: [regex]

README.md

@@ -1,6 +1,6 @@
 <p align="center">
   <picture>
-    <source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
+    <source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
     <img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
   </picture>
 </p>
@@ -9,13 +9,19 @@ 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="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
 </p>
 
+
 ---
 
 *Latest News* 🔥
 
+- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
+- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
+- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
+- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
+- [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.
 
 ---
@@ -66,9 +72,26 @@ 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).
 
+## Star History
+
+[![Star History Chart](https://api.star-history.com/svg?repos=vllm-project/vllm-omni&type=date&legend=top-left)](https://www.star-history.com/#vllm-project/vllm-omni&type=date&legend=top-left)
+
 ## License
 
 Apache License 2.0, as found in the [LICENSE](./LICENSE) file.

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,184 @@
+# 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
+- **Generate Video** for text-to-video and image-to-video tasks
+- **Multimodality Understanding** 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**)
+
+> [!NOTE]
+> 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**, **utils->Preview as Text**.
+- 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.
+
+## Screenshots and Examples
+
+### Multimodal understanding (e.g., Qwen Omni series, BAGEL)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Multimodal Understanding**)
+
+<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-understanding.jpg">
+    <img alt="vLLM-Omni multimodal understanding" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-understanding.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 multiple stages" 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 image generation" 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.
+
+### Text-to-video and image-to-video generation (e.g., Wan)
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video 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-video-generation.jpg">
+    <img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
+  </picture>
+</p>
+
+> [!TIP]
+> The node automatically choose text-to-video or image-to-video 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 TTS" 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.
+
+### Chaining multiple model services
+
+(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Chaining Services**)
+
+<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-chaining-services.jpg">
+    <img alt="vLLM-Omni TTS" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-chaining-services.jpg" width=55%>
+  </picture>
+</p>
+
+## 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,61 @@
+"""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__ = """vLLM-Omni Team"""
+__email__ = "vllm-omni@vllm.ai"
+__version__ = "0.0.1"
+
+from .comfyui_vllm_omni.nodes import (
+    VLLMOmniARSampling,
+    VLLMOmniDiffusionSampling,
+    VLLMOmniGenerateImage,
+    VLLMOmniGenerateVideo,
+    VLLMOmniQwenTTSParams,
+    VLLMOmniRemoteLoRA,
+    VLLMOmniSamplingParamsList,
+    VLLMOmniTTS,
+    VLLMOmniUnderstanding,
+    VLLMOmniVoiceClone,
+    VLLMOmniWanParams,
+)
+
+# A dictionary that contains all nodes you want to export with their names
+NODE_CLASS_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": VLLMOmniGenerateImage,
+    "VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
+    "VLLMOmniUnderstanding": VLLMOmniUnderstanding,
+    "VLLMOmniTTS": VLLMOmniTTS,
+    "VLLMOmniVoiceClone": VLLMOmniVoiceClone,
+    # === Params ===
+    "VLLMOmniARSampling": VLLMOmniARSampling,
+    "VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
+    "VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
+    "VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
+    "VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
+    "VLLMOmniWanParams": VLLMOmniWanParams,
+}
+
+# A dictionary that contains the friendly/humanly readable titles for the nodes
+NODE_DISPLAY_NAME_MAPPINGS = {
+    # === Generation ===
+    "VLLMOmniGenerateImage": "Generate Image",
+    "VLLMOmniGenerateVideo": "Generate Video",
+    "VLLMOmniUnderstanding": "Multimodality Understanding",
+    "VLLMOmniTTS": "TTS (Text to Speech)",
+    "VLLMOmniVoiceClone": "TTS Voice Cloning",
+    # === Params ===
+    "VLLMOmniARSampling": "AR Sampling Params",
+    "VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
+    "VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
+    "VLLMOmniRemoteLoRA": "LoRA",
+    "VLLMOmniQwenTTSParams": "Qwen TTS Params",
+    "VLLMOmniWanParams": "Wan Video 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,736 @@
+from typing import Literal
+
+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,
+    AutoregressionSamplingParams,
+    DiffusionSamplingParams,
+    QwenTTSModelSpecificParams,
+    WanModelSpecificParams,
+)
+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",),
+                "lora": ("REMOTE_LORA",),
+            },
+        }
+
+    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,
+        lora: dict | None = None,
+        **kwargs,
+    ):
+        if 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:
+            # The number of sampling parameter groups should have been validated.
+            # Now, simply convert single-item list to dict.
+            if isinstance(sampling_params, list):
+                sampling_params = sampling_params[0]
+            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,
+                    lora=lora,
+                )
+                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,
+                    lora=lora,
+                )
+                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,
+            lora=lora,
+        )
+
+        return (output,)
+
+
+class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "url": ("STRING", {"default": "http://localhost:8000/v1"}),
+                "model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
+                "prompt": ("STRING", {"multiline": True}),
+                "negative_prompt": ("STRING", {"multiline": True, "default": ""}),
+                "width": ("INT", {"default": 832, "min": 1}),
+                "height": ("INT", {"default": 480, "min": 1}),
+                "fps": ("INT", {"default": 16, "min": 1}),
+                "num_frames": ("INT", {"default": 41, "min": 1}),
+            },
+            "optional": {
+                "image": ("IMAGE",),
+                "sampling_params": ("SAMPLING_PARAMS",),
+                "lora": ("REMOTE_LORA",),
+                "model_params": ("VIDEO_PARAMS",),
+            },
+        }
+
+    RETURN_TYPES = ("VIDEO",)
+    RETURN_NAMES = ("video",)
+    FUNCTION = "generate"
+
+    async def generate(
+        self,
+        url: str,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        fps: int,
+        num_frames: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | list[dict] | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **kwargs,
+    ):
+        if kwargs:
+            logger.info("Uncaught kwargs: %s", kwargs)
+        logger.debug("Got sampling params: %s", sampling_params)
+        logger.debug("Got model params: %s", model_params)
+        validate_model_and_sampling_params_types(model, sampling_params)
+
+        # Currently, all video generation models are single-stage diffusion models
+        if isinstance(sampling_params, list):
+            if len(sampling_params) != 1:
+                raise ValueError(
+                    "Video generation expects a single sampling params group. "
+                    "Please provide one Diffusion sampling node."
+                )
+            sampling_params = sampling_params[0]
+
+        if sampling_params is not None:
+            sampling_params.pop("type", None)  # internal fields
+        if model_params is not None:
+            model_params.pop("type", None)  # internal fields
+
+        client = VLLMOmniClient(url)
+        output = await client.generate_video(
+            model=model,
+            prompt=prompt,
+            image=image,  # image present => i2v, absent => t2v
+            width=width,
+            height=height,
+            num_frames=num_frames,
+            fps=fps,
+            negative_prompt=negative_prompt,
+            sampling_params=sampling_params,
+            lora=lora,
+            model_params=model_params,
+        )
+        return (output,)
+
+
+class VLLMOmniUnderstanding(_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]:
+        if kwargs:
+            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_understanding_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_understanding_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()
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
+            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 {})}
+
+        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()
+        if not is_qwen_tts and isinstance(model_specific_params, QwenTTSModelSpecificParams):
+            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 {}),
+        }
+
+        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 = AutoregressionSamplingParams(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)",
+                    },
+                ),
+                "vae_use_tiling": (
+                    "BOOLEAN",
+                    {
+                        "default": False,
+                        "tooltip": "Enable VAE tiling 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 = DiffusionSamplingParams(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 VLLMOmniRemoteLoRA:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "local_path": ("STRING", {"default": ""}),
+                "name": ("STRING", {"default": ""}),
+                "scale": (
+                    "FLOAT",
+                    {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
+                ),
+                "int_id": (
+                    "INT",
+                    {
+                        "default": 0,
+                        "min": 0,
+                        "step": 1,
+                        "tooltip": "0 means it is not set and the server can derive it.",
+                    },
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("REMOTE_LORA",)
+    RETURN_NAMES = ("lora",)
+    FUNCTION = "get_lora"
+    CATEGORY = "vLLM-Omni"
+
+    @classmethod
+    def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
+        if not local_path.strip() or not name.strip():
+            return "Both local_path and name must be provided."
+        return True
+
+    def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
+        local_path = local_path.strip()
+        name = name.strip()
+        lora = {
+            "local_path": local_path or None,
+            "name": name or None,
+            "scale": float(scale),
+            "int_id": int(int_id) if int_id > 0 else None,
+        }
+        return (lora,)
+
+
+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 (QwenTTSModelSpecificParams(kwargs),)
+
+
+class VLLMOmniWanParams:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "guidance_scale_2": (
+                    "FLOAT",
+                    {"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
+                ),
+                "boundary_ratio": (
+                    "FLOAT",
+                    {"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
+                ),
+                "flow_shift": (
+                    "FLOAT",
+                    {"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
+                ),
+            }
+        }
+
+    RETURN_TYPES = ("VIDEO_PARAMS",)
+    RETURN_NAMES = ("Wan video params",)
+    FUNCTION = "get_params"
+    CATEGORY = "vLLM-Omni/Video Params"
+
+    def get_params(self, **kwargs):
+        return (WanModelSpecificParams(kwargs),)

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

@@ -0,0 +1,585 @@
+"""
+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.
+"""
+
+import asyncio
+import json
+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,
+    bytes_to_video,
+    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__)
+
+
+async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            try:
+                return await response.json()
+            except aiohttp.ContentTypeError as e:
+                raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
+async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
+    try:
+        async with getattr(session, verb)(url, **kwargs) as response:
+            if not response.ok:
+                error_text = await response.text()
+                raise (ValueError if response.status < 500 else RuntimeError)(
+                    f"vLLM-Omni API returned status {response.status}: {error_text}"
+                )
+            return await response.read()
+    except aiohttp.ClientError as e:
+        raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
+
+
+class VLLMOmniClient:
+    def __init__(
+        self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
+    ):
+        self.base_url = base_url
+        self.timeout = aiohttp.ClientTimeout(total=timeout)
+        self.poll_interval = poll_interval
+        self.max_poll_duration = max_poll_duration
+
+    async def generate_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        negative_prompt: str | None = None,
+        sampling_params: dict | None = None,
+        lora: 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: dict[str, Any] = {
+            "model": model,
+            "prompt": prompt,
+            "size": size,
+            "response_format": "b64_json",
+        }
+        if negative_prompt:
+            payload["negative_prompt"] = negative_prompt
+        if sampling_params is not None:
+            payload.update(sampling_params)
+        if lora is not None:
+            payload["lora"] = lora
+        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,
+        lora: 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:
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
+        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,
+        lora: 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"],
+            # === below are additional `extra_body` fields, handled by **kwargs ===
+            lora=lora,
+        )
+        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_video(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        num_frames: int,
+        fps: int,
+        negative_prompt: str | None = None,
+        image: torch.Tensor | None = None,
+        sampling_params: dict | None = None,
+        model_params: dict | None = None,
+        lora: dict | None = None,
+        **extra_body,
+    ) -> VideoInput:
+        form = aiohttp.FormData()
+        form.add_field("model", model)
+        form.add_field("prompt", prompt)
+        form.add_field("width", str(width))
+        form.add_field("height", str(height))
+        form.add_field("num_frames", str(num_frames))
+        form.add_field("fps", str(fps))
+        if negative_prompt:
+            form.add_field("negative_prompt", negative_prompt)
+        if sampling_params is not None:
+            for k, v in sampling_params.items():
+                form.add_field(k, str(v))
+        if model_params is not None:
+            for k, v in model_params.items():
+                form.add_field(k, str(v))
+        if lora is not None:
+            form.add_field("lora", json.dumps(lora, ensure_ascii=False))
+        if extra_body:
+            form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
+
+        if image is not None:
+            image_filename = "image.png"  # Required for multipart form
+            form.add_field(
+                "input_reference",
+                image_tensor_to_png_bytes(image, image_filename),
+                filename=image_filename,
+                content_type="image/png",
+            )
+
+        async with aiohttp.ClientSession(timeout=self.timeout) as session:
+            # Start the video generation job
+            url = f"{self.base_url}/videos"
+            data = await url_json(session, url, "post", data=form)
+            if (job_id := data.get("id", None)) is None:
+                raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
+            if (job_status := data.get("status", None)) is None:
+                raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+
+            # Poll for video generation job completion
+            deadline = asyncio.get_running_loop().time() + self.max_poll_duration
+            url = f"{self.base_url}/videos/{job_id}"
+            while job_status not in {"completed", "failed"}:
+                await asyncio.sleep(self.poll_interval)
+
+                data = await url_json(session, url)
+                if (job_status := data.get("status", None)) is None:
+                    raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
+                if asyncio.get_running_loop().time() >= deadline:
+                    raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
+
+            if job_status == "failed":
+                raise RuntimeError(f"Video job failed: {data}")
+
+            # Retrieve completed content
+            video_bytes = await url_bytes(session, f"{url}/content")
+
+            # Decode video and make a best effort at cleaning up server resources
+            try:
+                return bytes_to_video(video_bytes)
+            finally:
+                try:
+                    await url_json(session, url, "delete")
+                except Exception as exc:
+                    logger.warning("Failed to clean up video job %s: %s", job_id, exc)
+
+    async def generate_understanding_chat_completion(
+        self,
+        *,
+        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}]
+
+        payload: dict[str, Any] = {"messages": messages, "model": model}
+        if modalities:
+            payload["modalities"] = modalities
+
+        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
+
+            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 has already taken care of matching sampling params' types and length. Safe to take [0].
+                # * 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.copy()
+                if "n" in combined_extra_body:
+                    combined_extra_body["num_outputs_per_prompt"] = combined_extra_body.pop("n")
+            else:
+                # AR format: the payload has a sampling_params_list field, containing a list.
+                sampling_params_list = sampling_params if isinstance(sampling_params, list) else [sampling_params]
+                payload["sampling_params_list"] = sampling_params_list
+
+        if negative_prompt:
+            combined_extra_body["negative_prompt"] = negative_prompt
+
+        if extra_body:
+            combined_extra_body.update(extra_body)
+
+        # Add extra_body only if it has any content.
+        if combined_extra_body:
+            payload["extra_body"] = combined_extra_body
+
+        # Place to inject any model-specific payload adjustment
+        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,304 @@
+"""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 fractions import Fraction
+from io import BytesIO
+
+import av
+import numpy as np
+import torch
+from av.audio.frame import AudioFrame
+from av.audio.resampler import AudioResampler
+from av.video.frame import VideoFrame
+from comfy_api.input import AudioInput, VideoInput
+from comfy_api.latest import InputImpl, Types
+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 bytes_to_video(video_bytes: bytes) -> VideoInput:
+    video_buffer = BytesIO(video_bytes)
+
+    try:
+        with av.open(video_buffer, mode="r") as container:
+            video_stream = next((s for s in container.streams if s.type == "video"), None)
+            if video_stream is None:
+                raise ValueError("No video stream found in decoded payload.")
+
+            frames = []
+            for frame in container.decode(video_stream):
+                if not isinstance(frame, VideoFrame):
+                    continue
+                image = frame.to_ndarray(format="rgb24")
+                frames.append(torch.from_numpy(image).float() / 255.0)
+
+            if len(frames) == 0:
+                raise ValueError("No video frames found in decoded payload.")
+
+            images = torch.stack(frames, dim=0)
+            frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
+
+            audio: AudioInput | None = None
+            if len(container.streams.audio):
+                audio_stream = container.streams.audio[-1]
+                audio_frames = []
+                resampler = AudioResampler(format="fltp")
+                for frame in container.decode(audio_stream):
+                    if not isinstance(frame, AudioFrame):
+                        continue
+                    resampled = resampler.resample(frame)
+                    if not isinstance(resampled, list):
+                        resampled = [resampled]
+                    for audio_frame in resampled:
+                        if audio_frame is not None:
+                            audio_frames.append(audio_frame.to_ndarray())
+
+                if len(audio_frames) > 0:
+                    audio_data = np.concatenate(audio_frames, axis=1)
+                    sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
+                    audio = {
+                        "waveform": torch.from_numpy(audio_data).unsqueeze(0),
+                        "sample_rate": sample_rate,
+                    }
+
+            components = Types.VideoComponents(
+                images=images,
+                frame_rate=frame_rate,
+                audio=audio,
+                metadata=container.metadata if container.metadata else None,
+            )
+    except Exception as e:
+        raise RuntimeError(f"Failed to decode video: {e}")
+
+    return InputImpl.VideoFromComponents(components)
+
+
+def base64_to_video(base64_str: str) -> VideoInput:
+    if base64_str.startswith("data:video"):
+        _, base64_str = base64_str.split(",", 1)
+
+    try:
+        video_bytes = base64.b64decode(base64_str)
+    except Exception as e:
+        raise ValueError(f"Invalid base64 string: {e}")
+
+    return bytes_to_video(video_bytes)
+
+
+def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
+    waveform = audio["waveform"][0]  # Shape: (C, T)
+    sample_rate = audio["sample_rate"]
+    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,118 @@
+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")
+    return payload
+
+
+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.UNDERSTANDING,
+                "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.UNDERSTANDING,
+                "input_modalities": [
+                    Modality.TEXT,
+                    Modality.IMAGE,
+                    Modality.VIDEO,
+                    Modality.AUDIO,
+                ],
+            }
+        ],
+    },
+    r"Qwen3-Omni*": {
+        "stages": ["autoregression", "autoregression", "autoregression"],
+        "modes": [
+            {
+                "mode": ModelMode.UNDERSTANDING,
+                "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,55 @@
+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 AutoregressionSamplingParams(dict):
+    pass
+
+
+class DiffusionSamplingParams(dict):
+    pass
+
+
+class QwenTTSModelSpecificParams(dict):
+    pass
+
+
+class WanModelSpecificParams(dict):
+    pass
+
+
+class ModelMode(Enum):
+    IMAGE_GENERATION = auto()
+    VIDEO_GENERATION = auto()
+    AUDIO_GENERATION = auto()
+    UNDERSTANDING = 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,93 @@
+from .logger import get_logger
+from .models import lookup_model_spec
+from .types import AutoregressionSamplingParams, DiffusionSamplingParams
+
+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(f"skipping sampling params check because spec for {model_name} 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 not _check_sampling_param_matches_stage(sp, stages[i]):
+                raise ValueError(
+                    f"Sampling parameter type ({sp.__class__.__name__}) does not match "
+                    f"stage type ({stages[i]}) at index {i} for model {model_name}."
+                )
+    elif isinstance(sampling_param_list, dict):
+        # Check that the provided single sampling param matches all stages
+        for i, stage in enumerate(stages):
+            if not _check_sampling_param_matches_stage(sampling_param_list, stage):
+                raise ValueError(
+                    f"Provided single sampling parameter type ({sampling_param_list.__class__.__name__}) must match "
+                    f"the types of all stages of the model. "
+                    f"However, stage {i} of model {model_name} is of type {stage}."
+                )
+
+
+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 = sampling_param_list.__class__(sampling_param_list)
+        sampling_param_list.update(params_to_add)
+    elif sampling_param_list is None:
+        sampling_param_list = params_to_add.copy()
+    else:
+        for i, stage in enumerate(stages):
+            if stage == stage_type:
+                stage_param = sampling_param_list[i]
+                stage_param = stage_param.__class__(stage_param)
+                stage_param.update(params_to_add)
+                sampling_param_list[i] = stage_param
+
+    return sampling_param_list
+
+
+def _check_sampling_param_matches_stage(sampling_param: dict, stage_type: str) -> bool:
+    if stage_type == "autoregression":
+        return isinstance(sampling_param, AutoregressionSamplingParams)
+    if stage_type == "diffusion":
+        return isinstance(sampling_param, DiffusionSamplingParams)
+    raise RuntimeError(f"Internal error: unknown stage type {stage_type}.")

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Chaining Services.json

@@ -0,0 +1,552 @@
+{
+    "id": "6643e5fd-fa2a-4f25-935a-483173a8097c",
+    "revision": 0,
+    "last_node_id": 8,
+    "last_link_id": 6,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "PreviewImage",
+            "pos": [
+                1446.2005205859512,
+                -316.7359686049902
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 1
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 8,
+            "type": "PreviewImage",
+            "pos": [
+                1010.1848219273943,
+                -311.54216706114363
+            ],
+            "size": [
+                305.6628685610174,
+                307.0336884172169
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "images",
+                    "name": "images",
+                    "type": "IMAGE",
+                    "link": 6
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewImage"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 7,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                234.61958293819296,
+                9.405681270043619
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "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": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "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": [
+                        5
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1525,
+                "randomize"
+            ]
+        },
+        {
+            "id": 4,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                666.8380026154548,
+                268.86271068330126
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 1,
+            "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": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "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": [
+                        2
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                4,
+                50,
+                7,
+                1,
+                false,
+                false,
+                42,
+                "fixed"
+            ]
+        },
+        {
+            "id": 5,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                984.723613585788,
+                63.376900027553276
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 4
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 2
+                },
+                {
+                    "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": [
+                        1
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8001/v1",
+                "/home/models/Qwen/Qwen-Image-Edit",
+                "A high-quality, high contrast, stylized portrait of the object in the uploaded reference image. Pop art and doodle style, with abundant scribbling patterns, such as a teal crown, orange lightning bolts, colorful handwritten scripts.",
+                "Realistic",
+                800,
+                800
+            ]
+        },
+        {
+            "id": 6,
+            "type": "VLLMOmniGenerateImage",
+            "pos": [
+                541.2702717429818,
+                -160.70392744708548
+            ],
+            "size": [
+                416.56628685610167,
+                372.1662621294205
+            ],
+            "flags": {},
+            "order": 2,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": null
+                },
+                {
+                    "localized_name": "mask",
+                    "name": "mask",
+                    "shape": 7,
+                    "type": "MASK",
+                    "link": null
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 5
+                },
+                {
+                    "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": [
+                        4,
+                        6
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateImage"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "/home/models/Tongyi-MAI/Z-Image-Turbo",
+                "A headshot of a cute Siamese kitty. Blurred background due to wide aperture. Close-up look. Realistic.",
+                "Cartoonish.",
+                800,
+                800
+            ]
+        }
+    ],
+    "links": [
+        [
+            1,
+            5,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            4,
+            0,
+            5,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            4,
+            6,
+            0,
+            5,
+            0,
+            "IMAGE"
+        ],
+        [
+            5,
+            7,
+            0,
+            6,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            6,
+            6,
+            0,
+            8,
+            0,
+            "IMAGE"
+        ]
+    ],
+    "groups": [],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 1.0426432563169903,
+            "offset": [
+                -90.4071374799346,
+                575.2948920069742
+            ]
+        }
+    },
+    "version": 0.4
+}

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":"图像","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":10,"type":"MarkdownNote","pos":[227.99306819575278,-231.24143306069843],"size":[240.3326478922474,136.3505820791881],"flags":{},"order":1,"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":2,"mode":0,"inputs":[{"localized_name":"图像","name":"image","type":"COMBO","widget":{"name":"image"},"link":null},{"localized_name":"选择文件上传","name":"upload","type":"IMAGEUPLOAD","widget":{"name":"upload"},"link":null}],"outputs":[{"localized_name":"图像","name":"IMAGE","type":"IMAGE","links":[21]},{"localized_name":"遮罩","name":"MASK","type":"MASK","links":null}],"properties":{"Node name for S&R":"LoadImage"},"widgets_values":["example.png","image"]},{"id":8,"type":"VLLMOmniDiffusionSampling","pos":[478.59266934006774,183.67711984955648],"size":[284.205078125,226],"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":"vae_use_tiling","name":"vae_use_tiling","type":"BOOLEAN","widget":{"name":"vae_use_tiling"},"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,false,42,"randomize"]}],"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":[-55.18047513099182,220.2553505195962]}},"version":0.4}

apps/ComfyUI-vLLM-Omni/example_workflows/vLLM-Omni Multimodal Understanding.json

@@ -0,0 +1,761 @@
+{
+    "id": "1c99f525-0a37-45ba-a28a-7df7c3af66b4",
+    "revision": 0,
+    "last_node_id": 12,
+    "last_link_id": 14,
+    "nodes": [
+        {
+            "id": 1,
+            "type": "VLLMOmniUnderstanding",
+            "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": "VLLMOmniUnderstanding"
+            },
+            "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",
+                    "type": "SAMPLING_PARAMS",
+                    "links": [
+                        14
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniSamplingParamsList"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 11,
+            "type": "MarkdownNote",
+            "pos": [
+                826.2328280438272,
+                569.1890318701705
+            ],
+            "size": [
+                333.8220435590464,
+                261.63596728060656
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Sampling Parameters",
+            "properties": {},
+            "widgets_values": [
+                "## Sampling Parameter Types\n\nThere are two types of sampling parameters: one for autoregression and one for diffusion.\nYou should ensure that you have chosen the correct type of sampling parameters for the model you request.\n\n## Stages & Shorthand\n\nFor multi-stage models such as Qwen Omni, you can either\n- connect one sampling parameter node, which is applied to all stages.\n- connect exactly the same number of sampling parameter nodes to a \"Multi-Stage Sampling Parameter List\" node, then connect this node to the primary request node.\n\nNote that this shorthand is intended to stay consistent with the [online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/qwen2_5_omni/)"
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 12,
+            "type": "MarkdownNote",
+            "pos": [
+                378.9866207003777,
+                152.59550252215752
+            ],
+            "size": [
+                319.7287574247016,
+                107.15904081906785
+            ],
+            "flags": {},
+            "order": 6,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Input",
+            "properties": {},
+            "widgets_values": [
+                "Note that not all models support every modality as input. For example, `ByteDance-Seed/BAGEL-7B-MoT` in Multimodality Understanding mode only support text and image input.\n\nYou should ensure that the input are supported by the model. You can check the corresponding [online serving documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/) for confirmation."
+            ],
+            "color": "#432",
+            "bgcolor": "#000"
+        },
+        {
+            "id": 2,
+            "type": "LoadImage",
+            "pos": [
+                394.4674804308822,
+                -207.6987397548834
+            ],
+            "size": [
+                282.798828125,
+                314
+            ],
+            "flags": {},
+            "order": 5,
+            "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": [
+                        1
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "example.png",
+                "image"
+            ]
+        },
+        {
+            "id": 10,
+            "type": "PreviewAudio",
+            "pos": [
+                1664.548345556043,
+                297.5921292054054
+            ],
+            "size": [
+                270,
+                88
+            ],
+            "flags": {},
+            "order": 10,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "audio",
+                    "name": "audio",
+                    "type": "AUDIO",
+                    "link": 9
+                },
+                {
+                    "localized_name": "audioUI",
+                    "name": "audioUI",
+                    "type": "AUDIO_UI",
+                    "widget": {
+                        "name": "audioUI"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {
+                "Node name for S&R": "PreviewAudio"
+            },
+            "widgets_values": []
+        },
+        {
+            "id": 9,
+            "type": "ShowText|pysssss",
+            "pos": [
+                1649.2506875091847,
+                66.22823888292349
+            ],
+            "size": [
+                318.7188464232943,
+                173.38502269972975
+            ],
+            "flags": {},
+            "order": 9,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "text",
+                    "name": "text",
+                    "type": "STRING",
+                    "link": 8
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "STRING",
+                    "name": "STRING",
+                    "shape": 6,
+                    "type": "STRING",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "ShowText|pysssss"
+            },
+            "widgets_values": []
+        }
+    ],
+    "links": [
+        [
+            1,
+            2,
+            0,
+            1,
+            0,
+            "IMAGE"
+        ],
+        [
+            2,
+            3,
+            0,
+            1,
+            1,
+            "VIDEO"
+        ],
+        [
+            3,
+            4,
+            0,
+            1,
+            2,
+            "AUDIO"
+        ],
+        [
+            5,
+            7,
+            0,
+            8,
+            0,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            8,
+            1,
+            0,
+            9,
+            0,
+            "STRING"
+        ],
+        [
+            9,
+            1,
+            1,
+            10,
+            0,
+            "AUDIO"
+        ],
+        [
+            12,
+            5,
+            0,
+            8,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            13,
+            7,
+            0,
+            8,
+            2,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            14,
+            8,
+            0,
+            1,
+            3,
+            "SAMPLING_PARAMS"
+        ]
+    ],
+    "groups": [
+        {
+            "id": 1,
+            "title": "Sampling Parameters",
+            "bounding": [
+                480.1649301181556,
+                341.08402513937995,
+                692.4113568972277,
+                510.48648853403665
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        },
+        {
+            "id": 2,
+            "title": "Input",
+            "bounding": [
+                344.2364817813528,
+                -287.5850484183313,
+                704.8832238768634,
+                559.5124009832894
+            ],
+            "color": "#3f789e",
+            "font_size": 24,
+            "flags": {}
+        }
+    ],
+    "config": {},
+    "extra": {
+        "workflowRendererVersion": "LG",
+        "ds": {
+            "scale": 0.9478575057427204,
+            "offset": [
+                33.81037136029557,
+                307.1974296197726
+            ]
+        }
+    },
+    "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/example_workflows/vLLM-Omni Video Generation.json

@@ -0,0 +1,513 @@
+{
+    "id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
+    "revision": 0,
+    "last_node_id": 14,
+    "last_link_id": 28,
+    "nodes": [
+        {
+            "id": 10,
+            "type": "MarkdownNote",
+            "pos": [
+                227.99306819575278,
+                -231.24143306069843
+            ],
+            "size": [
+                240.3326478922474,
+                136.3505820791881
+            ],
+            "flags": {},
+            "order": 0,
+            "mode": 0,
+            "inputs": [],
+            "outputs": [],
+            "title": "Note: Task and Input",
+            "properties": {},
+            "widgets_values": [
+                "vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload 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": [
+                        25
+                    ]
+                },
+                {
+                    "localized_name": "MASK",
+                    "name": "MASK",
+                    "type": "MASK",
+                    "links": null
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "LoadImage"
+            },
+            "widgets_values": [
+                "cute-cat.jpg",
+                "image"
+            ]
+        },
+        {
+            "id": 8,
+            "type": "VLLMOmniDiffusionSampling",
+            "pos": [
+                465.5140218220927,
+                148.80072646828987
+            ],
+            "size": [
+                284.205078125,
+                226
+            ],
+            "flags": {},
+            "order": 2,
+            "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": "vae_use_tiling",
+                    "name": "vae_use_tiling",
+                    "type": "BOOLEAN",
+                    "widget": {
+                        "name": "vae_use_tiling"
+                    },
+                    "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": [
+                        27
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniDiffusionSampling"
+            },
+            "widgets_values": [
+                1,
+                50,
+                1,
+                1,
+                false,
+                false,
+                1992,
+                "randomize"
+            ]
+        },
+        {
+            "id": 14,
+            "type": "VLLMOmniWanParams",
+            "pos": [
+                518.6005521845747,
+                428.4455663719664
+            ],
+            "size": [
+                270,
+                106
+            ],
+            "flags": {},
+            "order": 3,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "guidance_scale_2",
+                    "name": "guidance_scale_2",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "guidance_scale_2"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "boundary_ratio",
+                    "name": "boundary_ratio",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "boundary_ratio"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "flow_shift",
+                    "name": "flow_shift",
+                    "type": "FLOAT",
+                    "widget": {
+                        "name": "flow_shift"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "Wan video params",
+                    "name": "Wan video params",
+                    "type": "VIDEO_PARAMS",
+                    "links": [
+                        28
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniWanParams"
+            },
+            "widgets_values": [
+                4,
+                0.875,
+                5
+            ]
+        },
+        {
+            "id": 12,
+            "type": "VLLMOmniGenerateVideo",
+            "pos": [
+                827.2566336087867,
+                -73.77449831827545
+            ],
+            "size": [
+                400,
+                346
+            ],
+            "flags": {},
+            "order": 4,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "image",
+                    "name": "image",
+                    "shape": 7,
+                    "type": "IMAGE",
+                    "link": 25
+                },
+                {
+                    "localized_name": "sampling_params",
+                    "name": "sampling_params",
+                    "shape": 7,
+                    "type": "SAMPLING_PARAMS",
+                    "link": 27
+                },
+                {
+                    "localized_name": "lora",
+                    "name": "lora",
+                    "shape": 7,
+                    "type": "REMOTE_LORA",
+                    "link": null
+                },
+                {
+                    "localized_name": "model_params",
+                    "name": "model_params",
+                    "shape": 7,
+                    "type": "VIDEO_PARAMS",
+                    "link": 28
+                },
+                {
+                    "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
+                },
+                {
+                    "localized_name": "fps",
+                    "name": "fps",
+                    "type": "INT",
+                    "widget": {
+                        "name": "fps"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "num_frames",
+                    "name": "num_frames",
+                    "type": "INT",
+                    "widget": {
+                        "name": "num_frames"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "links": [
+                        26
+                    ]
+                }
+            ],
+            "properties": {
+                "Node name for S&R": "VLLMOmniGenerateVideo"
+            },
+            "widgets_values": [
+                "http://localhost:8000/v1",
+                "Wan-AI/Wan2.2-TI2V-5B-Diffusers",
+                "Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
+                "Cartoonish",
+                900,
+                1200,
+                12,
+                41
+            ]
+        },
+        {
+            "id": 13,
+            "type": "SaveVideo",
+            "pos": [
+                1257.5441369501657,
+                -44.12956394419863
+            ],
+            "size": [
+                374,
+                372.34896826553796
+            ],
+            "flags": {},
+            "order": 5,
+            "mode": 0,
+            "inputs": [
+                {
+                    "localized_name": "video",
+                    "name": "video",
+                    "type": "VIDEO",
+                    "link": 26
+                },
+                {
+                    "localized_name": "filename_prefix",
+                    "name": "filename_prefix",
+                    "type": "STRING",
+                    "widget": {
+                        "name": "filename_prefix"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "format",
+                    "name": "format",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "format"
+                    },
+                    "link": null
+                },
+                {
+                    "localized_name": "codec",
+                    "name": "codec",
+                    "type": "COMBO",
+                    "widget": {
+                        "name": "codec"
+                    },
+                    "link": null
+                }
+            ],
+            "outputs": [],
+            "properties": {},
+            "widgets_values": [
+                "video/ComfyUI",
+                "auto",
+                "auto"
+            ]
+        }
+    ],
+    "links": [
+        [
+            25,
+            4,
+            0,
+            12,
+            0,
+            "IMAGE"
+        ],
+        [
+            26,
+            12,
+            0,
+            13,
+            0,
+            "VIDEO"
+        ],
+        [
+            27,
+            8,
+            0,
+            12,
+            1,
+            "SAMPLING_PARAMS"
+        ],
+        [
+            28,
+            14,
+            0,
+            12,
+            3,
+            "VIDEO_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": [
+                101.7632950847089,
+                492.29121889347806
+            ]
+        }
+    },
+    "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/README.md

@@ -0,0 +1,44 @@
+# Benchmarks Overview and Architecture
+
+This document explains the benchmark architecture across all benchmark assets in this repo. It describes what we measure, and where to find or plug in new scenarios. Per-task details remain in subfolder READMEs (e.g., `benchmarks/<model>/README.md`).
+
+## Scope and goals
+- Establish repeatable latency/throughput measurements for multimodal LLM pipelines.
+- Provide both HF Transformers (offline) and vLLM-Omni (multi-stage/pipeline) baselines.
+- Make it easy to plug in new datasets and models with minimal changes to the runner scripts.
+
+## Dataset and inputs
+- Default example: SeedTTS top-100 prompts (`benchmarks/build_dataset/top100.txt`) via `benchmarks/build_dataset/`.
+- Extensible: drop in new prompt files or modality-aligned payloads; keep the expected format for the consuming scripts (e.g., one prompt per line).
+- If you add a new dataset, document it under `benchmarks/<model>/README.md` and point scripts to your data path.
+
+## Directory layout
+- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
+- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
+- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
+  benchmark suites to vLLM-Omni serving and evaluation flows.
+- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
+
+## Reference workflows
+- **HF Transformers (offline, single process)**  
+  Script (example): `benchmarks/<model>/transformers/eval_qwen3_moe_omni_transformers.sh`  
+  Outputs: `benchmark_results/perf_stats.json`, `benchmark_results/results.json`, `benchmark_results/audio/` (if audio is produced).
+
+- **vLLM-Omni end-to-end pipeline**  
+  Script (example): `benchmarks/<model>/vllm_omni/eval_qwen3_moe_omni.sh`  
+  Outputs: `vllm_omni/logs/*.stats.jsonl` (per-stage/overall latency & TPS), `vllm_omni/logs/stage*.log`, `vllm_omni/outputs/` (text/audio artifacts).
+
+- **Adding a new task/model**  
+  1) Create `benchmarks/<model>/transformers/` and/or `benchmarks/<model>/vllm_omni/` with scripts referencing your model and dataset.  
+  2) Add a task README describing dataset, configs, and expected outputs.  
+  3) Keep the output/log structure similar for easy comparison (perf_stats/results/audio or text outputs; stats.jsonl/logs for pipeline).
+
+## Metrics to watch
+- **Throughput**: `overall_tps`, `*_tps_avg` per stage.
+- **Latency distribution**: look for long tails in `*.stats.jsonl`.
+- **Quality/completeness**: missing outputs or errors in stage logs indicate pipeline failures or misconfigurations.
+
+## Troubleshooting
+- Verify GPU/driver/FlashAttention2 requirements for your chosen model/config.
+- Ensure network access for dataset/model downloads (Google Drive, Hugging Face, etc.).
+- If outputs are missing or slow, inspect per-stage logs and `*.stats.jsonl` for errors, stragglers, or contention.

benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Benchmark helpers and runnable benchmark entrypoints."""

benchmarks/accuracy/README.md

@@ -0,0 +1,27 @@
+# Accuracy Benchmarks
+
+This directory hosts accuracy benchmark integrations that run entirely through a
+local `vllm-omni serve` deployment.
+
+Current integrations:
+
+- `text_to_image/`: GEBench generation + local judge scoring flow.
+- `image_to_image/`: GEdit-Bench generation + local VIEScore-style scoring flow.
+
+Design notes:
+
+- Generation is executed through the OpenAI-compatible endpoints exposed by
+  `vllm-omni serve`.
+- Evaluation is also executed through a local OpenAI-compatible judge model
+  served by `vllm-omni`.
+- Both generation and judge requests accept either `http://host:port` or
+  `http://host:port/v1`.
+- Output directory layout intentionally stays close to the upstream repos.
+
+Test guidance:
+
+- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
+- End-to-end generation/evaluation should be validated in a remote GPU
+  environment. In the current repo marker system there is `L4` but no `L5`
+  marker, so benchmark smoke tests should be wired as `advanced_model +
+  benchmark + L4` when GPU capacity is available.

benchmarks/accuracy/__init__.py

@@ -0,0 +1 @@
+"""Accuracy benchmark integrations for vLLM-Omni."""

benchmarks/accuracy/common.py

@@ -0,0 +1,199 @@
+from __future__ import annotations
+
+import base64
+import io
+import json
+from pathlib import Path
+from typing import Any
+
+import requests
+from PIL import Image
+
+
+def ensure_dir(path: Path) -> Path:
+    path.mkdir(parents=True, exist_ok=True)
+    return path
+
+
+def load_json(path: Path) -> dict[str, Any]:
+    with path.open("r", encoding="utf-8") as handle:
+        return json.load(handle)
+
+
+def write_json(path: Path, payload: dict[str, Any]) -> None:
+    ensure_dir(path.parent)
+    with path.open("w", encoding="utf-8") as handle:
+        json.dump(payload, handle, indent=2, ensure_ascii=False)
+
+
+def save_image(path: Path, image: Image.Image) -> None:
+    ensure_dir(path.parent)
+    image.save(path)
+
+
+def find_first_image(folder: Path, stem: str | None = None) -> Path | None:
+    patterns = [f"{stem}.*"] if stem else ["*.png", "*.jpg", "*.jpeg", "*.webp"]
+    for pattern in patterns:
+        for candidate in sorted(folder.glob(pattern)):
+            if candidate.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
+                return candidate
+    return None
+
+
+def extract_json_object(raw_text: str) -> dict[str, Any]:
+    raw_text = raw_text.strip()
+    delimiter = "||V^=^V||"
+    if raw_text.count(delimiter) >= 2:
+        start = raw_text.find(delimiter) + len(delimiter)
+        end = raw_text.rfind(delimiter)
+        raw_text = raw_text[start:end].strip()
+
+    start = raw_text.find("{")
+    end = raw_text.rfind("}")
+    if start == -1 or end == -1 or end < start:
+        raise ValueError(f"Could not find JSON object in: {raw_text[:200]}")
+    return json.loads(raw_text[start : end + 1])
+
+
+def build_openai_url(base_url: str, api_path: str) -> str:
+    base = base_url.rstrip("/")
+    normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
+    if base.endswith(normalized_path):
+        return base
+    if base.endswith("/v1"):
+        return f"{base}{normalized_path}"
+    return f"{base}/v1{normalized_path}"
+
+
+def pil_to_base64(image: Image.Image, image_format: str = "PNG") -> str:
+    buffer = io.BytesIO()
+    image.save(buffer, format=image_format)
+    return base64.b64encode(buffer.getvalue()).decode("utf-8")
+
+
+def pil_to_data_url(image: Image.Image, image_format: str = "PNG") -> str:
+    return f"data:image/{image_format.lower()};base64,{pil_to_base64(image, image_format=image_format)}"
+
+
+def decode_base64_image(encoded: str) -> Image.Image:
+    image = Image.open(io.BytesIO(base64.b64decode(encoded)))
+    image.load()
+    return image.convert("RGB")
+
+
+def pil_to_png_bytes(image: Image.Image) -> bytes:
+    buffer = io.BytesIO()
+    image.save(buffer, format="PNG")
+    return buffer.getvalue()
+
+
+class VllmOmniImageClient:
+    """Thin OpenAI-compatible image client for vLLM-Omni serving."""
+
+    def __init__(self, base_url: str, api_key: str = "EMPTY", timeout: int = 600):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.timeout = timeout
+
+    @property
+    def _headers(self) -> dict[str, str]:
+        return {
+            "Authorization": f"Bearer {self.api_key}",
+            "Content-Type": "application/json",
+        }
+
+    def generate_text_to_image(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        width: int,
+        height: int,
+        num_inference_steps: int = 20,
+        guidance_scale: float | None = None,
+        seed: int | None = None,
+        output_compression: int | None = None,
+    ) -> Image.Image:
+        payload: dict[str, Any] = {
+            "model": model,
+            "prompt": prompt,
+            "n": 1,
+            "size": f"{width}x{height}",
+            "response_format": "b64_json",
+            "num_inference_steps": num_inference_steps,
+        }
+        if guidance_scale is not None:
+            payload["guidance_scale"] = guidance_scale
+        if seed is not None:
+            payload["seed"] = seed
+        if output_compression is not None:
+            payload["output_compression"] = output_compression
+
+        response = requests.post(
+            build_openai_url(self.base_url, "/images/generations"),
+            json=payload,
+            headers=self._headers,
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        return decode_base64_image(response.json()["data"][0]["b64_json"])
+
+    def generate_image_edit(
+        self,
+        *,
+        model: str,
+        prompt: str,
+        images: Image.Image | list[Image.Image],
+        width: int,
+        height: int,
+        num_inference_steps: int = 20,
+        guidance_scale: float | None = None,
+        seed: int | None = None,
+        negative_prompt: str | None = None,
+        output_compression: int | None = None,
+    ) -> Image.Image:
+        if not isinstance(images, list):
+            images = [images]
+        data: dict[str, Any] = {
+            "model": model,
+            "prompt": prompt,
+            "n": 1,
+            "size": f"{width}x{height}",
+            "response_format": "b64_json",
+            "num_inference_steps": str(num_inference_steps),
+        }
+        if guidance_scale is not None:
+            data["guidance_scale"] = str(guidance_scale)
+        if seed is not None:
+            data["seed"] = str(seed)
+        if negative_prompt:
+            data["negative_prompt"] = negative_prompt
+        if output_compression is not None:
+            data["output_compression"] = str(output_compression)
+
+        files = [
+            (
+                "image[]" if len(images) > 1 else "image",
+                (f"image_{index}.png", pil_to_png_bytes(image), "image/png"),
+            )
+            for index, image in enumerate(images)
+        ]
+
+        edit_paths = ["/images/edits", "/images/edit"]
+        last_response: requests.Response | None = None
+        for api_path in edit_paths:
+            response = requests.post(
+                build_openai_url(self.base_url, api_path),
+                data=data,
+                files=files,
+                headers={"Authorization": f"Bearer {self.api_key}"},
+                timeout=self.timeout,
+            )
+            last_response = response
+            if response.status_code != 404:
+                response.raise_for_status()
+                return decode_base64_image(response.json()["data"][0]["b64_json"])
+
+        assert last_response is not None
+        last_response.raise_for_status()
+        raise ValueError("No image payload returned from image edit endpoint")

benchmarks/accuracy/image_to_image/README.md

@@ -0,0 +1,103 @@
+# GEdit-Bench on vLLM-Omni
+
+This integration adapts the upstream `stepfun-ai/Step1X-Edit/GEdit-Bench`
+evaluation flow into `vllm-omni/benchmarks/accuracy/image_to_image`.
+
+Upstream mapping:
+
+- `run_gedit_score.py` -> `run_gedit_bench.py evaluate`
+- `calculate_statistics.py` -> `run_gedit_bench.py summarize`
+- upstream output layout under `results/<method>/fullset/...` is preserved
+
+What changed:
+
+- The upstream repo mainly ships evaluation scripts. This integration adds a
+  generation runner that uses the local `vllm-omni` OpenAI-compatible
+  `/v1/images/edits` endpoint to produce benchmark outputs in the expected
+  directory structure.
+- The evaluator keeps the same VIEScore-style decomposition:
+  - `semantics_score`
+  - `quality_score`
+  - `overall_score = sqrt(semantics_score * quality_score)`
+- Judge calls are routed to a local OpenAI-compatible model served by
+  `vllm-omni`, not a remote provider.
+
+Dataset:
+
+- Default `--dataset-ref` is `stepfun-ai/GEdit-Bench`
+- You can also pass a local dataset directory previously saved with
+  Hugging Face `datasets`
+
+Example usage:
+
+```bash
+python benchmarks/accuracy/image_to_image/run_gedit_bench.py generate \
+  --output-root benchmarks/accuracy/image_to_image/results \
+  --base-url http://127.0.0.1:8000 \
+  --model Qwen/Qwen-Image-Edit \
+  --model-name qwen_image_edit \
+  --dataset-ref stepfun-ai/GEdit-Bench \
+  --task-type all \
+  --instruction-language en
+```
+
+```bash
+python benchmarks/accuracy/image_to_image/run_gedit_bench.py evaluate \
+  --output-root benchmarks/accuracy/image_to_image/results \
+  --model-name qwen_image_edit \
+  --save-dir benchmarks/accuracy/image_to_image/scores \
+  --dataset-ref stepfun-ai/GEdit-Bench \
+  --judge-base-url http://127.0.0.1:8000 \
+  --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
+  --judge-api-key EMPTY
+```
+
+```bash
+python benchmarks/accuracy/image_to_image/run_gedit_bench.py summarize \
+  --csv-path benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv \
+  --language en
+```
+
+Example summary output:
+
+```json
+{
+  "language": "all",
+  "languages": {
+    "en": {
+      "overall": {"count": 110, "Q_SC": 6.58, "Q_PQ": 5.89, "Q_O": 5.86},
+      "intersection": {"count": 78, "Q_SC": 6.50, "Q_PQ": 5.66, "Q_O": 5.65}
+    },
+    "cn": {
+      "overall": {"count": 110, "Q_SC": 6.90, "Q_PQ": 5.78, "Q_O": 6.11},
+      "intersection": {"count": 63, "Q_SC": 7.22, "Q_PQ": 5.59, "Q_O": 6.28}
+    }
+  }
+}
+```
+
+Example generated images to inspect:
+
+- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/background_change/en/<sample>.png`
+- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/text_change/en/<sample>.png`
+- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/subject-replace/cn/<sample>.png`
+
+Example score artifacts to inspect together with the images:
+
+- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv`
+- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_summary.json`
+
+What to expect:
+
+- `Q_SC` measures instruction following and content preservation.
+- `Q_PQ` measures image naturalness and artifact quality.
+- `Q_O` is the combined overall score; higher is better.
+- `overall.count` is the number of evaluated samples for that language, while `intersection.count` is the subset with `intersection_exist == True`.
+
+Notes:
+
+- This flow requires the optional Hugging Face `datasets` package.
+- `generate` writes `generation_manifest.json` with local output coverage.
+- The current repo marker set exposes `L4` but not `L5`, so if you promote an
+  end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
+  and `L4` markers or introduce a new repo-wide marker explicitly first.

benchmarks/accuracy/image_to_image/__init__.py

@@ -0,0 +1 @@
+"""GEdit-Bench integration for vLLM-Omni."""

benchmarks/accuracy/image_to_image/gedit_bench.py

@@ -0,0 +1,718 @@
+from __future__ import annotations
+
+import argparse
+import csv
+import json
+import logging
+import math
+import statistics
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from datetime import datetime, timezone
+from io import BytesIO
+from pathlib import Path
+from typing import Any
+
+import requests
+from PIL import Image
+from tqdm.auto import tqdm
+
+from benchmarks.accuracy.common import (
+    VllmOmniImageClient,
+    build_openai_url,
+    extract_json_object,
+    write_json,
+)
+
+GROUPS = [
+    "background_change",
+    "color_alter",
+    "material_alter",
+    "motion_change",
+    "ps_human",
+    "style_change",
+    "subject-add",
+    "subject-remove",
+    "subject-replace",
+    "text_change",
+    "tone_transfer",
+]
+DEFAULT_SAMPLES_PER_GROUP = 10
+logger = logging.getLogger(__name__)
+
+
+def infer_model_name(model: str) -> str:
+    normalized = model.rstrip("/\\")
+    name = Path(normalized).name
+    return name or normalized
+
+
+def resolve_model_name(*, model_name: str | None, model: str | None = None, output_root: Path | None = None) -> str:
+    if model_name:
+        return model_name
+    if model:
+        return infer_model_name(model)
+    if output_root is not None:
+        candidates = sorted(path.name for path in output_root.iterdir() if path.is_dir())
+        if len(candidates) == 1:
+            return candidates[0]
+        if not candidates:
+            raise ValueError(f"Could not infer model-name from empty output root: {output_root}")
+        raise ValueError(
+            f"Could not infer model-name from output root {output_root}; multiple candidates found: {candidates}"
+        )
+    raise ValueError("model-name is required when it cannot be inferred from model or output-root")
+
+
+def parse_score_payload(raw_text: str) -> dict[str, Any]:
+    try:
+        parsed = extract_json_object(raw_text)
+    except Exception:
+        stripped = raw_text.strip()
+        if stripped.startswith("[") and stripped.endswith("]"):
+            parsed = {"score": json.loads(stripped), "reasoning": ""}
+        else:
+            raise
+    score = parsed.get("score", [])
+    if not isinstance(score, list):
+        score = [score]
+    parsed["score"] = [int(value) for value in score]
+    return parsed
+
+
+def _utc_timestamp() -> str:
+    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+
+
+def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
+    by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for record in records:
+        by_task[record["task_type"]].append(record)
+        by_language[record["instruction_language"]].append(record)
+
+    return {
+        "count": len(records),
+        "by_task": {
+            task_type: {
+                "count": len(rows),
+                "samples": sorted(row["key"] for row in rows),
+            }
+            for task_type, rows in sorted(by_task.items())
+        },
+        "by_language": {
+            language: {
+                "count": len(rows),
+                "samples": sorted(row["key"] for row in rows),
+            }
+            for language, rows in sorted(by_language.items())
+        },
+    }
+
+
+def select_balanced_gedit_rows(
+    rows: list[dict[str, Any]],
+    *,
+    task_type: str = "all",
+    instruction_language: str = "all",
+    samples_per_group: int | None,
+) -> list[dict[str, Any]]:
+    filtered_rows = []
+    for row in rows:
+        if task_type != "all" and row["task_type"] != task_type:
+            continue
+        if instruction_language != "all" and row["instruction_language"] != instruction_language:
+            continue
+        filtered_rows.append(row)
+
+    if samples_per_group is None:
+        return filtered_rows
+
+    def _select_rows_for_group(group_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
+        if instruction_language != "all":
+            return group_rows[:samples_per_group]
+
+        rows_by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
+        for row in group_rows:
+            rows_by_language[str(row.get("instruction_language", "")).strip()].append(row)
+
+        ordered_languages = ["en", "cn"]
+        per_language_quota = samples_per_group // max(len(ordered_languages), 1)
+        remainder = samples_per_group % max(len(ordered_languages), 1)
+
+        selected_rows: list[dict[str, Any]] = []
+        leftovers: list[dict[str, Any]] = []
+        for index, language in enumerate(ordered_languages):
+            language_rows = rows_by_language.get(language, [])
+            quota = per_language_quota + (1 if index < remainder else 0)
+            selected_rows.extend(language_rows[:quota])
+            leftovers.extend(language_rows[quota:])
+
+        if len(selected_rows) < samples_per_group:
+            selected_rows.extend(leftovers[: samples_per_group - len(selected_rows)])
+
+        return selected_rows
+
+    if task_type != "all":
+        return _select_rows_for_group(filtered_rows)
+
+    grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for row in filtered_rows:
+        grouped_rows[row["task_type"]].append(row)
+
+    selected: list[dict[str, Any]] = []
+    for group in GROUPS:
+        selected.extend(_select_rows_for_group(grouped_rows.get(group, [])))
+    return selected
+
+
+def summarize_gedit_rows(rows: list[dict[str, Any]], language: str = "all") -> dict[str, Any]:
+    return summarize_gedit_rows_with_backbone(rows, language=language)
+
+
+def _mean_or_none(values: list[float]) -> float | None:
+    return statistics.fmean(values) if values else None
+
+
+def _to_q_metrics(section: dict[str, float | None]) -> dict[str, float | None]:
+    return {
+        "count": section.get("count"),
+        "Q_SC": section.get("avg_semantics"),
+        "Q_PQ": section.get("avg_quality"),
+        "Q_O": section.get("avg_overall"),
+    }
+
+
+def _summarize_gedit_rows_single_language(
+    rows: list[dict[str, Any]],
+    *,
+    language: str,
+) -> dict[str, Any]:
+    filtered_rows = [row for row in rows if str(row.get("instruction_language", "")).strip() == language]
+
+    def _to_bool(value: Any) -> bool:
+        if isinstance(value, bool):
+            return value
+        return str(value).strip().lower() in {"1", "true", "yes"}
+
+    per_group: dict[str, dict[str, float | None]] = {}
+    intersection_group: dict[str, dict[str, float | None]] = {}
+
+    for group in GROUPS:
+        group_rows = [row for row in filtered_rows if row.get("task_type") == group]
+        semantics = [float(row["semantics_score"]) for row in group_rows]
+        quality = [float(row["quality_score"]) for row in group_rows]
+        overall = [math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in group_rows]
+
+        per_group[group] = {
+            "count": len(group_rows),
+            "avg_semantics": _mean_or_none(semantics),
+            "avg_quality": _mean_or_none(quality),
+            "avg_overall": _mean_or_none(overall),
+        }
+
+        intersection_rows = [row for row in group_rows if _to_bool(row.get("intersection_exist", False))]
+        intersection_semantics = [float(row["semantics_score"]) for row in intersection_rows]
+        intersection_quality = [float(row["quality_score"]) for row in intersection_rows]
+        intersection_overall = [
+            math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in intersection_rows
+        ]
+        intersection_group[group] = {
+            "count": len(intersection_rows),
+            "avg_semantics": _mean_or_none(intersection_semantics),
+            "avg_quality": _mean_or_none(intersection_quality),
+            "avg_overall": _mean_or_none(intersection_overall),
+        }
+
+    overall_section = {
+        "count": len(filtered_rows),
+        "avg_semantics": _mean_or_none(
+            [score["avg_semantics"] for score in per_group.values() if score["avg_semantics"] is not None]
+        ),
+        "avg_quality": _mean_or_none(
+            [score["avg_quality"] for score in per_group.values() if score["avg_quality"] is not None]
+        ),
+        "avg_overall": _mean_or_none(
+            [score["avg_overall"] for score in per_group.values() if score["avg_overall"] is not None]
+        ),
+    }
+    intersection_section = {
+        "count": sum(score["count"] for score in intersection_group.values()),
+        "avg_semantics": _mean_or_none(
+            [score["avg_semantics"] for score in intersection_group.values() if score["avg_semantics"] is not None]
+        ),
+        "avg_quality": _mean_or_none(
+            [score["avg_quality"] for score in intersection_group.values() if score["avg_quality"] is not None]
+        ),
+        "avg_overall": _mean_or_none(
+            [score["avg_overall"] for score in intersection_group.values() if score["avg_overall"] is not None]
+        ),
+    }
+
+    return {
+        "language": language,
+        "by_group": {group: _to_q_metrics(section) for group, section in per_group.items()},
+        "overall": _to_q_metrics(overall_section),
+        "intersection": _to_q_metrics(intersection_section),
+    }
+
+
+def summarize_gedit_rows_with_backbone(
+    rows: list[dict[str, Any]],
+    *,
+    language: str = "all",
+) -> dict[str, Any]:
+    if language == "all":
+        return {
+            "language": "all",
+            "languages": {
+                single_language: _summarize_gedit_rows_single_language(
+                    rows,
+                    language=single_language,
+                )
+                for single_language in ["en", "cn"]
+            },
+        }
+
+    return _summarize_gedit_rows_single_language(rows, language=language)
+
+
+def _require_datasets():
+    try:
+        from datasets import load_dataset, load_from_disk
+    except ImportError as exc:
+        raise ImportError("GEdit-Bench requires the optional `datasets` package.") from exc
+    return load_dataset, load_from_disk
+
+
+def _load_gedit_dataset(dataset_ref: str):
+    load_dataset, load_from_disk = _require_datasets()
+    dataset_path = Path(dataset_ref)
+    if dataset_path.exists():
+        if (dataset_path / "state.json").exists() and (dataset_path / "dataset_info.json").exists():
+            return load_from_disk(str(dataset_path))
+        return load_dataset(str(dataset_path))
+    return load_dataset(dataset_ref)
+
+
+def _resolve_gedit_split(dataset_obj: Any) -> Any:
+    if isinstance(dataset_obj, dict):
+        if "train" in dataset_obj:
+            return dataset_obj["train"]
+        return dataset_obj
+    try:
+        return dataset_obj["train"]
+    except Exception:
+        return dataset_obj
+
+
+def _to_pil_image(value: Any) -> Image.Image:
+    if isinstance(value, Image.Image):
+        return value.convert("RGB")
+    if isinstance(value, dict) and "bytes" in value:
+        image = Image.open(BytesIO(value["bytes"]))
+        image.load()
+        return image.convert("RGB")
+    raise TypeError(f"Unsupported image payload type: {type(value)!r}")
+
+
+class LocalVIEScorer:
+    def __init__(self, *, base_url: str, api_key: str, model: str, timeout: int = 600):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.model = model
+        self.timeout = timeout
+        self.sc_prompt = (
+            "You are evaluating image editing quality.\n"
+            "Two images are provided: the source image and the edited image.\n"
+            'Return JSON only in the format {"score": [edit_success, content_preservation], '
+            '"reasoning": "..."}.\n'
+            "Each score must be an integer from 0 to 10.\n"
+            "Editing instruction: <instruction>"
+        )
+        self.pq_prompt = (
+            "You are evaluating image quality.\n"
+            'Return JSON only in the format {"score": [naturalness, artifact_free], "reasoning": "..."}.\n'
+            "Each score must be an integer from 0 to 10."
+        )
+
+    def _request(self, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
+        from benchmarks.accuracy.common import pil_to_data_url
+
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        for image in images:
+            content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
+
+        response = requests.post(
+            build_openai_url(self.base_url, "/chat/completions"),
+            json={
+                "model": self.model,
+                "messages": [{"role": "user", "content": content}],
+                "temperature": 0,
+            },
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        message_content = response.json()["choices"][0]["message"]["content"]
+        if isinstance(message_content, list):
+            text = "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
+        else:
+            text = str(message_content)
+        return parse_score_payload(text)
+
+    def evaluate(self, source_image: Image.Image, edited_image: Image.Image, instruction: str) -> dict[str, float]:
+        sc_payload = self._request(self.sc_prompt.replace("<instruction>", instruction), [source_image, edited_image])
+        pq_payload = self._request(self.pq_prompt, [edited_image])
+        semantics = float(min(sc_payload["score"])) if sc_payload["score"] else 0.0
+        quality = float(min(pq_payload["score"])) if pq_payload["score"] else 0.0
+        overall = math.sqrt(semantics * quality)
+        return {
+            "semantics_score": semantics,
+            "quality_score": quality,
+            "overall_score": overall,
+        }
+
+
+class GEditBenchRunner:
+    def __init__(
+        self,
+        *,
+        dataset_ref: str,
+        output_root: Path,
+        base_url: str,
+        model: str,
+        api_key: str = "EMPTY",
+        width: int = 512,
+        height: int = 512,
+        num_inference_steps: int = 20,
+        guidance_scale: float | None = None,
+        seed: int | None = 42,
+    ):
+        self.dataset_ref = dataset_ref
+        self.output_root = output_root
+        self.model = model
+        self.width = width
+        self.height = height
+        self.num_inference_steps = num_inference_steps
+        self.guidance_scale = guidance_scale
+        self.seed = seed
+        self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
+
+    def generate(
+        self,
+        *,
+        model_name: str,
+        task_type: str = "all",
+        instruction_language: str = "all",
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_group: int | None = None,
+    ) -> list[dict[str, Any]]:
+        dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
+        rows = select_balanced_gedit_rows(
+            list(dataset),
+            task_type=task_type,
+            instruction_language=instruction_language,
+            samples_per_group=samples_per_group,
+        )
+        if max_samples is not None:
+            rows = rows[:max_samples]
+
+        outputs: list[dict[str, Any]] = []
+        total = len(rows)
+        if workers <= 1:
+            with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
+                for item in rows:
+                    result = self._safe_generate_one(model_name, item)
+                    if result:
+                        outputs.append(result)
+                    progress.update(1)
+            return outputs
+
+        with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
+            with ThreadPoolExecutor(max_workers=workers) as executor:
+                futures = [executor.submit(self._safe_generate_one, model_name, item) for item in rows]
+                for future in as_completed(futures):
+                    result = future.result()
+                    if result:
+                        outputs.append(result)
+                    progress.update(1)
+        return outputs
+
+    def _safe_generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        try:
+            return self._generate_one(model_name, item)
+        except Exception:
+            logger.exception("Failed to generate GEdit-Bench sample %s", item.get("key", "<unknown>"))
+            return None
+
+    def _generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        output_path = (
+            self.output_root
+            / model_name
+            / "fullset"
+            / item["task_type"]
+            / item["instruction_language"]
+            / f"{item['key']}.png"
+        )
+        output_path.parent.mkdir(parents=True, exist_ok=True)
+        if output_path.exists():
+            return {
+                "task_type": item["task_type"],
+                "instruction_language": item["instruction_language"],
+                "key": item["key"],
+                "output_path": str(output_path),
+            }
+
+        source_image = _to_pil_image(item["input_image_raw"])
+        edited_image = self.client.generate_image_edit(
+            model=self.model,
+            prompt=item["instruction"],
+            images=source_image,
+            width=self.width,
+            height=self.height,
+            num_inference_steps=self.num_inference_steps,
+            guidance_scale=self.guidance_scale,
+            seed=self.seed,
+        )
+        edited_image.save(output_path)
+        return {
+            "task_type": item["task_type"],
+            "instruction_language": item["instruction_language"],
+            "key": item["key"],
+            "output_path": str(output_path),
+        }
+
+
+class GEditBenchEvaluator:
+    def __init__(self, *, dataset_ref: str, output_root: Path, scorer: LocalVIEScorer):
+        self.dataset_ref = dataset_ref
+        self.output_root = output_root
+        self.scorer = scorer
+
+    def evaluate(
+        self,
+        *,
+        model_name: str,
+        save_dir: Path,
+        task_type: str = "all",
+        instruction_language: str = "all",
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_group: int | None = None,
+    ) -> dict[str, Any]:
+        dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
+        rows = select_balanced_gedit_rows(
+            list(dataset),
+            task_type=task_type,
+            instruction_language=instruction_language,
+            samples_per_group=samples_per_group,
+        )
+        if max_samples is not None:
+            rows = rows[:max_samples]
+
+        results: list[dict[str, Any]] = []
+        total = len(rows)
+        if workers <= 1:
+            with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
+                for item in rows:
+                    result = self._safe_evaluate_one(model_name, item)
+                    if result:
+                        results.append(result)
+                    progress.update(1)
+        else:
+            with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
+                with ThreadPoolExecutor(max_workers=workers) as executor:
+                    futures = [executor.submit(self._safe_evaluate_one, model_name, item) for item in rows]
+                    for future in as_completed(futures):
+                        result = future.result()
+                        if result:
+                            results.append(result)
+                        progress.update(1)
+
+        save_dir.mkdir(parents=True, exist_ok=True)
+        base_name = f"{model_name}_{task_type}_{instruction_language}"
+        timestamp = _utc_timestamp()
+        csv_path = save_dir / f"{base_name}_vie_score.csv"
+        timestamped_csv_path = save_dir / f"{base_name}_vie_score_{timestamp}.csv"
+
+        def _write_csv(path: Path) -> None:
+            with path.open("w", encoding="utf-8", newline="") as handle:
+                writer = csv.DictWriter(
+                    handle,
+                    fieldnames=[
+                        "key",
+                        "task_type",
+                        "edited_image",
+                        "instruction",
+                        "semantics_score",
+                        "quality_score",
+                        "overall_score",
+                        "intersection_exist",
+                        "instruction_language",
+                    ],
+                )
+                writer.writeheader()
+                for row in results:
+                    writer.writerow(row)
+
+        _write_csv(csv_path)
+        _write_csv(timestamped_csv_path)
+
+        summary = summarize_gedit_rows_with_backbone(
+            results,
+            language=instruction_language,
+        )
+        summary_path = save_dir / f"{base_name}_summary.json"
+        timestamped_summary_path = save_dir / f"{base_name}_summary_{timestamp}.json"
+        write_json(summary_path, summary)
+        write_json(timestamped_summary_path, summary)
+        return {
+            "results": results,
+            "summary": summary,
+            "csv_path": str(csv_path),
+            "summary_path": str(summary_path),
+            "timestamped_csv_path": str(timestamped_csv_path),
+            "timestamped_summary_path": str(timestamped_summary_path),
+        }
+
+    def _safe_evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        try:
+            return self._evaluate_one(model_name, item)
+        except Exception:
+            logger.exception("Failed to evaluate GEdit-Bench sample %s", item.get("key", "<unknown>"))
+            return None
+
+    def _evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
+        edited_image_path = (
+            self.output_root
+            / model_name
+            / "fullset"
+            / item["task_type"]
+            / item["instruction_language"]
+            / f"{item['key']}.png"
+        )
+        if not edited_image_path.exists():
+            return None
+
+        source_image = _to_pil_image(item["input_image_raw"])
+        edited_image = Image.open(edited_image_path).convert("RGB")
+        scores = self.scorer.evaluate(source_image, edited_image, item["instruction"])
+        return {
+            "key": item["key"],
+            "task_type": item["task_type"],
+            "edited_image": str(edited_image_path),
+            "instruction": item["instruction"],
+            "semantics_score": scores["semantics_score"],
+            "quality_score": scores["quality_score"],
+            "overall_score": scores["overall_score"],
+            "intersection_exist": item.get("Intersection_exist", False),
+            "instruction_language": item["instruction_language"],
+        }
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run the GEdit-Bench integration against a local vLLM-Omni server.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    generate = subparsers.add_parser("generate")
+    generate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
+    generate.add_argument("--output-root", type=Path, required=True)
+    generate.add_argument("--base-url", type=str, required=True)
+    generate.add_argument("--model", type=str, required=True)
+    generate.add_argument("--model-name", type=str, default=None)
+    generate.add_argument("--api-key", type=str, default="EMPTY")
+    generate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
+    generate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
+    generate.add_argument("--width", type=int, default=512)
+    generate.add_argument("--height", type=int, default=512)
+    generate.add_argument("--num-inference-steps", type=int, default=20)
+    generate.add_argument("--guidance-scale", type=float, default=None)
+    generate.add_argument("--seed", type=int, default=42)
+    generate.add_argument("--workers", type=int, default=1)
+    generate.add_argument("--max-samples", type=int, default=None)
+    generate.add_argument("--samples-per-group", type=int, default=None)
+
+    evaluate = subparsers.add_parser("evaluate")
+    evaluate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
+    evaluate.add_argument("--output-root", type=Path, required=True)
+    evaluate.add_argument("--model-name", type=str, default=None)
+    evaluate.add_argument("--save-dir", type=Path, required=True)
+    evaluate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
+    evaluate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
+    evaluate.add_argument("--judge-base-url", type=str, required=True)
+    evaluate.add_argument("--judge-model", type=str, required=True)
+    evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
+    evaluate.add_argument("--workers", type=int, default=1)
+    evaluate.add_argument("--max-samples", type=int, default=None)
+    evaluate.add_argument("--samples-per-group", type=int, default=None)
+
+    summarize = subparsers.add_parser("summarize")
+    summarize.add_argument("--csv-path", type=Path, required=True)
+    summarize.add_argument("--language", choices=["all", "en", "cn"], default="all")
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "generate":
+        model_name = resolve_model_name(model_name=args.model_name, model=args.model)
+        runner = GEditBenchRunner(
+            dataset_ref=args.dataset_ref,
+            output_root=args.output_root,
+            base_url=args.base_url,
+            model=args.model,
+            api_key=args.api_key,
+            width=args.width,
+            height=args.height,
+            num_inference_steps=args.num_inference_steps,
+            guidance_scale=args.guidance_scale,
+            seed=args.seed,
+        )
+        records = runner.generate(
+            model_name=model_name,
+            task_type=args.task_type,
+            instruction_language=args.instruction_language,
+            workers=args.workers,
+            max_samples=args.max_samples,
+            samples_per_group=args.samples_per_group,
+        )
+        payload = {"records": records, "summary": summarize_generated_records(records)}
+        write_json(args.output_root / model_name / "generation_manifest.json", payload)
+        return 0
+
+    if args.command == "evaluate":
+        model_name = resolve_model_name(model_name=args.model_name, output_root=args.output_root)
+        scorer = LocalVIEScorer(
+            base_url=args.judge_base_url,
+            api_key=args.judge_api_key,
+            model=args.judge_model,
+        )
+        evaluator = GEditBenchEvaluator(dataset_ref=args.dataset_ref, output_root=args.output_root, scorer=scorer)
+        evaluator.evaluate(
+            model_name=model_name,
+            save_dir=args.save_dir,
+            task_type=args.task_type,
+            instruction_language=args.instruction_language,
+            workers=args.workers,
+            max_samples=args.max_samples,
+            samples_per_group=args.samples_per_group,
+        )
+        return 0
+
+    if args.command == "summarize":
+        with args.csv_path.open("r", encoding="utf-8", newline="") as handle:
+            rows = list(csv.DictReader(handle))
+        summary = summarize_gedit_rows_with_backbone(rows, language=args.language)
+        print(json.dumps(summary, indent=2, ensure_ascii=False))
+        return 0
+
+    parser.error(f"Unknown command: {args.command}")
+    return 1

benchmarks/accuracy/image_to_image/run_gedit_bench.py

@@ -0,0 +1,13 @@
+# ruff: noqa: E402, I001
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from benchmarks.accuracy.image_to_image.gedit_bench import main
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

benchmarks/accuracy/text_to_image/README.md

@@ -0,0 +1,103 @@
+# GEBench on vLLM-Omni
+
+This integration adapts the upstream `stepfun-ai/GEBench` scripts into
+`vllm-omni/benchmarks/accuracy/text_to_image`.
+
+Upstream mapping:
+
+- `scripts/generate.py` -> `run_gebench.py generate`
+- `scripts/evaluate.py` -> `run_gebench.py evaluate`
+- upstream prompt / judge logic -> `gbench.py`
+
+What changed:
+
+- Generation calls the local OpenAI-compatible `vllm-omni` endpoints:
+  - `/v1/images/generations` for text-only frame generation
+  - `/v1/images/edits` for image-conditioned GUI transition generation
+- Evaluation still keeps the GEBench scoring dimensions:
+  - `goal`
+  - `logic`
+  - `cons`
+  - `ui`
+  - `qual`
+- Judge calls are also routed to a local OpenAI-compatible model served by
+  `vllm-omni` instead of a remote service.
+
+Expected dataset layout:
+
+- Clone the benchmark dataset from Hugging Face into a local directory:
+
+```bash
+git clone https://huggingface.co/datasets/stepfun-ai/GEBench /path/to/GEBench
+```
+
+Example usage:
+
+```bash
+python benchmarks/accuracy/text_to_image/run_gebench.py generate \
+  --dataset-root /path/to/GEBench \
+  --output-root benchmarks/accuracy/text_to_image/outputs \
+  --base-url http://127.0.0.1:8000 \
+  --model Tongyi-MAI/Z-Image-Turbo \
+  --data-type type3
+```
+
+```bash
+python benchmarks/accuracy/text_to_image/run_gebench.py evaluate \
+  --dataset-root /path/to/GEBench \
+  --output-root benchmarks/accuracy/text_to_image/outputs \
+  --data-type type3 \
+  --judge-base-url http://127.0.0.1:8000 \
+  --judge-model Qwen/Qwen2.5-VL-7B-Instruct \
+  --judge-api-key EMPTY
+```
+
+```bash
+python benchmarks/accuracy/text_to_image/run_gebench.py summarize \
+  --output-root benchmarks/accuracy/text_to_image/outputs
+```
+
+Example summary output:
+
+```json
+{
+  "generation": {
+    "count": 20,
+    "by_type": {
+      "type3": {"count": 10},
+      "type4": {"count": 10}
+    }
+  },
+  "evaluation": {
+    "count": 20,
+    "overall_mean": 0.52,
+    "by_type": {
+      "type3": {"count": 10, "overall_mean": 0.50, "overall_mean_100": 50.0},
+      "type4": {"count": 10, "overall_mean": 0.54, "overall_mean_100": 54.0}
+    }
+  }
+}
+```
+
+Example generated images to inspect:
+
+- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame0.png`
+- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame5.png`
+- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame0.png`
+- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame5.png`
+
+What to expect:
+
+- `overall_mean` is normalized to `0.0 ~ 1.0`; higher is better.
+- `frame0.png` is the initial GUI frame; `frame5.png` is the final trajectory frame most often used for quick inspection.
+- For full debugging, inspect the whole `frame0.png` ... `frame5.png` sequence for one sample directory.
+
+Notes:
+
+- GEBench upstream leaves type3/type4 generation unfinished. This integration
+  fills that gap with a trajectory runner that generates `frame0.png` followed
+  by `frame1.png` ... `frame5.png`.
+- Type1/2/5 require an image-edit capable model exposed through
+  `vllm-omni serve`.
+- `summarize` will report both generated coverage and any existing evaluation
+  summary files.

benchmarks/accuracy/text_to_image/__init__.py

@@ -0,0 +1 @@
+"""GEBench integration for vLLM-Omni."""

benchmarks/accuracy/text_to_image/gbench.py

@@ -0,0 +1,927 @@
+from __future__ import annotations
+
+import argparse
+import json
+import statistics
+from collections import defaultdict
+from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
+from datetime import datetime, timezone
+from pathlib import Path
+from typing import Any
+
+import requests
+from PIL import Image
+
+from benchmarks.accuracy.common import (
+    VllmOmniImageClient,
+    build_openai_url,
+    ensure_dir,
+    extract_json_object,
+    find_first_image,
+    load_json,
+    pil_to_data_url,
+    save_image,
+    write_json,
+)
+
+TYPE_TO_FOLDER = {
+    "type1": "01_single_step",
+    "type2": "02_multi_step",
+    "type3": "03_trajectory_text_fictionalapp",
+    "type4": "04_trajectory_text_realapp",
+    "type5": "05_grounding_data",
+}
+SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
+DEFAULT_SAMPLES_PER_TYPE = 10
+
+
+def _utc_timestamp() -> str:
+    return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
+
+
+def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
+    write_json(path, payload)
+    timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
+    write_json(timestamped_path, payload)
+    return timestamped_path
+
+
+@dataclass(frozen=True)
+class GEBenchSampleSpec:
+    sample_path: Path
+    metadata: dict[str, Any]
+    sample_name: str
+    lang_device: str
+
+
+def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
+    by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for record in records:
+        by_type[record["data_type"]].append(record)
+
+    return {
+        "count": len(records),
+        "by_type": {
+            data_type: {
+                "count": len(rows),
+            }
+            for data_type, rows in sorted(by_type.items())
+        },
+    }
+
+
+def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
+    by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
+    for result in results:
+        by_type[result["data_type"]].append(result)
+
+    summary: dict[str, Any] = {
+        "count": len(results),
+        "overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
+        "by_type": {},
+    }
+    for data_type, rows in by_type.items():
+        score_means: dict[str, float] = {}
+        all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
+        for score_key in all_score_keys:
+            values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
+            score_means[score_key] = statistics.fmean(values) if values else 0.0
+        overall_mean = statistics.fmean(row["overall"] for row in rows)
+        summary["by_type"][data_type] = {
+            "count": len(rows),
+            "overall_mean": overall_mean,
+            "overall_mean_100": overall_mean * 100.0,
+            "score_means": score_means,
+        }
+    return summary
+
+
+def select_balanced_gebench_samples(
+    sample_paths_by_type: dict[str, list[Any]],
+    *,
+    samples_per_type: int | None,
+) -> dict[str, list[Any]]:
+    if samples_per_type is None:
+        return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
+    return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
+
+
+def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
+    records: list[dict[str, Any]] = []
+    for data_type, folder_name in TYPE_TO_FOLDER.items():
+        type_root = output_root / folder_name
+        if not type_root.exists():
+            continue
+        for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
+            for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
+                expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
+                if expected is None:
+                    expected = find_first_image(sample_dir)
+                elif not expected.exists():
+                    expected = None
+                if expected is None:
+                    continue
+                records.append(
+                    {
+                        "data_type": data_type,
+                        "sample_name": f"{lang_dir.name}/{sample_dir.name}",
+                        "output_path": str(expected),
+                    }
+                )
+    return summarize_generated_records(records)
+
+
+def _normalize_score_key(key: str) -> str:
+    mapping = {
+        "goal": "goal",
+        "logic": "logic",
+        "cons": "cons",
+        "consistency": "cons",
+        "ui": "ui",
+        "qual": "qual",
+        "quality": "qual",
+    }
+    return mapping.get(key.lower(), key.lower())
+
+
+def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
+    scores: dict[str, int] = {}
+    for key, value in raw_scores.items():
+        normalized = _normalize_score_key(key)
+        if normalized not in SCORE_KEYS:
+            continue
+        scalar = value.get("s", 0) if isinstance(value, dict) else value
+        try:
+            scores[normalized] = int(scalar)
+        except (TypeError, ValueError):
+            scores[normalized] = 0
+    for key in SCORE_KEYS:
+        scores.setdefault(key, 0)
+    return scores
+
+
+def _compute_overall(scores: dict[str, int]) -> float:
+    return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
+
+
+def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
+    data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
+    if not data_dir.exists():
+        raise FileNotFoundError(f"Dataset path not found: {data_dir}")
+
+    samples: list[Path] = []
+    for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
+        for child in sorted(lang_dir.iterdir()):
+            if child.is_dir():
+                samples.append(child)
+            elif child.suffix.lower() == ".json":
+                samples.append(child)
+    return samples
+
+
+def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
+    sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
+    if sample_id:
+        return str(sample_id)
+    if item_index is not None:
+        return f"{sample_path.stem}_{item_index:04d}"
+    return sample_path.stem if sample_path.is_file() else sample_path.name
+
+
+def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
+    if sample_path.is_dir():
+        metadata = _load_metadata(sample_path)
+        return [
+            GEBenchSampleSpec(
+                sample_path=sample_path,
+                metadata=metadata,
+                sample_name=_sample_name(sample_path),
+                lang_device=_lang_device(sample_path, metadata),
+            )
+        ]
+
+    payload = load_json(sample_path)
+    if isinstance(payload, dict):
+        return [
+            GEBenchSampleSpec(
+                sample_path=sample_path,
+                metadata=payload,
+                sample_name=_sample_name_from_metadata(payload, sample_path),
+                lang_device=_lang_device(sample_path, payload),
+            )
+        ]
+
+    if isinstance(payload, list):
+        specs: list[GEBenchSampleSpec] = []
+        for item_index, item in enumerate(payload):
+            if not isinstance(item, dict):
+                continue
+            specs.append(
+                GEBenchSampleSpec(
+                    sample_path=sample_path,
+                    metadata=item,
+                    sample_name=_sample_name_from_metadata(item, sample_path, item_index),
+                    lang_device=_lang_device(sample_path, item),
+                )
+            )
+        return specs
+
+    raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
+
+
+def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
+    sample_specs: list[GEBenchSampleSpec] = []
+    for sample_path in _iter_sample_paths(dataset_root, data_type):
+        sample_specs.extend(_expand_sample_path(sample_path))
+    return sample_specs
+
+
+def _load_metadata(sample_path: Path) -> dict[str, Any]:
+    if sample_path.is_file():
+        return load_json(sample_path)
+    for candidate in ("meta_data.json", "metadata.json"):
+        meta_path = sample_path / candidate
+        if meta_path.exists():
+            return load_json(meta_path)
+    raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
+
+
+def _sample_name(sample_path: Path) -> str:
+    return sample_path.stem if sample_path.is_file() else sample_path.name
+
+
+def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
+    return str(metadata.get("lang_device") or sample_path.parent.name)
+
+
+def _resolve_referenced_image(
+    *,
+    metadata: dict[str, Any],
+    sample_path: Path,
+    dataset_root: Path,
+    data_type: str,
+) -> Image.Image | None:
+    for key in ("image", "input_image", "initial_image", "reference_image"):
+        image_ref = metadata.get(key)
+        if not image_ref:
+            continue
+        candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
+        if candidate.exists():
+            image = Image.open(candidate)
+            image.load()
+            return image.convert("RGB")
+    if sample_path.is_dir():
+        local_image = find_first_image(sample_path)
+        if local_image:
+            image = Image.open(local_image)
+            image.load()
+            return image.convert("RGB")
+    return None
+
+
+def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
+    for key in ("trajectory", "steps", "frames"):
+        value = metadata.get(key)
+        if isinstance(value, list):
+            return [step for step in value if isinstance(step, dict)]
+    extracted: list[dict[str, Any]] = []
+    for index in range(1, 6):
+        value = metadata.get(f"step{index}") or metadata.get(str(index))
+        if isinstance(value, dict):
+            extracted.append(value)
+    return extracted
+
+
+def _text_or_default(value: Any, default: str = "") -> str:
+    return str(value).strip() if value is not None else default
+
+
+def _type1_prompt(metadata: dict[str, Any]) -> str:
+    caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
+    return (
+        "Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
+        f"Requested change:\n{caption}\n\n"
+        "Requirements:\n"
+        "- Preserve layout, visual identity, and unrelated regions.\n"
+        "- Only apply the requested state change.\n"
+        "- Keep all text and controls readable.\n"
+    )
+
+
+def _type2_prompt(goal: str, step_num: int) -> str:
+    return (
+        "Generate the next GUI state for a multi-step task.\n\n"
+        f"Overall goal: {goal}\n"
+        f"Current progress step: {step_num}/5\n\n"
+        "Requirements:\n"
+        "- The change should be incremental and plausible.\n"
+        "- Preserve layout and visual identity.\n"
+        "- Make text/buttons readable.\n"
+    )
+
+
+def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
+    app_name = _text_or_default(metadata.get("app_name"), "App")
+    final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
+    visual_description = _text_or_default(
+        metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
+        "A clean product-quality app home screen.",
+    )
+    return (
+        "Generate the first GUI frame for a task trajectory.\n\n"
+        f"App name: {app_name}\n"
+        f"Final goal: {final_goal}\n"
+        f"Visual description:\n{visual_description}\n\n"
+        "Requirements:\n"
+        "- Generate a production-looking UI screenshot only.\n"
+        "- Keep the layout coherent and readable.\n"
+    )
+
+
+def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
+    action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
+    visual_description = _text_or_default(
+        step_info.get("visual_description") or step_info.get("description"),
+        "Reflect the expected next GUI state.",
+    )
+    return (
+        "Using the previous frame as reference, generate the next GUI frame.\n\n"
+        f"Step {step_num} action: {action}\n"
+        f"Expected visual state:\n{visual_description}\n\n"
+        "Requirements:\n"
+        "- Only change UI regions affected by this action.\n"
+        "- Preserve persistent bars, layout, and style.\n"
+        "- Keep text and icons readable.\n"
+    )
+
+
+def _type5_prompt(metadata: dict[str, Any]) -> str:
+    grounding = metadata.get("grounding") or {}
+    explanation = _text_or_default(
+        metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
+        "Predict the immediate GUI reaction to the indicated target.",
+    )
+    return (
+        "Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
+        f"Expected effect: {explanation}\n"
+        f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
+        "Requirements:\n"
+        "- Apply only the interaction-triggered change.\n"
+        "- Preserve unrelated regions.\n"
+        "- Keep the UI realistic and readable.\n"
+    )
+
+
+def _make_storyboard_image(
+    frames: list[Image.Image],
+    *,
+    columns: int = 3,
+    background_color: tuple[int, int, int] = (255, 255, 255),
+) -> Image.Image:
+    if not frames:
+        raise ValueError("Expected at least one frame to build a storyboard image.")
+
+    normalized = [frame.convert("RGB") for frame in frames]
+    frame_width = max(frame.width for frame in normalized)
+    frame_height = max(frame.height for frame in normalized)
+    rows = (len(normalized) + columns - 1) // columns
+    storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
+
+    for index, frame in enumerate(normalized):
+        x_offset = (index % columns) * frame_width
+        y_offset = (index // columns) * frame_height
+        if frame.size != (frame_width, frame_height):
+            frame = frame.resize((frame_width, frame_height))
+        storyboard.paste(frame, (x_offset, y_offset))
+    return storyboard
+
+
+def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
+    storyboard = _make_storyboard_image(frames, columns=3)
+    prompt_suffix = (
+        "The attached image is a storyboard containing six frames arranged left-to-right, "
+        "top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
+    )
+    return prompt_suffix, [storyboard]
+
+
+class LocalJudgeClient:
+    def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
+        self.base_url = base_url.rstrip("/")
+        self.api_key = api_key
+        self.model = model
+        self.timeout = timeout
+
+    def _build_scoring_prompt(self, task_prompt: str) -> str:
+        return (
+            "You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
+            "Evaluate whether the generated image(s) satisfy the task.\n\n"
+            "Score these five dimensions from 0 to 5:\n"
+            "- goal: whether the user goal is completed correctly\n"
+            "- logic: whether the transition/state change is logically correct\n"
+            "- cons: whether unrelated regions remain consistent\n"
+            "- ui: whether the UI layout/components remain realistic and coherent\n"
+            "- qual: whether the images are visually clear and artifact-free\n\n"
+            "Return JSON only. Do not add any prose outside JSON.\n"
+            "Use exactly this schema:\n"
+            "{\n"
+            '  "goal": 0,\n'
+            '  "logic": 0,\n'
+            '  "cons": 0,\n'
+            '  "ui": 0,\n'
+            '  "qual": 0,\n'
+            '  "reasoning": "short explanation"\n'
+            "}\n\n"
+            "Scoring task:\n"
+            f"{task_prompt}"
+        )
+
+    def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
+        content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
+        for image in images:
+            content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
+
+        response = requests.post(
+            build_openai_url(self.base_url, "/chat/completions"),
+            json={
+                "model": self.model,
+                "messages": [{"role": "user", "content": content}],
+                "temperature": 0,
+            },
+            headers={
+                "Authorization": f"Bearer {self.api_key}",
+                "Content-Type": "application/json",
+            },
+            timeout=self.timeout,
+        )
+        response.raise_for_status()
+        message_content = response.json()["choices"][0]["message"]["content"]
+        if isinstance(message_content, list):
+            return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
+        return str(message_content)
+
+    def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
+        primary_prompt = self._build_scoring_prompt(prompt)
+        text = self._request_text(primary_prompt, images)
+        try:
+            return extract_json_object(text)
+        except ValueError:
+            retry_prompt = (
+                self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
+                "Return only the JSON object with integer scores."
+            )
+            retry_text = self._request_text(retry_prompt, images)
+            try:
+                return extract_json_object(retry_text)
+            except ValueError:
+                return {
+                    "goal": 0,
+                    "logic": 0,
+                    "cons": 0,
+                    "ui": 0,
+                    "qual": 0,
+                    "reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
+                }
+
+
+class GEBenchRunner:
+    def __init__(
+        self,
+        *,
+        dataset_root: Path,
+        output_root: Path,
+        base_url: str,
+        model: str,
+        api_key: str = "EMPTY",
+        width: int = 768,
+        height: int = 576,
+        num_inference_steps: int = 8,
+        output_compression: int | None = 98,
+        guidance_scale: float | None = None,
+        seed: int | None = 42,
+    ):
+        self.dataset_root = dataset_root
+        self.output_root = output_root
+        self.model = model
+        self.width = width
+        self.height = height
+        self.num_inference_steps = num_inference_steps
+        self.output_compression = output_compression
+        self.guidance_scale = guidance_scale
+        self.seed = seed
+        self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
+
+    def generate(
+        self,
+        *,
+        data_type: str,
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_type: int | None = None,
+    ) -> list[dict[str, Any]]:
+        sample_specs = select_balanced_gebench_samples(
+            {data_type: _iter_sample_specs(self.dataset_root, data_type)},
+            samples_per_type=samples_per_type,
+        )[data_type]
+        if max_samples is not None:
+            sample_specs = sample_specs[:max_samples]
+
+        results: list[dict[str, Any]] = []
+        if workers <= 1:
+            for sample_spec in sample_specs:
+                result = self._generate_one(data_type, sample_spec)
+                if result:
+                    results.append(result)
+            return results
+
+        with ThreadPoolExecutor(max_workers=workers) as executor:
+            futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
+            for future in as_completed(futures):
+                result = future.result()
+                if result:
+                    results.append(result)
+        return results
+
+    def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
+        sample_path = sample_spec.sample_path
+        metadata = sample_spec.metadata
+        lang_device = sample_spec.lang_device
+        sample_name = sample_spec.sample_name
+        output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
+
+        if data_type == "type1":
+            output_path = output_dir / "generated.png"
+            if output_path.exists():
+                return {
+                    "data_type": data_type,
+                    "sample_name": f"{lang_device}/{sample_name}",
+                    "output_path": str(output_path),
+                }
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
+            )
+            if source is None:
+                return None
+            generated = self.client.generate_image_edit(
+                model=self.model,
+                prompt=_type1_prompt(metadata),
+                images=source,
+                width=self.width,
+                height=self.height,
+                num_inference_steps=self.num_inference_steps,
+                output_compression=self.output_compression,
+                guidance_scale=self.guidance_scale,
+                seed=self.seed,
+            )
+            save_image(output_path, generated)
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        if data_type == "type2":
+            goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
+            )
+            if source is None:
+                return None
+            frame0_path = output_dir / "frame0.png"
+            if not frame0_path.exists():
+                save_image(frame0_path, source)
+            previous = source
+            for step_num in range(1, 6):
+                frame_path = output_dir / f"frame{step_num}.png"
+                if frame_path.exists():
+                    previous = Image.open(frame_path).convert("RGB")
+                    continue
+                generated = self.client.generate_image_edit(
+                    model=self.model,
+                    prompt=_type2_prompt(goal, step_num),
+                    images=previous,
+                    width=self.width,
+                    height=self.height,
+                    num_inference_steps=self.num_inference_steps,
+                    output_compression=self.output_compression,
+                    guidance_scale=self.guidance_scale,
+                    seed=self.seed,
+                )
+                save_image(frame_path, generated)
+                previous = generated
+            output_path = output_dir / "frame5.png"
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        if data_type in {"type3", "type4"}:
+            steps = _trajectory_steps(metadata)
+            frame0_path = output_dir / "frame0.png"
+            if frame0_path.exists():
+                previous = Image.open(frame0_path).convert("RGB")
+            else:
+                previous = self.client.generate_text_to_image(
+                    model=self.model,
+                    prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
+                    width=self.width,
+                    height=self.height,
+                    num_inference_steps=self.num_inference_steps,
+                    output_compression=self.output_compression,
+                    guidance_scale=self.guidance_scale,
+                    seed=self.seed,
+                )
+                save_image(frame0_path, previous)
+
+            for step_num in range(1, 6):
+                frame_path = output_dir / f"frame{step_num}.png"
+                if frame_path.exists():
+                    previous = Image.open(frame_path).convert("RGB")
+                    continue
+                step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
+                generated = self.client.generate_image_edit(
+                    model=self.model,
+                    prompt=_type34_next_prompt(step_num, step_info),
+                    images=previous,
+                    width=self.width,
+                    height=self.height,
+                    num_inference_steps=self.num_inference_steps,
+                    output_compression=self.output_compression,
+                    guidance_scale=self.guidance_scale,
+                    seed=self.seed,
+                )
+                save_image(frame_path, generated)
+                previous = generated
+            output_path = output_dir / "frame5.png"
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        if data_type == "type5":
+            output_path = output_dir / "generated.png"
+            if output_path.exists():
+                return {
+                    "data_type": data_type,
+                    "sample_name": f"{lang_device}/{sample_name}",
+                    "output_path": str(output_path),
+                }
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
+            )
+            if source is None:
+                return None
+            generated = self.client.generate_image_edit(
+                model=self.model,
+                prompt=_type5_prompt(metadata),
+                images=source,
+                width=self.width,
+                height=self.height,
+                num_inference_steps=self.num_inference_steps,
+                output_compression=self.output_compression,
+                guidance_scale=self.guidance_scale,
+                seed=self.seed,
+            )
+            save_image(output_path, generated)
+            return {
+                "data_type": data_type,
+                "sample_name": f"{lang_device}/{sample_name}",
+                "output_path": str(output_path),
+            }
+
+        raise ValueError(f"Unsupported data type: {data_type}")
+
+
+class GEBenchEvaluator:
+    def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
+        self.dataset_root = dataset_root
+        self.output_root = output_root
+        self.judge = judge
+
+    def evaluate(
+        self,
+        *,
+        data_type: str,
+        workers: int = 1,
+        max_samples: int | None = None,
+        samples_per_type: int | None = None,
+    ) -> dict[str, Any]:
+        output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
+        sample_specs_by_name = {
+            (spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
+        }
+        if not output_type_dir.exists():
+            payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
+            write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
+            return payload
+        sample_dirs = [
+            sample_dir
+            for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
+            for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
+            if (lang_dir.name, sample_dir.name) in sample_specs_by_name
+        ]
+        sample_dirs = select_balanced_gebench_samples(
+            {data_type: sample_dirs},
+            samples_per_type=samples_per_type,
+        )[data_type]
+        if max_samples is not None:
+            sample_dirs = sample_dirs[:max_samples]
+        results: list[dict[str, Any]] = []
+        if workers <= 1:
+            for sample_dir in sample_dirs:
+                result = self._evaluate_one(
+                    data_type,
+                    sample_dir,
+                    sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
+                )
+                if result:
+                    results.append(result)
+        else:
+            with ThreadPoolExecutor(max_workers=workers) as executor:
+                futures = [
+                    executor.submit(
+                        self._evaluate_one,
+                        data_type,
+                        sample_dir,
+                        sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
+                    )
+                    for sample_dir in sample_dirs
+                ]
+                for future in as_completed(futures):
+                    result = future.result()
+                    if result:
+                        results.append(result)
+
+        payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
+        write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
+        return payload
+
+    def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
+        lang_device = sample_dir.parent.name
+        sample_name = sample_dir.name
+        dataset_sample = sample_spec.sample_path
+        metadata = sample_spec.metadata
+
+        if data_type == "type1":
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
+            )
+            generated_path = find_first_image(sample_dir)
+            if source is None or generated_path is None:
+                return None
+            generated = Image.open(generated_path).convert("RGB")
+            raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
+        elif data_type == "type2":
+            frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
+            goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
+            prompt_suffix, judge_images = _trajectory_judge_payload(frames)
+            raw_scores = self.judge.evaluate(
+                prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
+                images=judge_images,
+            )
+        elif data_type in {"type3", "type4"}:
+            frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
+            instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
+            prompt_suffix, judge_images = _trajectory_judge_payload(frames)
+            raw_scores = self.judge.evaluate(
+                prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
+                images=judge_images,
+            )
+        elif data_type == "type5":
+            source = _resolve_referenced_image(
+                metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
+            )
+            generated_path = find_first_image(sample_dir)
+            if source is None or generated_path is None:
+                return None
+            generated = Image.open(generated_path).convert("RGB")
+            raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
+        else:
+            raise ValueError(f"Unsupported data type: {data_type}")
+
+        scores = _normalize_scores(raw_scores)
+        return {
+            "sample_name": f"{lang_device}/{sample_name}",
+            "data_type": data_type,
+            "scores": scores,
+            "overall": _compute_overall(scores),
+            "raw_scores": raw_scores,
+        }
+
+
+def _data_types_arg(value: str) -> list[str]:
+    return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
+    subparsers = parser.add_subparsers(dest="command", required=True)
+
+    generate = subparsers.add_parser("generate")
+    generate.add_argument("--dataset-root", type=Path, required=True)
+    generate.add_argument("--output-root", type=Path, required=True)
+    generate.add_argument("--base-url", type=str, required=True)
+    generate.add_argument("--model", type=str, required=True)
+    generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
+    generate.add_argument("--api-key", type=str, default="EMPTY")
+    generate.add_argument("--width", type=int, default=768)
+    generate.add_argument("--height", type=int, default=576)
+    generate.add_argument("--num-inference-steps", type=int, default=8)
+    generate.add_argument("--output-compression", type=int, default=98)
+    generate.add_argument("--guidance-scale", type=float, default=None)
+    generate.add_argument("--seed", type=int, default=42)
+    generate.add_argument("--workers", type=int, default=1)
+    generate.add_argument("--max-samples", type=int, default=None)
+    generate.add_argument("--samples-per-type", type=int, default=None)
+
+    evaluate = subparsers.add_parser("evaluate")
+    evaluate.add_argument("--dataset-root", type=Path, required=True)
+    evaluate.add_argument("--output-root", type=Path, required=True)
+    evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
+    evaluate.add_argument("--judge-base-url", type=str, required=True)
+    evaluate.add_argument("--judge-model", type=str, required=True)
+    evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
+    evaluate.add_argument("--workers", type=int, default=1)
+    evaluate.add_argument("--max-samples", type=int, default=None)
+    evaluate.add_argument("--samples-per-type", type=int, default=None)
+
+    summarize = subparsers.add_parser("summarize")
+    summarize.add_argument("--output-root", type=Path, required=True)
+
+    return parser
+
+
+def main(argv: list[str] | None = None) -> int:
+    parser = build_parser()
+    args = parser.parse_args(argv)
+
+    if args.command == "generate":
+        runner = GEBenchRunner(
+            dataset_root=args.dataset_root,
+            output_root=args.output_root,
+            base_url=args.base_url,
+            model=args.model,
+            api_key=args.api_key,
+            width=args.width,
+            height=args.height,
+            num_inference_steps=args.num_inference_steps,
+            output_compression=args.output_compression,
+            guidance_scale=args.guidance_scale,
+            seed=args.seed,
+        )
+        records: list[dict[str, Any]] = []
+        for data_type in _data_types_arg(args.data_type):
+            records.extend(
+                runner.generate(
+                    data_type=data_type,
+                    workers=args.workers,
+                    max_samples=args.max_samples,
+                    samples_per_type=args.samples_per_type,
+                )
+            )
+        payload = {"records": records, "summary": summarize_generated_records(records)}
+        write_json(args.output_root / "generation_manifest.json", payload)
+        return 0
+
+    if args.command == "evaluate":
+        judge = LocalJudgeClient(
+            base_url=args.judge_base_url,
+            api_key=args.judge_api_key,
+            model=args.judge_model,
+        )
+        evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
+        combined_results: list[dict[str, Any]] = []
+        for data_type in _data_types_arg(args.data_type):
+            payload = evaluator.evaluate(
+                data_type=data_type,
+                workers=args.workers,
+                max_samples=args.max_samples,
+                samples_per_type=args.samples_per_type,
+            )
+            combined_results.extend(payload["results"])
+        _write_json_with_timestamp(
+            args.output_root / "evaluations" / "summary.json",
+            {"summary": summarize_gebench_results(combined_results)},
+        )
+        return 0
+
+    if args.command == "summarize":
+        generation_summary = collect_gebench_generation_summary(args.output_root)
+        evaluation_dir = args.output_root / "evaluations"
+        result_records: list[dict[str, Any]] = []
+        if evaluation_dir.exists():
+            for file_path in sorted(evaluation_dir.glob("type*.json")):
+                payload = load_json(file_path)
+                result_records.extend(payload.get("results", []))
+        payload: dict[str, Any] = {"generation": generation_summary}
+        if result_records:
+            payload["evaluation"] = summarize_gebench_results(result_records)
+        _write_json_with_timestamp(args.output_root / "summary.json", payload)
+        print(json.dumps(payload, indent=2, ensure_ascii=False))
+        return 0
+
+    parser.error(f"Unknown command: {args.command}")
+    return 1

benchmarks/accuracy/text_to_image/run_gebench.py

@@ -0,0 +1,13 @@
+# ruff: noqa: E402, I001
+import sys
+from pathlib import Path
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+if str(REPO_ROOT) not in sys.path:
+    sys.path.insert(0, str(REPO_ROOT))
+
+from benchmarks.accuracy.text_to_image.gbench import main
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

benchmarks/build_dataset/download_process_data_seedtts.md

@@ -0,0 +1,82 @@
+# Benchmark Dataset Preparation Guide
+
+This guide describes how to download and prepare the SeedTTS test dataset for benchmarking Qwen-Omni models.
+
+## Prerequisites
+
+- Python 3.8+
+- `gdown` for downloading from Google Drive
+- Access to the benchmark scripts
+
+## Steps
+
+### 1. Navigate to the Dataset Directory
+
+```bash
+cd benchmarks/build_dataset
+```
+
+### 2. Install Dependencies
+
+```bash
+pip install gdown
+```
+
+### 3. Download the SeedTTS Test Dataset
+
+Download the dataset from Google Drive:
+
+```bash
+gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+```
+
+### 4. Extract the Dataset
+
+```bash
+tar -xf seedtts_testset.tar
+```
+
+### 5. Prepare the Metadata File
+
+Copy the English metadata file to the working directory:
+
+```bash
+cp seedtts_testset/en/meta.lst meta.lst
+```
+
+### 6. Extract Prompts
+
+Extract the first N prompts from the metadata file:
+
+```bash
+# Extract top 100 prompts (adjust -n for different amounts)
+python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
+```
+
+**Options:**
+- `-i, --input`: Input metadata file (default: `meta.lst`)
+- `-o, --output`: Output prompts file (default: `prompts.txt`)
+- `-n, --num-lines`: Number of prompts to extract (required)
+
+### 7. Clean Up (Optional)
+
+Remove temporary files to save disk space:
+
+```bash
+rm -rf seedtts_testset
+rm seedtts_testset.tar
+rm meta.lst
+```
+
+## Quick Start (All-in-One)
+
+```bash
+# Full setup and benchmark
+cd benchmarks/build_dataset
+pip install gdown
+gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+tar -xf seedtts_testset.tar
+cp seedtts_testset/en/meta.lst meta.lst
+python extract_tts_prompts.py -i meta.lst -o top100.txt -n 100
+rm -rf seedtts_testset seedtts_testset.tar meta.lst
+```

benchmarks/build_dataset/extract_tts_prompts.py

@@ -0,0 +1,73 @@
+#!/usr/bin/env python3
+"""
+Extract prompts from meta.lst and save them to a txt file.
+
+Each line in meta.lst has the format:
+ID|prompt_text|audio_path|target_text
+
+This script extracts the prompt_text (second field) from the first N lines.
+"""
+
+import argparse
+from pathlib import Path
+
+
+def extract_prompts(input_file: str, output_file: str, num_lines: int) -> None:
+    """
+    Extract prompts from meta.lst and save to output file.
+
+    Args:
+        input_file: Path to the meta.lst file
+        output_file: Path to the output txt file
+        num_lines: Number of lines to process
+    """
+    prompts = []
+
+    with open(input_file, encoding="utf-8") as f:
+        for i, line in enumerate(f):
+            if i >= num_lines:
+                break
+
+            line = line.strip()
+            if not line:  # Skip empty lines
+                continue
+
+            parts = line.split("|")
+            if len(parts) >= 2:
+                prompt = parts[1]  # The prompt is the second field
+                prompts.append(prompt)
+
+    # Write prompts to output file
+    with open(output_file, "w", encoding="utf-8") as f:
+        for prompt in prompts:
+            f.write(prompt + "\n")
+
+    # Print result stats
+    print(f"Extracted {len(prompts)} prompts from first {num_lines} lines")
+    print(f"Saved to: {output_file}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Extract prompts from meta.lst file")
+    parser.add_argument(
+        "-i", "--input", type=str, default="meta.lst", help="Input meta.lst file path (default: meta.lst)"
+    )
+    parser.add_argument(
+        "-o", "--output", type=str, default="prompts.txt", help="Output txt file path (default: prompts.txt)"
+    )
+    parser.add_argument(
+        "-n", "--num-lines", type=int, required=True, help="Number of lines to extract from the beginning"
+    )
+
+    args = parser.parse_args()
+
+    # Check if input file exists
+    if not Path(args.input).exists():
+        print(f"Error: Input file '{args.input}' not found")
+        return
+
+    extract_prompts(args.input, args.output, args.num_lines)
+
+
+if __name__ == "__main__":
+    main()

benchmarks/diffusion/README.md

@@ -0,0 +1,138 @@
+
+# Diffusion Serving Benchmark (Image/Video)
+
+This folder contains an online-serving benchmark script for diffusion models.
+It sends requests to a vLLM OpenAI-compatible endpoint and reports throughput,
+latency percentiles, and optional SLO attainment.
+
+The main entrypoint is:
+
+- `benchmarks/diffusion/diffusion_benchmark_serving.py`
+
+## 1. Quick Start
+
+1. Start the server:
+
+```bash
+vllm serve Qwen/Qwen-Image --omni --port 8099
+```
+
+2. Run a minimal benchmark:
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+	--base-url http://localhost:8099 \
+	--model Qwen/Qwen-Image \
+	--task t2i \
+	--dataset vbench \
+	--num-prompts 5
+```
+
+**Notes**
+
+- The benchmark talks to `http://<host>:<port>/v1/chat/completions`.
+- If you run the server on another host or port, pass `--base-url` accordingly.
+
+## 2. Supported Datasets
+
+The benchmark supports three dataset modes via `--dataset`:
+
+- `vbench`: Built-in prompt/data loader.
+- `trace`: Heterogeneous request traces (each request can have different resolution/frames/steps).
+- `random`: Synthetic prompts for quick smoke tests.
+
+### VBench dataset
+
+`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
+per-request generation fields. In this mode, all requests share CLI values:
+`--width --height --num-frames --fps --num-inference-steps`
+(pass `--width` and `--height` together).
+
+Example (`t2v`):
+
+```bash
+python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
+	--base-url http://localhost:8099 \
+	--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+	--task t2v \
+	--dataset vbench \
+	--num-prompts 50 \
+	--width 640 --height 480 \
+	--num-frames 81 --fps 16 \
+	--num-inference-steps 40
+```
+
+Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
+
+If you use i2v/i2i bench datasets and need auto-download support, you may need:
+
+```bash
+uv pip install gdown
+```
+
+### Trace dataset
+
+Use `--dataset trace` to replay a trace file. The trace can specify per-request fields such as:
+
+- `width`, `height`
+- `num_frames` (video)
+- `num_inference_steps`
+- `seed`, `fps`
+- optional `slo_ms` (per-request SLO target)
+
+By default (when `--dataset-path` is not provided), the script downloads a default trace from
+the HuggingFace dataset repo `asukaqaqzz/Dit_Trace`. The default filename can depend on `--task`
+(e.g., `t2v` uses a video trace).
+
+Current defaults:
+
+- `--task t2i` -> `sd3_trace.txt`
+- `--task t2v` -> `cogvideox_trace.txt`
+
+You can point to your own trace using `--dataset-path`.
+
+## 3. Benchmark Parameters
+
+### Basic flags
+
+- `--base-url`: Server address (the script calls `.../v1/chat/completions`).
+- `--model`: The OpenAI-compatible `model` field.
+- `--task`: Task type (e.g., `t2i`, `t2v`, `i2i`, `i2v`).
+- `--dataset`: Dataset mode (`vbench` / `trace` / `random`).
+- `--num-prompts`: Number of requests to send.
+
+Common optional flags:
+
+- `--output-file`: Write metrics to a JSON file.
+- `--disable-tqdm`: Disable the progress bar.
+
+### Resolution / frames / steps: CLI defaults vs dataset fields
+
+Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
+
+- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
+- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
+
+Precedence rules for `trace` (i.e., what actually gets sent):
+
+- `width/height`: if either `--width` or `--height` is explicitly set, it overrides per-request values from the trace; otherwise per-request values are used when present.
+- `num_frames`: per-request `num_frames` takes precedence; otherwise fall back to `--num-frames`.
+- `num_inference_steps`: per-request `num_inference_steps` takes precedence; otherwise fall back to `--num-inference-steps`.
+
+### SLO, warmup, and max concurrency
+
+Enable SLO evaluation with `--slo`.
+
+- If a request in the trace already has `slo_ms`, that value is used.
+- Otherwise, the script runs warmup requests to infer a base unit time, estimates `expected_ms` by linearly scaling with area/frames/steps, and then sets `slo_ms = expected_ms * --slo-scale`.
+
+Warmup flags:
+
+- `--warmup-requests`: Number of warmup requests.
+- `--warmup-num-inference-steps`: Steps used during warmup.
+- For `--task t2v`: warmup requests are forced to use `num_frames=1` to make warmup faster and less noisy.
+
+Traffic / concurrency flags:
+
+- `--request-rate`: Target request rate (requests/second). If set to `inf`, the script sends all requests immediately.
+- `--max-concurrency`: Max number of in-flight requests (default: `1`). This can hard-cap the achieved QPS: if it is too small, requests will queue behind the semaphore, and both achieved throughput and observed SLO attainment can be skewed.

benchmarks/diffusion/backends.py

@@ -0,0 +1,345 @@
+import asyncio
+import base64
+import mimetypes
+import os
+import time
+import uuid
+from dataclasses import dataclass, field
+from typing import Any
+
+import aiohttp
+from tqdm import tqdm
+
+
+@dataclass
+class RequestFuncInput:
+    prompt: str
+    api_url: str
+    model: str
+    width: int | None = None
+    height: int | None = None
+    num_frames: int | None = None
+    num_inference_steps: int | None = None
+    seed: int | None = None
+    fps: int | None = None
+    timestamp: float | None = None
+    slo_ms: float | None = None
+    extra_body: dict[str, Any] = field(default_factory=dict)
+    image_paths: list[str] | None = None
+    request_id: str = field(default_factory=lambda: str(uuid.uuid4()))
+
+
+@dataclass
+class RequestFuncOutput:
+    success: bool = False
+    latency: float = 0.0
+    error: str = ""
+    start_time: float = 0.0
+    response_body: dict[str, Any] = field(default_factory=dict)
+    stage_durations: dict[str, float] = field(default_factory=dict)
+    peak_memory_mb: float = 0.0
+    slo_achieved: bool | None = None
+
+
+def _guess_mime_type(path: str) -> str:
+    mime, _ = mimetypes.guess_type(path)
+    return mime or "application/octet-stream"
+
+
+def _encode_image_as_data_url(path: str) -> str:
+    with open(path, "rb") as f:
+        encoded = base64.b64encode(f.read()).decode("utf-8")
+    mime = _guess_mime_type(path)
+    return f"data:{mime};base64,{encoded}"
+
+
+async def async_request_chat_completions(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+    enable_diffusion_pipeline_profiler: bool = False,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    extra_body = dict(input.extra_body)
+    if input.width and input.height:
+        extra_body.setdefault("height", input.height)
+        extra_body.setdefault("width", input.width)
+    if input.num_frames:
+        extra_body.setdefault("num_frames", input.num_frames)
+    if input.num_inference_steps:
+        extra_body.setdefault("num_inference_steps", input.num_inference_steps)
+    if input.seed is not None:
+        extra_body.setdefault("seed", input.seed)
+    if input.fps:
+        extra_body.setdefault("fps", input.fps)
+
+    if input.image_paths and len(input.image_paths) > 0:
+        content = []
+        if input.prompt:
+            content.append({"type": "text", "text": input.prompt})
+        for img_path in input.image_paths:
+            if not os.path.exists(img_path):
+                output.error = f"Image file not found: {img_path}"
+                output.success = False
+                if pbar:
+                    pbar.update(1)
+                return output
+            content.append(
+                {
+                    "type": "image_url",
+                    "image_url": {"url": _encode_image_as_data_url(img_path)},
+                }
+            )
+        messages = [{"role": "user", "content": content}]
+    else:
+        messages = [{"role": "user", "content": input.prompt}]
+
+    payload = {
+        "model": input.model,
+        "messages": messages,
+    }
+    if extra_body:
+        payload["extra_body"] = extra_body
+
+    try:
+        async with session.post(input.api_url, json=payload) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                try:
+                    choices = resp_json.get("choices", [])
+                    if choices and isinstance(choices, list):
+                        msg = choices[0].get("message", {})
+                        if isinstance(msg, dict):
+                            content = msg.get("content", [])
+                            if content and isinstance(content, list) and len(content) > 0:
+                                first_item = content[0]
+                                if isinstance(first_item, dict):
+                                    output.stage_durations = first_item.get("stage_durations") or {}
+                                    output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
+                except (IndexError, TypeError, AttributeError):
+                    pass
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+async def async_request_openai_images(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    """
+    Send request to OpenAI's /v1/images/generations endpoint.
+    """
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    # Build size string from width/height
+    width = input.width or 1024
+    height = input.height or 1024
+    size = f"{width}x{height}"
+
+    payload: dict[str, Any] = {
+        "model": input.model,
+        "prompt": input.prompt,
+        "n": 1,
+        "size": size,
+        "response_format": "b64_json",
+    }
+
+    # Add optional parameters
+    if input.seed is not None:
+        payload["seed"] = input.seed
+    if input.num_inference_steps is not None:
+        payload["num_inference_steps"] = input.num_inference_steps
+
+    # Add any extra body parameters
+    if input.extra_body:
+        for key, value in input.extra_body.items():
+            if key not in payload:
+                payload[key] = value
+
+    headers = {
+        "Content-Type": "application/json",
+        "Authorization": "Bearer EMPTY",
+    }
+
+    try:
+        async with session.post(input.api_url, json=payload, headers=headers) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                output.response_body = resp_json
+                output.success = True
+                # Check for usage/memory info if available
+                if "usage" in resp_json and "peak_memory_mb" in resp_json.get("usage", {}):
+                    output.peak_memory_mb = resp_json["usage"]["peak_memory_mb"]
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+async def async_request_v1_videos(
+    input: RequestFuncInput,
+    session: aiohttp.ClientSession,
+    pbar: tqdm | None = None,
+) -> RequestFuncOutput:
+    output = RequestFuncOutput()
+    output.start_time = time.perf_counter()
+
+    files = dict(input.extra_body)
+    if input.prompt:
+        files.setdefault("prompt", input.prompt)
+    if input.width and input.height:
+        files.setdefault("height", input.height)
+        files.setdefault("width", input.width)
+    if input.num_frames:
+        files.setdefault("num_frames", input.num_frames)
+    if input.num_inference_steps:
+        files.setdefault("num_inference_steps", input.num_inference_steps)
+    if input.seed is not None:
+        files.setdefault("seed", input.seed)
+    if input.fps:
+        files.setdefault("fps", input.fps)
+
+    form = aiohttp.FormData()
+    for k, v in files.items():
+        form.add_field(k, str(v))
+
+    image_file = None
+    if input.image_paths and len(input.image_paths) > 0:
+        image_path = input.image_paths[0]
+        image_file = open(image_path, "rb")
+        form.add_field(
+            "input_reference",
+            image_file,
+            filename=os.path.basename(image_path),
+            content_type="application/octet-stream",
+        )
+
+    job_id = None
+    job_status = None
+    poll_json = {}
+    resp_json = {}
+
+    try:
+        # invoke a post request (POST /v1/videos)
+        async with session.post(input.api_url, data=form) as response:
+            if response.status == 200:
+                resp_json = await response.json()
+                job_id = resp_json.get("id")
+                job_status = resp_json.get("status")
+                if not job_id or not job_status:
+                    output.error = "API response missing job 'id' or 'status' field."
+                    output.success = False
+                    return output
+            else:
+                output.error = f"HTTP {response.status}: {await response.text()}"
+                output.success = False
+                return output
+
+        # invoke a poll request (GET /v1/videos/{video_id})
+        poll_interval = 2.0  # Unit(s)
+        timeout_seconds = 600.0
+        deadline = time.perf_counter() + timeout_seconds
+        job_url = f"{input.api_url}/{job_id}"
+
+        while job_status not in {"completed", "failed"}:
+            await asyncio.sleep(poll_interval)
+
+            async with session.get(job_url) as poll_response:
+                if poll_response.status != 200:
+                    output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
+                    output.success = False
+                    return output
+
+                poll_json = await poll_response.json()
+                job_status = poll_json.get("status")
+
+                if time.perf_counter() >= deadline:
+                    output.error = f"Timed out waiting for video job {job_id} to complete."
+                    output.success = False
+                    return output
+
+        if job_status == "failed":
+            output.error = f"Video job failed: {poll_json}"
+            output.success = False
+            return output
+
+        # invoke a get request (GET /v1/videos/{video_id}/content)
+        content_url = f"{job_url}/content"
+        async with session.get(content_url) as content_response:
+            if content_response.status != 200:
+                output.error = (
+                    f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
+                )
+                output.success = False
+                return output
+
+            video_bytes = await content_response.read()
+            output.response_body = video_bytes
+            output.success = True
+            if "peak_memory_mb" in poll_json:
+                output.peak_memory_mb = poll_json["peak_memory_mb"]
+            elif "peak_memory_mb" in resp_json:
+                output.peak_memory_mb = resp_json["peak_memory_mb"]
+    except Exception as e:
+        output.error = str(e)
+        output.success = False
+    finally:
+        if image_file is not None:
+            image_file.close()
+
+        if job_id is not None:
+            try:
+                async with session.delete(f"{input.api_url}/{job_id}") as _:
+                    pass
+            except Exception as e:
+                print(f"Failed to clean up video job {job_id}: {e}")
+
+    output.latency = time.perf_counter() - output.start_time
+
+    if output.success and input.slo_ms is not None:
+        output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
+
+    if pbar:
+        pbar.update(1)
+    return output
+
+
+backends_function_mapping = {
+    "2i": {
+        "vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
+        "openai": (async_request_openai_images, "/v1/images/generations"),
+    },
+    "2v": {
+        "v1/videos": (async_request_v1_videos, "/v1/videos"),
+    },
+}

benchmarks/diffusion/performance_dashboard/qwen_image_serving_performance.md

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

benchmarks/diffusion/performance_dashboard/wan_2_2_serving_performance.md

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

benchmarks/diffusion/quantization_quality.py

@@ -0,0 +1,460 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Benchmark quantization quality loss for diffusion models (image & video).
+
+Generates outputs with BF16 (baseline) and a quantized config using the same
+seed, then computes LPIPS perceptual distance between them. Results are printed
+as a Markdown table ready to paste into a PR description.
+
+Requirements:
+    pip install lpips Pillow numpy
+
+Image example (text-to-image):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --task t2i \
+        --quantization fp8 \
+        --prompts \
+            "an aerial view of a coral reef with crystal clear turquoise water" \
+            "a campfire in a dark forest with sparks rising into a starry sky" \
+            "a gourmet dessert plate with chocolate mousse and gold leaf" \
+        --height 1024 --width 1024 \
+        --num-inference-steps 50 --seed 42
+
+Video example (text-to-video):
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
+        --task t2v \
+        --quantization fp8 \
+        --prompts \
+            "A serene lakeside sunrise with mist over the water" \
+            "A cat walking across a wooden bridge in autumn" \
+        --height 720 --width 1280 \
+        --num-frames 81 --num-inference-steps 40 --seed 42
+
+Multiple quantization methods:
+    python benchmarks/diffusion/quantization_quality.py \
+        --model Tongyi-MAI/Z-Image-Turbo \
+        --task t2i \
+        --quantization fp8 int8 bitsandbytes \
+        --prompts "a cup of coffee on the table" \
+        --height 1024 --width 1024 \
+        --num-inference-steps 50 --seed 42
+
+Output directory structure (--output-dir, default: ./quant_bench_output):
+    quant_bench_output/
+        baseline/           # BF16 outputs
+        <method>/           # Quantized outputs per method
+        results.md          # Markdown table
+"""
+
+import argparse
+import gc
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+
+
+def compute_lpips_images(
+    baseline_images: list,
+    quantized_images: list,
+    net: str = "alex",
+) -> list[float]:
+    """Compute LPIPS between paired lists of PIL images."""
+    import lpips
+    from torchvision import transforms
+
+    loss_fn = lpips.LPIPS(net=net).eval()
+    if torch.cuda.is_available():
+        loss_fn = loss_fn.cuda()
+
+    transform = transforms.Compose(
+        [
+            transforms.Resize((256, 256)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
+        ]
+    )
+
+    scores = []
+    for img_bl, img_qt in zip(baseline_images, quantized_images):
+        t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
+        t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
+        if torch.cuda.is_available():
+            t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
+        with torch.no_grad():
+            score = loss_fn(t_bl, t_qt).item()
+        scores.append(score)
+    return scores
+
+
+def compute_lpips_video(
+    baseline_frames: np.ndarray,
+    quantized_frames: np.ndarray,
+    net: str = "alex",
+) -> float:
+    """Compute mean per-frame LPIPS for a video pair.
+
+    Args:
+        baseline_frames: (F, H, W, C) float array in [0, 1].
+        quantized_frames: same shape.
+
+    Returns:
+        Mean LPIPS across all frames.
+    """
+    import lpips
+
+    loss_fn = lpips.LPIPS(net=net).eval()
+    if torch.cuda.is_available():
+        loss_fn = loss_fn.cuda()
+
+    num_frames = min(len(baseline_frames), len(quantized_frames))
+    scores = []
+    for i in range(num_frames):
+        # Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
+        f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
+        f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
+        if torch.cuda.is_available():
+            f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
+        with torch.no_grad():
+            score = loss_fn(f_bl, f_qt).item()
+        scores.append(score)
+    return float(np.mean(scores))
+
+
+def _build_omni_kwargs(args, quantization=None):
+    """Build kwargs dict for Omni() constructor."""
+    from vllm_omni.diffusion.data import DiffusionParallelConfig
+
+    parallel_config = DiffusionParallelConfig(
+        ulysses_degree=args.ulysses_degree,
+        ring_degree=args.ring_degree,
+        tensor_parallel_size=args.tensor_parallel_size,
+    )
+    kwargs = {
+        "model": args.model,
+        "parallel_config": parallel_config,
+        "enforce_eager": args.enforce_eager,
+    }
+    if quantization:
+        kwargs["quantization_config"] = quantization
+    return kwargs
+
+
+def _generate_image(omni, args, prompt, seed):
+    """Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+    torch.cuda.reset_peak_memory_stats()
+    start = time.perf_counter()
+    outputs = omni.generate(
+        {"prompt": prompt},
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            num_inference_steps=args.num_inference_steps,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+    first = outputs[0]
+    req_out = first.request_output[0] if hasattr(first, "request_output") else first
+    img = req_out.images[0]
+    return img, elapsed, peak_mem
+
+
+def _generate_video(omni, args, prompt, seed):
+    """Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
+    from vllm_omni.inputs.data import OmniDiffusionSamplingParams
+    from vllm_omni.outputs import OmniRequestOutput
+    from vllm_omni.platforms import current_omni_platform
+
+    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
+    torch.cuda.reset_peak_memory_stats()
+    start = time.perf_counter()
+    outputs = omni.generate(
+        {"prompt": prompt, "negative_prompt": ""},
+        OmniDiffusionSamplingParams(
+            height=args.height,
+            width=args.width,
+            generator=generator,
+            guidance_scale=args.guidance_scale,
+            num_inference_steps=args.num_inference_steps,
+            num_frames=args.num_frames,
+        ),
+    )
+    elapsed = time.perf_counter() - start
+    peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
+
+    first = outputs[0]
+    if hasattr(first, "request_output") and isinstance(first.request_output, list):
+        inner = first.request_output[0]
+        if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
+            frames = inner.images[0] if inner.images else None
+        else:
+            frames = inner
+    elif hasattr(first, "images") and first.images:
+        frames = first.images
+    else:
+        raise ValueError("Could not extract video frames from output.")
+
+    if isinstance(frames, torch.Tensor):
+        video = frames.detach().cpu()
+        if video.dim() == 5:
+            video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
+        elif video.dim() == 4 and video.shape[0] in (3, 4):
+            video = video.permute(1, 2, 3, 0)
+        if video.is_floating_point():
+            video = video.clamp(-1, 1) * 0.5 + 0.5
+        frames_array = video.float().numpy()
+    else:
+        frames_array = np.asarray(frames)
+        if frames_array.ndim == 5:
+            frames_array = frames_array[0]
+
+    return frames_array, elapsed, peak_mem
+
+
+def _unload_omni(omni):
+    """Delete Omni instance and free GPU memory."""
+    del omni
+    gc.collect()
+    if torch.cuda.is_available():
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+
+def run_benchmark(args):
+    from vllm_omni.entrypoints.omni import Omni
+
+    output_dir = Path(args.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    is_video = args.task == "t2v"
+    prompts = args.prompts
+    seed = args.seed
+
+    # Determine configs to benchmark
+    configs = []  # list of (label, quantization_method)
+    for method in args.quantization:
+        configs.append((method, method))
+
+    # --- Baseline run ---
+    print("\n" + "=" * 60)
+    print("Running BF16 baseline...")
+    print("=" * 60)
+    bl_kwargs = _build_omni_kwargs(args, quantization=None)
+    omni_bl = Omni(**bl_kwargs)
+
+    baseline_outputs = {}  # prompt -> (output, time, mem)
+    for prompt in prompts:
+        print(f"  Generating: {prompt[:60]}...")
+        if is_video:
+            out, t, mem = _generate_video(omni_bl, args, prompt, seed)
+        else:
+            out, t, mem = _generate_image(omni_bl, args, prompt, seed)
+        baseline_outputs[prompt] = (out, t, mem)
+
+    bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
+    bl_mem = baseline_outputs[prompts[0]][2]  # use first prompt's memory
+    _unload_omni(omni_bl)
+
+    # Save baseline outputs
+    bl_dir = output_dir / "baseline"
+    bl_dir.mkdir(parents=True, exist_ok=True)
+    for i, prompt in enumerate(prompts):
+        out = baseline_outputs[prompt][0]
+        if is_video:
+            try:
+                from diffusers.utils import export_to_video
+
+                frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
+                export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
+            except ImportError:
+                np.save(bl_dir / f"prompt_{i}.npy", out)
+        else:
+            out.save(bl_dir / f"prompt_{i}.png")
+
+    # --- Quantized runs ---
+    all_results = []  # list of dicts
+
+    for config_label, quant_method in configs:
+        print(f"\n{'=' * 60}")
+        print(f"Running: {config_label}...")
+        print("=" * 60)
+
+        qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
+        omni_qt = Omni(**qt_kwargs)
+
+        qt_outputs = {}
+        for prompt in prompts:
+            print(f"  Generating: {prompt[:60]}...")
+            if is_video:
+                out, t, mem = _generate_video(omni_qt, args, prompt, seed)
+            else:
+                out, t, mem = _generate_image(omni_qt, args, prompt, seed)
+            qt_outputs[prompt] = (out, t, mem)
+
+        qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
+        qt_mem = qt_outputs[prompts[0]][2]
+        _unload_omni(omni_qt)
+
+        # Save quantized outputs
+        qt_dir = output_dir / config_label.replace(" ", "_")
+        qt_dir.mkdir(parents=True, exist_ok=True)
+
+        # Compute LPIPS per prompt
+        per_prompt = []
+        for i, prompt in enumerate(prompts):
+            bl_out = baseline_outputs[prompt][0]
+            qt_out = qt_outputs[prompt][0]
+            if is_video:
+                lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
+                try:
+                    from diffusers.utils import export_to_video
+
+                    frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
+                    export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
+                except ImportError:
+                    np.save(qt_dir / f"prompt_{i}.npy", qt_out)
+            else:
+                lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
+                qt_out.save(qt_dir / f"prompt_{i}.png")
+            per_prompt.append({"prompt": prompt, "lpips": lpips_score})
+
+        mean_lpips = np.mean([p["lpips"] for p in per_prompt])
+        speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
+        mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
+
+        all_results.append(
+            {
+                "config": config_label,
+                "avg_time": qt_avg_time,
+                "speedup": speedup,
+                "memory_gib": qt_mem,
+                "mem_reduction_pct": mem_reduction,
+                "mean_lpips": mean_lpips,
+                "per_prompt": per_prompt,
+            }
+        )
+
+    # --- Print results ---
+    print("\n\n")
+    print("=" * 80)
+    print("RESULTS")
+    print("=" * 80)
+
+    # Summary table
+    lines = []
+    lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
+    lines.append(
+        f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
+        f"seed={args.seed}, LPIPS ({args.lpips_net})"
+    )
+    if is_video:
+        lines.append(f"Video: {args.num_frames} frames")
+    lines.append("")
+    lines.append("### Summary")
+    lines.append("")
+    lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
+    lines.append("|--------|----------|---------|--------------|---------------|------------|")
+    lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
+    for r in all_results:
+        lines.append(
+            f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
+            f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
+            f"| {r['mean_lpips']:.4f} |"
+        )
+    lines.append("")
+    lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
+    lines.append("")
+
+    # Per-prompt table
+    if len(prompts) > 1:
+        lines.append("### Per-Prompt LPIPS")
+        lines.append("")
+        header = "| Prompt |"
+        sep = "|--------|"
+        for r in all_results:
+            header += f" {r['config']} |"
+            sep += "--------|"
+        lines.append(header)
+        lines.append(sep)
+        for i, prompt in enumerate(prompts):
+            short = prompt[:50] + "..." if len(prompt) > 50 else prompt
+            row = f"| {short} |"
+            for r in all_results:
+                row += f" {r['per_prompt'][i]['lpips']:.4f} |"
+            lines.append(row)
+        lines.append("")
+
+    md = "\n".join(lines)
+    print(md)
+
+    # Save markdown
+    results_path = output_dir / "results.md"
+    results_path.write_text(md, encoding="utf-8")
+    print(f"\nResults saved to {results_path}")
+    print(f"Baseline outputs in {bl_dir}")
+    for r in all_results:
+        qt_dir = output_dir / r["config"].replace(" ", "_")
+        print(f"Quantized outputs in {qt_dir}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="Benchmark quantization quality loss for diffusion models.",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("--model", required=True, help="Model name or local path.")
+    parser.add_argument(
+        "--task",
+        default="t2i",
+        choices=["t2i", "t2v"],
+        help="Task type: t2i (text-to-image) or t2v (text-to-video).",
+    )
+    parser.add_argument(
+        "--quantization",
+        nargs="+",
+        required=True,
+        help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
+    )
+    parser.add_argument(
+        "--prompts",
+        nargs="+",
+        default=["a cup of coffee on the table"],
+        help="One or more prompts to generate.",
+    )
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--height", type=int, default=1024)
+    parser.add_argument("--width", type=int, default=1024)
+    parser.add_argument("--num-inference-steps", type=int, default=50)
+    parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
+    parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
+    parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
+    parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
+    parser.add_argument(
+        "--lpips-net",
+        type=str,
+        default="alex",
+        choices=["alex", "vgg", "squeeze"],
+        help="LPIPS backbone network.",
+    )
+    parser.add_argument("--ulysses-degree", type=int, default=1)
+    parser.add_argument("--ring-degree", type=int, default=1)
+    parser.add_argument("--tensor-parallel-size", type=int, default=1)
+    parser.add_argument("--enforce-eager", action="store_true")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_benchmark(args)

benchmarks/distributed/omni_connectors/README.md

@@ -0,0 +1,397 @@
+# RDMA Test Configuration Guide
+
+This document explains how to configure the RDMA environment and run tests for `MooncakeTransferEngineConnector`.
+
+## Table of Contents
+
+- [Docker Container Permissions](#docker-container-permissions)
+- [Single-Node Testing](#single-node-testing)
+- [Multi-Node Testing](#multi-node-testing)
+- [Running Tests](#running-tests)
+- [Cross-Node Testing](#cross-node-testing)
+- [Troubleshooting](#troubleshooting)
+
+---
+
+## Docker Container Permissions
+
+RDMA tests require access to InfiniBand/RoCE devices and system topology. Add the following permissions when running `docker run`.
+
+### Option 1: Minimal Permissions (Recommended)
+
+```bash
+docker run -it \
+    --cap-add=SYS_PTRACE \
+    --cap-add=IPC_LOCK \
+    --security-opt seccomp=unconfined \
+    --network=host \
+    --device=/dev/infiniband \
+    -v /sys/class/infiniband:/sys/class/infiniband:ro \
+    your-image:tag
+```
+
+Parameter explanation:
+- `--cap-add=SYS_PTRACE`: Allow reading system topology information
+- `--cap-add=IPC_LOCK`: Allow memory locking (required for RDMA memory registration)
+- `--security-opt seccomp=unconfined`: Disable seccomp restrictions
+- `--network=host`: Use host network (required for RDMA)
+- `--device=/dev/infiniband`: Mount InfiniBand devices
+- `-v /sys/class/infiniband`: Mount IB device info (read-only)
+
+### Option 2: Full Permissions (Quick but not recommended for production)
+
+```bash
+docker run -it \
+    --privileged \
+    --network=host \
+    your-image:tag
+```
+
+`--privileged` grants full host permissions. Suitable for quick testing but not recommended for production.
+
+---
+
+## Single-Node Testing
+
+When running single-node tests (producer and consumer on the same machine), ensure they use the **same RDMA device**.
+
+### Problem Background
+
+InfiniBand devices use LID (Local Identifier) for routing. Different devices have different LIDs and cannot communicate directly. If no device is specified, Mooncake may assign different devices to connectors, causing handshake failures.
+
+Common error:
+```
+[Handshake] Failed to modify QP to RTR, check mtu, gid, peer lid, peer qp num: Invalid argument [22]
+```
+
+### Solution
+
+**Method 1: Set Environment Variable (Recommended)**
+
+```bash
+# List available RDMA devices
+ibstat
+
+# Select a device (e.g., mlx5_0)
+export RDMA_DEVICE_NAME='mlx5_0'
+
+# Run tests
+pytest test_mooncake_transfer_engine_rdma.py -v -s
+```
+
+**Method 2: Use RoCE Devices**
+
+If the system has RoCE devices (using IPv4 routing), the test code will automatically detect and prefer them. RoCE device GIDs start with `00:00:00:00:00:00:00:00:00:00:ff:ff` (IPv4-mapped).
+
+**Method 3: Ensure MTU Consistency**
+
+Make sure both endpoints use the same MTU:
+
+```bash
+# Check device MTU
+ibstatus mlx5_0
+```
+
+---
+
+## Multi-Node Testing
+
+For multi-node tests, producer and consumer run on different machines connected via InfiniBand switch.
+
+### Prerequisites
+
+1. Both machines have Mooncake and RDMA drivers installed
+2. Both machines are in the same InfiniBand subnet
+3. Switch is properly configured
+
+### Configuration
+
+**Machine A (Producer):**
+
+```bash
+# Set RDMA host IP (InfiniBand interface IP)
+export RDMA_TEST_HOST='10.0.0.1'
+
+# Optional: Specify device
+export RDMA_DEVICE_NAME='mlx5_0'
+```
+
+**Machine B (Consumer):**
+
+```bash
+# Set RDMA host IP
+export RDMA_TEST_HOST='10.0.0.2'
+
+# Optional: Specify device
+export RDMA_DEVICE_NAME='mlx5_0'
+```
+
+### Verify Connectivity
+
+```bash
+# Ping IB interface
+ping 10.0.0.2
+
+# Test RDMA connectivity with ibping
+# On Machine B (server)
+ibping -S
+
+# On Machine A (client)
+ibping -G <Machine_B_GID>
+```
+
+---
+
+## Running Tests
+
+### Run All RDMA Tests (Single-Node, fast suite)
+
+Slow tests (large payloads, stress, concurrency integrity) are marked `@pytest.mark.slow`. Use `-m "not slow"` to skip them in quick CI or local fast iteration.
+
+```bash
+cd tests/distributed/omni_connectors
+
+# Fast suite only (excludes slow/stress tests)
+pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
+```
+
+### Run Including Slow Tests
+
+```bash
+# Run ALL tests including slow/stress tests
+pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s
+
+# Run ONLY the slow/stress tests
+pytest test_mooncake_transfer_engine_rdma.py test_mooncake_transfer_engine_buffer.py -v -s -m slow
+```
+
+### Run Buffer Management Tests
+
+```bash
+# Fast only
+pytest test_mooncake_transfer_engine_buffer.py -v -s -m "not slow"
+
+# Including allocator invariant tests (double-free, overlap, merge)
+pytest test_mooncake_transfer_engine_buffer.py -v -s
+```
+
+### Run Specific Test Classes
+
+```bash
+# Basic connector tests
+pytest test_mooncake_transfer_engine_rdma.py::TestBasicConnector -v -s
+
+# End-to-end RDMA transfer tests
+pytest test_mooncake_transfer_engine_rdma.py::TestEndToEnd -v -s
+
+# Lifecycle & resource management tests
+pytest test_mooncake_transfer_engine_rdma.py::TestLifecycle -v -s
+
+# GPU memory pool tests (requires CUDA)
+pytest test_mooncake_transfer_engine_rdma.py::TestGPUPool -v -s
+
+# Stress / correctness tests (slow)
+pytest test_mooncake_transfer_engine_rdma.py::TestStressCorrectness -v -s
+```
+
+### RDMA Environment Diagnostics
+
+For quick diagnostics (device status, Mooncake availability, env vars, etc.),
+see the [Troubleshooting section](../../../docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md#troubleshooting)
+in the connector documentation.
+
+---
+
+## Cross-Node Testing
+
+The `cross_node_mooncake_transfer_engine.py` script enables testing RDMA transfers between two separate physical machines. This script is **not** auto-discovered by `pytest` (it does not start with `test_`) — it must be run manually on each node.
+
+### Prerequisites
+
+1. Both machines have Mooncake installed
+2. Both machines are connected via InfiniBand/RoCE switch
+3. Firewall allows ZMQ ports (default: 15500, 15501)
+4. Same RDMA device name on both nodes (if multiple devices exist)
+
+### Running Cross-Node Tests
+
+**On Machine A (Producer) — start first:**
+
+```bash
+cd benchmarks/distributed/omni_connectors/
+
+# Optional: specify device if multiple exist
+export RDMA_DEVICE_NAME='mlx5_0'
+
+python cross_node_mooncake_transfer_engine.py \
+    --role producer \
+    --local-host <PRODUCER_IP> \
+    --remote-host <CONSUMER_IP> \
+    --tensor-size-mb 100 \
+    --num-transfers 3
+```
+
+**On Machine B (Consumer) — start after producer:**
+
+```bash
+cd benchmarks/distributed/omni_connectors/
+
+export RDMA_DEVICE_NAME='mlx5_0'
+
+python cross_node_mooncake_transfer_engine.py \
+    --role consumer \
+    --local-host <CONSUMER_IP> \
+    --remote-host <PRODUCER_IP> \
+    --tensor-size-mb 100 \
+    --num-transfers 3
+```
+
+### Transfer Modes
+
+| Mode | Description | Example |
+|------|-------------|---------|
+| `copy` | Normal path — tensor copied to RDMA pool (default) | `--mode copy` |
+| `zerocopy` | Zero-copy path — data created directly in RDMA pool | `--mode zerocopy` |
+| `gpu` | GPU transfer — RDMA pool on GPU, uses GPUDirect | `--mode gpu --gpu-id 0` |
+
+### Benchmark Mode
+
+Skip MD5 verification and measure pure RDMA throughput:
+
+```bash
+# Producer
+python cross_node_mooncake_transfer_engine.py \
+    --role producer \
+    --local-host <PRODUCER_IP> \
+    --remote-host <CONSUMER_IP> \
+    --tensor-size-mb 1024 \
+    --num-transfers 20 \
+    --benchmark
+
+# Consumer
+python cross_node_mooncake_transfer_engine.py \
+    --role consumer \
+    --local-host <CONSUMER_IP> \
+    --remote-host <PRODUCER_IP> \
+    --tensor-size-mb 1024 \
+    --num-transfers 20 \
+    --benchmark
+```
+
+### Cross-Node Test Options
+
+| Option | Description | Default |
+|--------|-------------|---------|
+| `--role` | `producer` or `consumer` | Required |
+| `--local-host` | Local RDMA IP address | Required |
+| `--remote-host` | Remote RDMA IP address | Required |
+| `--local-port` | Local ZMQ port for RDMA data | 15500 |
+| `--remote-port` | Remote ZMQ port for RDMA data | 15500 |
+| `--ctrl-port` | Control channel port | 15501 |
+| `--tensor-size-mb` | Tensor size in MB | 100 |
+| `--num-transfers` | Number of transfers | 3 |
+| `--mode` | `copy`, `zerocopy`, or `gpu` | `copy` |
+| `--gpu-id` | GPU ID for GPU mode | 0 |
+| `--benchmark` | Skip MD5, pure performance test | off |
+
+---
+
+## Troubleshooting
+
+### 1. "Failed to modify QP to RTR" Error
+
+**Cause**: QP handshake failed, usually due to device configuration mismatch.
+
+**Solution**:
+```bash
+# Force using the same device
+export RDMA_DEVICE_NAME='mlx5_0'
+```
+
+### 2. "Mooncake TransferEngine is not available"
+
+**Cause**: Mooncake not installed or import failed.
+
+**Solution**:
+```bash
+# Check Mooncake installation
+python -c "from mooncake.engine import TransferEngine; print('OK')"
+
+# Reinstall if needed
+pip install mooncake-transfer-engine
+# Or using uv
+uv pip install mooncake-transfer-engine
+
+```
+
+### 3. "Permission denied" accessing /dev/infiniband
+
+**Cause**: Container lacks IB device access permissions.
+
+**Solution**:
+```bash
+docker run --device=/dev/infiniband --cap-add=IPC_LOCK ...
+```
+
+### 4. Test Timeout
+
+**Cause**: RDMA connection establishment failed or network latency.
+
+**Solution**:
+```bash
+# Check network status
+ibstat
+ibstatus
+```
+
+### 5. GPU Test Failed "CUDA is not available"
+
+**Cause**: CUDA environment not configured or GPU unavailable.
+
+**Solution**:
+```bash
+# Check CUDA
+python -c "import torch; print(torch.cuda.is_available())"
+
+# Docker needs NVIDIA runtime
+docker run --gpus all ...
+```
+
+---
+
+## Environment Variables Reference
+
+| Variable | Description | Example |
+|----------|-------------|---------|
+| `RDMA_DEVICE_NAME` | Specify RDMA device name | `mlx5_0` |
+| `RDMA_TEST_HOST` | Specify test host IP | `10.0.0.1` |
+| `MC_TE_METRIC` | Enable Mooncake metrics | `1` |
+| `MC_IB_PCI_RELAXED_ORDERING` | Enable PCIe relaxed ordering | `1` |
+
+---
+
+## Test Files Overview
+
+| File | Description | Auto-discovered by pytest |
+|------|-------------|--------------------------|
+| `test_mooncake_transfer_engine_rdma.py` | Integration tests for MooncakeTransferEngineConnector (basic, E2E, lifecycle, GPU) | Yes |
+| `test_mooncake_transfer_engine_buffer.py` | Memory pool and buffer management unit tests | Yes |
+| `cross_node_mooncake_transfer_engine.py` | Cross-node (multi-machine) testing script — run manually | No (filename does not start with `test_`) |
+
+### test_mooncake_transfer_engine_rdma.py — Test Classes
+
+| Test Class | Memory Pool | Marker | Description |
+|------------|-------------|--------|-------------|
+| `TestBasicConnector` | CPU | — | Initialization, put tensor/bytes/object, cleanup, pool exhaustion |
+| `TestEndToEnd` | CPU | — | E2E RDMA transfer: tensor, bytes, object, zero-copy, large payload (100MB), mixed types, concurrency |
+| `TestLifecycle` | CPU | — | Close, context manager, double-close safety |
+| `TestGPUPool` | GPU | — | GPU pool init, put CPU/GPU tensor, GPU E2E transfer |
+| `TestStressCorrectness` | CPU | `slow` | Concurrent put+get with MD5 integrity, bidirectional concurrency, edge cases (1-element tensor, empty bytes), 500MB payload, rapid alloc/free cycles |
+
+### test_mooncake_transfer_engine_buffer.py — Test Classes
+
+| Test Class | Marker | Description |
+|------------|--------|-------------|
+| `TestBufferAllocator` | — | Basic alloc/free, alignment, exhaustion/recovery, thread safety |
+| `TestAllocatorInvariants` | `slow` | Double-free safety, overlap corruption detection, adjacent-block merging, fragmentation/defrag |
+| `TestManagedBuffer` | — | Tensor views, context manager |

benchmarks/distributed/omni_connectors/cross_node_mooncake_transfer_engine.py

@@ -0,0 +1,644 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+"""
+Cross-Node RDMA Test Script (Automated Version)
+
+This script enables testing RDMA transfers between two separate machines.
+Supports three transfer modes:
+  - copy:      Normal path - tensor copied to RDMA pool (default)
+  - zerocopy:  Zero-copy path - data created directly in RDMA pool
+  - gpu:       GPU transfer - RDMA pool on GPU, uses GPUDirect
+
+Usage:
+    # On Machine A (Producer) - start first:
+    python cross_node_mooncake_transfer_engine.py --role producer --local-host hostname_A --remote-host hostname_B
+
+    # On Machine B (Consumer) - start after producer:
+    python cross_node_mooncake_transfer_engine.py --role consumer --local-host hostname_B --remote-host hostname_A
+
+    # Zero-copy mode:
+    python cross_node_mooncake_transfer_engine.py --role producer ... --mode zerocopy
+
+    # GPU mode (requires GPUDirect RDMA support):
+    python cross_node_mooncake_transfer_engine.py --role producer ... --mode gpu --gpu-id 0
+
+    # Benchmark mode (skip random data generation and MD5 verification,
+    # measures pure RDMA throughput):
+    python cross_node_mooncake_transfer_engine.py --role producer ... --benchmark
+
+Environment Variables:
+    RDMA_DEVICE_NAME:              Specify RDMA device (e.g., mlx5_0)
+    MC_IB_PCI_RELAXED_ORDERING:    Set to 1 to enable PCIe relaxed ordering
+                                   for higher RDMA throughput
+"""
+
+import argparse
+import hashlib
+import os
+import sys
+import time
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import Any
+
+import msgspec
+import torch
+import zmq
+
+# Add parent path for imports
+sys.path.insert(
+    0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))))
+)
+
+from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import (
+    ManagedBuffer,
+    MooncakeTransferEngineConnector,
+    TransferEngine,
+)
+
+
+def compute_md5(tensor: torch.Tensor) -> str:
+    """Compute MD5 checksum of a tensor."""
+    if tensor.is_cuda:
+        tensor = tensor.cpu()
+    data = tensor.contiguous().view(torch.uint8).numpy().tobytes()
+    return hashlib.md5(data).hexdigest()
+
+
+# Control channel message types
+class CtrlMsg(msgspec.Struct):
+    """Control channel message."""
+
+    msg_type: str  # "READY", "TRANSFER", "ACK", "DONE", "ERROR"
+    request_id: str = ""
+    md5: str = ""
+    data_size: int = 0
+    error: str = ""
+
+
+@dataclass
+class TransferConfig:
+    """Configuration for cross-node transfer test."""
+
+    local_host: str
+    remote_host: str
+    local_port: int
+    remote_port: int
+    ctrl_port: int
+    num_transfers: int
+    tensor_size_mb: int
+    mode: str  # "copy", "zerocopy", "gpu"
+    gpu_id: int = 0
+    benchmark: bool = False  # Skip MD5 verification for pure performance test
+
+
+@dataclass
+class TransferStats:
+    """Statistics for transfer operations."""
+
+    success_count: int = 0
+    fail_count: int = 0
+    total_bytes: int = 0
+    elapsed_time: float = 0.0
+
+    @property
+    def throughput_mbps(self) -> float:
+        if self.elapsed_time > 0:
+            return (self.total_bytes / (1024 * 1024)) / self.elapsed_time
+        return 0.0
+
+    def print_summary(self, role: str):
+        print(f"\n{'=' * 60}")
+        print(f" {role.upper()} SUMMARY")
+        print(f"  Successful: {self.success_count}/{self.success_count + self.fail_count}")
+        print(f"  Failed:     {self.fail_count}/{self.success_count + self.fail_count}")
+        print(f"  Total:      {self.total_bytes / (1024 * 1024):.2f} MB")
+        print(f"  Time:       {self.elapsed_time:.2f} s")
+        print(f"  Throughput: {self.throughput_mbps:.2f} MB/s")
+        print(f"{'=' * 60}")
+
+
+class CrossNodeTester(ABC):
+    """Abstract base class for cross-node RDMA testing."""
+
+    def __init__(self, config: TransferConfig):
+        self.config = config
+        self.connector: MooncakeTransferEngineConnector | None = None
+        self.zmq_ctx: zmq.Context | None = None
+        self.ctrl_socket: zmq.Socket | None = None
+        self.stats = TransferStats()
+
+    def get_connector_config(self) -> dict:
+        """Get connector configuration based on mode."""
+        pool_size = int(self.config.tensor_size_mb * 1.5) * 1024 * 1024
+        pool_size = max(pool_size, 128 * 1024 * 1024)
+
+        conn_config = {
+            "host": self.config.local_host,
+            "zmq_port": self.config.local_port,
+            "protocol": "rdma",
+            "memory_pool_size": pool_size,
+        }
+
+        # Set device based on mode
+        if self.config.mode == "gpu":
+            conn_config["memory_pool_device"] = f"cuda:{self.config.gpu_id}"
+        else:
+            conn_config["memory_pool_device"] = "cpu"
+
+        # RDMA device name from environment
+        device_name = os.environ.get("RDMA_DEVICE_NAME")
+        if device_name:
+            conn_config["device_name"] = device_name
+            print(f"[CONFIG] Using RDMA device: {device_name}")
+
+        return conn_config
+
+    def initialize(self):
+        """Initialize connector and ZMQ context."""
+        print(f"[{self.role}] Initializing connector...")
+        conn_config = self.get_connector_config()
+        self.connector = MooncakeTransferEngineConnector(conn_config)
+        self.zmq_ctx = zmq.Context()
+        print(f"[{self.role}] Ready at {self.config.local_host}:{self.config.local_port}")
+
+    def cleanup(self):
+        """Cleanup resources."""
+        if self.ctrl_socket:
+            self.ctrl_socket.close()
+        if self.zmq_ctx:
+            self.zmq_ctx.term()
+        if self.connector:
+            self.connector.close()
+        print(f"[{self.role}] Closed.")
+
+    @property
+    @abstractmethod
+    def role(self) -> str:
+        pass
+
+    @abstractmethod
+    def run(self):
+        pass
+
+
+class Producer(CrossNodeTester):
+    """Producer node - sends data to consumer."""
+
+    @property
+    def role(self) -> str:
+        return "PRODUCER"
+
+    def print_header(self):
+        print(f"\n{'=' * 60}")
+        print(f" PRODUCER MODE ({self.config.mode.upper()})")
+        print(f" Local:  {self.config.local_host}:{self.config.local_port}")
+        print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
+        print(f" Control Port: {self.config.ctrl_port}")
+        print(f" Transfer Mode: {self.config.mode}")
+        if self.config.mode == "gpu":
+            print(f" GPU ID: {self.config.gpu_id}")
+        print(f"{'=' * 60}\n")
+
+    def setup_control_channel(self):
+        """Setup ZMQ control channel as server (REP socket)."""
+        self.ctrl_socket = self.zmq_ctx.socket(zmq.REP)
+        self.ctrl_socket.bind(f"tcp://*:{self.config.ctrl_port}")
+        print(f"[PRODUCER] Control channel listening on port {self.config.ctrl_port}")
+        print("[PRODUCER] Waiting for consumer to connect...")
+
+    def wait_for_consumer(self) -> bool:
+        """Wait for consumer READY signal."""
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+        if msg.msg_type != "READY":
+            print(f"[PRODUCER] Unexpected message: {msg.msg_type}")
+            return False
+        print("[PRODUCER] Consumer connected!")
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
+        return True
+
+    def create_test_data(self, transfer_idx: int) -> tuple[Any, str, int]:
+        """
+        Create test data based on transfer mode.
+
+        Returns:
+            (data, md5, size) tuple
+        """
+        num_elements = (self.config.tensor_size_mb * 1024 * 1024) // 4
+        data_size = num_elements * 4
+
+        if self.config.mode == "zerocopy":
+            # Zero-Copy Path: Allocate directly from connector's pool
+            offset = self.connector.allocator.alloc(data_size)
+            managed_buf = ManagedBuffer(self.connector.allocator, offset, data_size, self.connector.pool)
+
+            if self.config.benchmark:
+                # In benchmark mode, skip random data generation (use uninitialized memory)
+                return managed_buf, "", data_size
+            else:
+                # Fill buffer with random data using tensor view
+                tensor_view = managed_buf.as_tensor(dtype=torch.float32, shape=(num_elements,))
+                random_data = torch.randn(num_elements, dtype=torch.float32)
+                if tensor_view.is_cuda:
+                    tensor_view.copy_(random_data.to(tensor_view.device))
+                else:
+                    tensor_view.copy_(random_data)
+                md5 = compute_md5(tensor_view)
+                return managed_buf, md5, data_size
+
+        elif self.config.mode == "gpu":
+            # GPU Path: Create tensor on GPU
+            device = f"cuda:{self.config.gpu_id}"
+            if self.config.benchmark:
+                # In benchmark mode, use empty tensor (no random generation)
+                gpu_tensor = torch.empty(num_elements, dtype=torch.float32, device=device)
+                return gpu_tensor, "", data_size
+            else:
+                cpu_tensor = torch.randn(num_elements, dtype=torch.float32)
+                md5 = compute_md5(cpu_tensor)
+                gpu_tensor = cpu_tensor.to(device)
+                return gpu_tensor, md5, data_size
+
+        else:
+            # Copy Path (default): Create regular CPU tensor
+            if self.config.benchmark:
+                # In benchmark mode, use empty tensor (no random generation)
+                tensor = torch.empty(num_elements, dtype=torch.float32)
+                return tensor, "", data_size
+            else:
+                tensor = torch.randn(num_elements, dtype=torch.float32)
+                md5 = compute_md5(tensor)
+                return tensor, md5, data_size
+
+    def do_transfer(self, transfer_idx: int) -> bool:
+        """Perform a single transfer."""
+        req_id = f"cross_node_transfer_{transfer_idx}"
+
+        if not self.config.benchmark:
+            print(f"\n[PRODUCER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
+
+        # Create test data
+        t0 = time.time()
+        data, md5, data_size = self.create_test_data(transfer_idx)
+        t_create = time.time() - t0
+
+        if not self.config.benchmark:
+            print(f"  Mode: {self.config.mode}")
+            print(f"  Size: {self.config.tensor_size_mb} MB")
+            if md5:
+                print(f"  MD5:  {md5[:16]}...")
+            print(f"  Create time: {t_create * 1000:.1f} ms")
+
+        # Put data
+        t1 = time.time()
+        success, size, metadata = self.connector.put("producer", "consumer", req_id, data)
+        t_put = time.time() - t1
+
+        if not success:
+            print("  [FAIL] Put failed")
+            return False
+
+        if not self.config.benchmark:
+            print(f"  [OK] Put successful, {size} bytes ({t_put * 1000:.1f} ms)")
+
+        # Wait for consumer to request transfer info
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+
+        if msg.msg_type != "READY":
+            print(f"  [ERROR] Unexpected message: {msg.msg_type}")
+            return False
+
+        # Send transfer metadata to consumer
+        transfer_msg = CtrlMsg(
+            msg_type="TRANSFER",
+            request_id=req_id,
+            md5=md5,
+            data_size=data_size,
+        )
+        self.ctrl_socket.send(msgspec.msgpack.encode(transfer_msg))
+
+        # Wait for consumer ACK (this includes RDMA transfer time)
+        t2 = time.time()
+        msg_data = self.ctrl_socket.recv()
+        t_rdma = time.time() - t2
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+
+        success = msg.msg_type == "ACK"
+        if success:
+            if not self.config.benchmark:
+                print(f"  [OK] RDMA transfer complete ({t_rdma * 1000:.1f} ms)")
+            self.stats.success_count += 1
+            self.stats.total_bytes += size
+        else:
+            print(f"  [WARN] Consumer reported error: {msg.error}")
+            self.stats.fail_count += 1
+
+        # Send ACK to allow consumer to continue
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="ACK")))
+
+        # Cleanup buffer
+        self.connector.cleanup(req_id)
+
+        return success
+
+    def run(self):
+        """Run the producer."""
+        self.print_header()
+        self.initialize()
+        self.setup_control_channel()
+
+        try:
+            if not self.wait_for_consumer():
+                return
+
+            if self.config.benchmark:
+                print(
+                    f"[BENCHMARK] Running {self.config.num_transfers} "
+                    f"transfers of {self.config.tensor_size_mb} MB each..."
+                )
+
+            start_time = time.time()
+
+            for i in range(self.config.num_transfers):
+                self.do_transfer(i)
+                if self.config.benchmark and (i + 1) % 10 == 0:
+                    elapsed = time.time() - start_time
+                    current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
+                    print(f"  Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
+
+            self.stats.elapsed_time = time.time() - start_time
+            self.stats.print_summary("PRODUCER")
+
+            # Wait for final consumer message and send DONE
+            self.ctrl_socket.recv()
+            self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="DONE")))
+
+        finally:
+            self.cleanup()
+
+
+class Consumer(CrossNodeTester):
+    """Consumer node - receives data from producer."""
+
+    @property
+    def role(self) -> str:
+        return "CONSUMER"
+
+    def print_header(self):
+        print(f"\n{'=' * 60}")
+        print(f" CONSUMER MODE ({self.config.mode.upper()})")
+        print(f" Local:  {self.config.local_host}:{self.config.local_port}")
+        print(f" Remote: {self.config.remote_host}:{self.config.remote_port}")
+        print(f" Control Port: {self.config.ctrl_port}")
+        print(f" Transfer Mode: {self.config.mode}")
+        if self.config.mode == "gpu":
+            print(f" GPU ID: {self.config.gpu_id}")
+        print(f"{'=' * 60}\n")
+
+    def setup_control_channel(self):
+        """Setup ZMQ control channel as client (REQ socket)."""
+        self.ctrl_socket = self.zmq_ctx.socket(zmq.REQ)
+        ctrl_addr = f"tcp://{self.config.remote_host}:{self.config.ctrl_port}"
+        print(f"[CONSUMER] Connecting to producer control channel at {ctrl_addr}...")
+        self.ctrl_socket.connect(ctrl_addr)
+
+    def connect_to_producer(self) -> bool:
+        """Connect to producer and send READY signal."""
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+        if msg.msg_type != "ACK":
+            print(f"[CONSUMER] Unexpected response: {msg.msg_type}")
+            return False
+        print("[CONSUMER] Connected to producer! Starting transfers...")
+        return True
+
+    def do_transfer(self, transfer_idx: int) -> bool:
+        """Perform a single transfer."""
+        if not self.config.benchmark:
+            print(f"\n[CONSUMER] Transfer {transfer_idx + 1}/{self.config.num_transfers}")
+
+        # Request next transfer info
+        self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
+        msg_data = self.ctrl_socket.recv()
+        msg = msgspec.msgpack.decode(msg_data, type=CtrlMsg)
+
+        if msg.msg_type == "DONE":
+            print("[CONSUMER] Producer signaled completion")
+            return False
+
+        if msg.msg_type != "TRANSFER":
+            print(f"[CONSUMER] Unexpected message: {msg.msg_type}")
+            return False
+
+        req_id = msg.request_id
+        expected_md5 = msg.md5
+        data_size = msg.data_size
+        num_elements = data_size // 4
+
+        if not self.config.benchmark:
+            print(f"  Request ID: {req_id}")
+            if expected_md5:
+                print(f"  Expected MD5: {expected_md5[:16]}...")
+            print(f"  Data Size: {data_size / (1024 * 1024):.2f} MB")
+
+        # Build metadata for get
+        metadata = {
+            "request_id": req_id,
+            "source_host": self.config.remote_host,
+            "source_port": self.config.remote_port,
+            "data_size": data_size,
+            "dtype": "float32",
+            "shape": [num_elements],
+            "is_fast_path": True,
+        }
+
+        if not self.config.benchmark:
+            print(f"  [INFO] Requesting from {self.config.remote_host}:{self.config.remote_port}")
+
+        # Get data with timing
+        t0 = time.time()
+        result = self.connector.get("producer", "consumer", req_id, metadata)
+        t_get = time.time() - t0
+
+        response_msg = CtrlMsg(msg_type="ERROR", error="Get failed")
+
+        if result is not None:
+            recv_buffer, recv_size = result
+            if not self.config.benchmark:
+                print(f"  [OK] Get successful, {recv_size} bytes ({t_get * 1000:.1f} ms)")
+
+            if isinstance(recv_buffer, ManagedBuffer):
+                # In benchmark mode, skip MD5 verification
+                if self.config.benchmark or not expected_md5:
+                    response_msg = CtrlMsg(msg_type="ACK")
+                    self.stats.success_count += 1
+                    self.stats.total_bytes += recv_size
+                else:
+                    # Verify data
+                    t1 = time.time()
+                    reconstructed = recv_buffer.as_tensor(dtype=torch.float32, shape=(num_elements,))
+                    recv_md5 = compute_md5(reconstructed)
+                    t_md5 = time.time() - t1
+                    print(f"  MD5: {recv_md5[:16]}... ({t_md5 * 1000:.1f} ms)")
+
+                    if recv_md5 == expected_md5:
+                        print("  [PASS] MD5 checksum verified!")
+                        response_msg = CtrlMsg(msg_type="ACK")
+                        self.stats.success_count += 1
+                        self.stats.total_bytes += recv_size
+                    else:
+                        print("  [FAIL] MD5 mismatch!")
+                        response_msg = CtrlMsg(msg_type="ERROR", error="MD5 mismatch")
+                        self.stats.fail_count += 1
+
+                recv_buffer.release()
+            else:
+                response_msg = CtrlMsg(msg_type="ACK")
+                self.stats.success_count += 1
+                self.stats.total_bytes += recv_size
+        else:
+            print("  [FAIL] Get failed")
+            self.stats.fail_count += 1
+
+        # Send response to producer
+        self.ctrl_socket.send(msgspec.msgpack.encode(response_msg))
+        # Wait for ACK
+        self.ctrl_socket.recv()
+
+        return response_msg.msg_type == "ACK"
+
+    def run(self):
+        """Run the consumer."""
+        self.print_header()
+        self.initialize()
+        self.setup_control_channel()
+
+        try:
+            if not self.connect_to_producer():
+                return
+
+            if self.config.benchmark:
+                print(
+                    f"[BENCHMARK] Running {self.config.num_transfers} "
+                    f"transfers of {self.config.tensor_size_mb} MB each..."
+                )
+
+            start_time = time.time()
+
+            for i in range(self.config.num_transfers):
+                if not self.do_transfer(i):
+                    break
+                if self.config.benchmark and (i + 1) % 10 == 0:
+                    elapsed = time.time() - start_time
+                    current_throughput = (self.stats.total_bytes / (1024 * 1024)) / elapsed
+                    print(f"  Progress: {i + 1}/{self.config.num_transfers}, Throughput: {current_throughput:.2f} MB/s")
+
+            self.stats.elapsed_time = time.time() - start_time
+            self.stats.print_summary("CONSUMER")
+
+            # Send final READY and wait for DONE
+            self.ctrl_socket.send(msgspec.msgpack.encode(CtrlMsg(msg_type="READY")))
+            self.ctrl_socket.recv()
+
+        finally:
+            self.cleanup()
+
+
+def main():
+    parser = argparse.ArgumentParser(
+        description="Cross-Node RDMA Test (Automated)",
+        formatter_class=argparse.RawDescriptionHelpFormatter,
+        epilog="""
+Transfer Modes:
+  copy      - Normal path: tensor copied to RDMA pool (default)
+  zerocopy  - Zero-copy path: data created directly in RDMA pool
+  gpu       - GPU transfer: RDMA pool on GPU, uses GPUDirect
+
+Examples:
+  # Copy mode (default):
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB
+
+  # Zero-copy mode:
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB --mode zerocopy
+
+  # GPU mode:
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB --mode gpu --gpu-id 0
+
+  # Benchmark mode (skip MD5, measure pure RDMA performance):
+  python cross_node_mooncake_transfer_engine.py --role producer \
+  --local-host hostA --remote-host hostB --benchmark
+
+  # With specific RDMA device:
+  RDMA_DEVICE_NAME=mlx5_0 python cross_node_mooncake_transfer_engine.py --role producer ...
+        """,
+    )
+
+    parser.add_argument(
+        "--role", required=True, choices=["producer", "consumer"], help="Role: producer (sends) or consumer (receives)"
+    )
+    parser.add_argument("--local-host", required=True, help="Local hostname or IP address")
+    parser.add_argument("--remote-host", required=True, help="Remote hostname or IP address")
+    parser.add_argument("--local-port", type=int, default=15500, help="Local ZMQ port for RDMA data (default: 15500)")
+    parser.add_argument("--remote-port", type=int, default=15500, help="Remote ZMQ port for RDMA data (default: 15500)")
+    parser.add_argument("--ctrl-port", type=int, default=15501, help="Control channel port (default: 15501)")
+    parser.add_argument("--num-transfers", type=int, default=20, help="Number of transfers to perform (default: 3)")
+    parser.add_argument("--tensor-size-mb", type=int, default=100, help="Tensor size in MB (default: 100)")
+    parser.add_argument(
+        "--mode",
+        choices=["copy", "zerocopy", "gpu"],
+        default="copy",
+        help="Transfer mode: copy, zerocopy, or gpu (default: copy)",
+    )
+    parser.add_argument("--gpu-id", type=int, default=0, help="GPU ID for GPU mode (default: 0)")
+    parser.add_argument(
+        "--benchmark", action="store_true", help="Benchmark mode: skip MD5 verification for pure performance test"
+    )
+
+    args = parser.parse_args()
+
+    # Check Mooncake
+    if TransferEngine is None:
+        print("[ERROR] Mooncake TransferEngine is not available.")
+        print("Install with: pip install mooncake")
+        sys.exit(1)
+
+    # Check CUDA for GPU mode
+    if args.mode == "gpu":
+        if not torch.cuda.is_available():
+            print("[ERROR] CUDA is not available but GPU mode was requested.")
+            sys.exit(1)
+        if args.gpu_id >= torch.cuda.device_count():
+            print(f"[ERROR] GPU {args.gpu_id} not available. Found {torch.cuda.device_count()} GPUs.")
+            sys.exit(1)
+        print(f"[INFO] Using GPU {args.gpu_id}: {torch.cuda.get_device_name(args.gpu_id)}")
+
+    config = TransferConfig(
+        local_host=args.local_host,
+        remote_host=args.remote_host,
+        local_port=args.local_port,
+        remote_port=args.remote_port,
+        ctrl_port=args.ctrl_port,
+        num_transfers=args.num_transfers,
+        tensor_size_mb=args.tensor_size_mb,
+        mode=args.mode,
+        gpu_id=args.gpu_id,
+        benchmark=args.benchmark,
+    )
+
+    if args.role == "producer":
+        tester = Producer(config)
+    else:
+        tester = Consumer(config)
+
+    tester.run()
+
+
+if __name__ == "__main__":
+    main()

benchmarks/qwen3-omni/README.md

@@ -0,0 +1,86 @@
+# Benchmarks Guide
+
+This README explains how to (1) prepare benchmark datasets and (2) run the provided Qwen3-Omni benchmarks.
+
+## 1) Prepare the dataset (SeedTTS top100)
+
+```bash
+cd benchmarks/build_dataset
+pip install gdown
+
+# Download SeedTTS test set from Google Drive
+gdown --id 1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP
+
+# Extract
+tar -xf seedtts_testset.tar
+
+# Copy metadata and extract top-100 prompts
+cp seedtts_testset/en/meta.lst meta.lst
+python extract_prompts.py -i meta.lst -o top100.txt -n 100
+
+# (Optional) clean up to save space
+rm -rf seedtts_testset seedtts_testset.tar meta.lst
+```
+
+Artifacts:
+- `benchmarks/build_dataset/top100.txt` — 100 text prompts (one per line).
+
+## 2) Run benchmarks
+
+All commands assume repo root (`vllm-omni`).
+
+### A. Transformers benchmark (offline, HF Transformers)
+
+```
+bash benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh
+```
+
+What it does:
+- Runs `qwen3_omni_moe_transformers.py` over `top100.txt` with `--num-prompts 100`.
+- Outputs to `benchmarks/qwen3-omni/transformers/benchmark_results/`:
+  - `perf_stats.json` — aggregated & per-prompt TPS/latency (thinker/talker/code2wav/overall).
+  - `results.json` — per-prompt outputs and audio paths.
+  - `audio/` — ~100 generated `.wav` files.
+
+Key checks:
+- `overall_tps` and `*_tps_avg` should be non-zero and reasonably stable.
+- Investigate any 0/NaN or unusually low TPS / long-tail latency.
+
+### B. vLLM Omni end-to-end benchmark (pipeline)
+
+```
+bash benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh
+```
+
+What it does:
+- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
+- Uses `benchmarks/build_dataset/top100.txt` and writes to:
+  - Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
+    - `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
+    - `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
+    - `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
+  - Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
+
+Key checks:
+- Overall stats: end-to-end latency/TPS should be reasonable.
+- Orchestrator stats: per-stage latency should be stable; investigate long tails.
+- Stage logs: ensure no errors and no unusually slow stages.
+
+
+## Performance snapshot
+
+The chart below summarizes our measured Qwen3-Omni MoE end-to-end benchmark, comparing vLLM-Omni against HF Transformers. It shows the overall throughput advantage for vLLM-Omni. These are actual experiment results—please refer to this performance when evaluating or reproducing the benchmark.
+
+![vLLM-Omni vs HF](./vllm-omni-vs-hf.png)
+
+## Directory layout
+- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
+- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
+- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
+- `benchmarks/<model>/vllm-omni-vs-hf.png` — current performance snapshot (overall throughput comparison).
+- `benchmarks/<model>/transformers/` — HF Transformers benchmarks (offline reference).
+
+## Troubleshooting
+- Make sure GPU/driver/FlashAttention2 requirements are met for the chosen model.
+- If downloads fail, confirm network access to Google Drive (`gdown`) and Hugging Face.
+- If audio files are missing, check for errors in stage logs or model generation.***

benchmarks/qwen3-omni/transformers/eval_qwen3_moe_omni_transformers.sh

@@ -0,0 +1,26 @@
+#!/bin/bash
+# Qwen3-Omni Transformers Benchmark Evaluation Script
+# This script must be run from the vllm-omni root directory
+
+# Get the directory where this script is located
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Navigate to vllm-omni root directory (4 levels up from script location)
+VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
+
+echo "Working directory: $(pwd)"
+# Verify we're in the correct directory and run benchmark
+if [[ ! -f "benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py" ]]; then
+    echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
+else
+    cd benchmarks/qwen3-omni/transformers
+
+    python qwen3_omni_moe_transformers.py --prompts-file ../../build_dataset/top100.txt --num-prompts 100
+
+    echo "Logs and outputs are saved to $(pwd)/benchmark_results:"
+    echo "  - perf_stats.json    Aggregated/per-prompt TPS and latency (thinker/talker/code2wav/overall)"
+    echo "  - results.json       Per-prompt outputs and audio paths"
+    echo "  - audio/             Generated wav files, there should be 100 wav file generated"
+    echo "Key checks: overall_tps and *_tps_avg should be non-zero and stable; investigate 0/NaN or unusually low TPS/long-tail latency."
+fi

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_model.py

@@ -0,0 +1,265 @@
+import time
+
+import torch
+from transformers import Qwen3OmniMoeForConditionalGeneration
+
+
+class Qwen3OmniMoeForConditionalGenerationWithLogging(Qwen3OmniMoeForConditionalGeneration):
+    @torch.no_grad()
+    def generate(
+        self,
+        input_ids: torch.Tensor | None = None,
+        speaker: str = "Ethan",
+        use_audio_in_video: bool = False,
+        return_audio: bool | None = None,
+        thinker_max_new_tokens: int = 1024,
+        thinker_eos_token_id: int = 151645,
+        talker_max_new_tokens: int = 4096,
+        talker_do_sample: bool = True,
+        talker_top_k: int = 50,
+        talker_top_p: float = 1.0,
+        talker_temperature: float = 0.9,
+        talker_repetition_penalty: float = 1.05,
+        **kwargs,
+    ):
+        total_t0 = time.time()
+        perf_stats = {
+            "thinker_tokens": 0,
+            "thinker_time_s": 0.0,
+            "thinker_tps": 0.0,
+            "talker_tokens": 0,
+            "talker_time_s": 0.0,
+            "talker_tps": 0.0,
+            "code2wav_tokens": 0,
+            "code2wav_time_s": 0.0,
+            "code2wav_tps": 0.0,
+            "total_tokens": 0,
+            "total_time_s": 0.0,
+            "total_tps": 0.0,
+        }
+        if return_audio and not self.has_talker:
+            raise ValueError(
+                "Cannot use talker when talker module not initialized. "
+                "Use `enable_talker` method or set enable_talker in config "
+                "to enable talker."
+            )
+        if return_audio is None:
+            return_audio = self.has_talker
+
+        shared_kwargs = {"use_audio_in_video": use_audio_in_video}
+        thinker_kwargs = {
+            "max_new_tokens": thinker_max_new_tokens,
+            "eos_token_id": thinker_eos_token_id,
+        }
+
+        talker_kwargs = {}
+        token2wav_kwargs = {}
+        if return_audio:
+            speaker_id = self.config.talker_config.speaker_id.get(speaker.lower())
+            if speaker_id is None:
+                raise NotImplementedError(f"Speaker {speaker} not implemented")
+            if input_ids.shape[0] != 1:
+                raise NotImplementedError("Qwen3-Omni currently does not support batched inference with audio output")
+            talker_suppressed_tokens = [
+                i
+                for i in range(
+                    self.config.talker_config.text_config.vocab_size - 1024,
+                    self.config.talker_config.text_config.vocab_size,
+                )
+                if i != self.config.talker_config.codec_eos_token_id
+            ]  # Suppress additional special tokens, should not be predicted
+            talker_kwargs = {
+                "max_new_tokens": talker_max_new_tokens,
+                "do_sample": talker_do_sample,
+                "top_k": talker_top_k,
+                "top_p": talker_top_p,
+                "temperature": talker_temperature,
+                "eos_token_id": self.config.talker_config.codec_eos_token_id,
+                "repetition_penalty": talker_repetition_penalty,
+                "suppress_tokens": talker_suppressed_tokens,
+                "output_hidden_states": True,
+                "return_dict_in_generate": True,
+            }
+            token2wav_kwargs = {}
+
+        for key, value in kwargs.items():
+            if key.startswith("thinker_"):
+                thinker_kwargs[key[len("thinker_") :]] = value
+            elif key.startswith("talker_"):
+                talker_kwargs[key[len("talker_") :]] = value
+            elif key.startswith("token2wav_"):
+                token2wav_kwargs[key[len("token2wav_") :]] = value
+            # Process special input values
+            elif key == "feature_attention_mask":
+                thinker_kwargs[key] = value
+                talker_kwargs["audio_feature_lengths"] = torch.sum(value, dim=1)
+            elif key in ("input_features", "attention_mask"):
+                thinker_kwargs[key] = value
+            # Put other key to shared kwargs
+            else:
+                shared_kwargs[key] = value
+
+        # Merge kwargs
+        for key, value in shared_kwargs.items():
+            if key not in thinker_kwargs:
+                thinker_kwargs[key] = value
+            if key not in talker_kwargs and key in ["image_grid_thw", "video_grid_thw", "video_second_per_grid"]:
+                talker_kwargs[key] = value
+            if key not in token2wav_kwargs:
+                token2wav_kwargs[key] = value
+
+        # 1. Generate from thinker module
+        generate_audio = return_audio and self.has_talker
+        if generate_audio:
+            thinker_kwargs["output_hidden_states"] = True
+            thinker_kwargs["return_dict_in_generate"] = True
+
+        t0 = time.time()
+        thinker_result = self.thinker.generate(input_ids=input_ids, **thinker_kwargs)
+        t1 = time.time()
+        perf_stats["thinker_time_s"] = max(0.0, t1 - t0)
+        try:
+            prompt_len = int(input_ids.shape[1]) if input_ids is not None else 0
+            total_len = int(thinker_result.sequences.shape[-1])
+            thinker_out_len = max(0, total_len - prompt_len)
+        except Exception:
+            thinker_out_len = 0
+        perf_stats["thinker_tokens"] = thinker_out_len
+        perf_stats["thinker_tps"] = (
+            (thinker_out_len / perf_stats["thinker_time_s"]) if perf_stats["thinker_time_s"] > 0 else 0.0
+        )
+
+        if not generate_audio:
+            perf_stats["total_tokens"] = perf_stats["thinker_tokens"]
+            perf_stats["total_time_s"] = time.time() - total_t0
+            perf_stats["total_tps"] = (
+                (perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
+            )
+            # attach stats to self
+            setattr(self, "_perf_stats_last", perf_stats)
+            if not hasattr(self, "_perf_stats_history"):
+                setattr(self, "_perf_stats_history", [])
+            self._perf_stats_history.append(perf_stats)
+            return thinker_result, None
+
+        # 2. Prepare talker input
+        thinker_embed = torch.cat([hidden_states[0] for hidden_states in thinker_result.hidden_states], dim=1).to(
+            self.talker.device
+        )  # [1 t d]
+        thinker_hidden = torch.cat(
+            [
+                hidden_states[self.config.talker_config.accept_hidden_layer]
+                for hidden_states in thinker_result.hidden_states
+            ],
+            dim=1,
+        ).to(self.talker.device)  # [1 t d]
+
+        im_start_indexes = torch.cat(
+            (
+                torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(),
+                torch.tensor([thinker_result.sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype),
+            ),
+            dim=-1,
+        ).to(self.talker.device)  # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here.
+        multimodal_mask = (
+            (thinker_result.sequences == self.config.thinker_config.audio_token_id) |
+            (thinker_result.sequences == self.config.thinker_config.image_token_id) |
+            (thinker_result.sequences == self.config.thinker_config.video_token_id)
+        ).to(self.talker.device)  # [1 t] # fmt: skip
+
+        talker_special_tokens = torch.tensor(
+            [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]],
+            device=self.thinker.device,
+            dtype=input_ids.dtype,
+        )
+        tts_bos_embed, tts_eos_embed, tts_pad_embed = (
+            self.talker.text_projection(self.thinker.get_input_embeddings()(talker_special_tokens))
+            .to(self.talker.device)
+            .chunk(3, dim=1)
+        )  # 3 * [1 1 d]
+
+        talker_input_embeds = []  # [1 t d]
+        talker_input_ids = []
+        # For every chatml parts
+        for i in range(len(im_start_indexes) - 1):
+            im_start_index = im_start_indexes[i]
+            segment_end_index = im_start_indexes[i + 1]
+            role_token = input_ids[0][im_start_index + 1]
+            # Talker should ignore thinker system prompt
+            if role_token == self.config.system_token_id:
+                continue
+            # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs
+            elif role_token == self.config.user_token_id:
+                talker_user_part = self._get_talker_user_parts(
+                    im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed
+                )
+                talker_input_embeds.append(talker_user_part)
+                talker_input_ids.append(thinker_result.sequences[:, im_start_index:segment_end_index])
+            # Take assistant output (for now)
+            elif role_token == self.config.assistant_token_id and i == len(im_start_indexes) - 2:
+                talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts(
+                    im_start_index,
+                    segment_end_index,
+                    speaker_id,
+                    thinker_embed,
+                    tts_pad_embed,
+                    tts_bos_embed,
+                    tts_eos_embed,
+                )
+                talker_input_embeds.append(talker_assistant_embeds)
+                talker_input_ids.append(talker_assistant_ids)
+            # History assistant output (ignore for now)
+            elif role_token == self.config.assistant_token_id and i != len(im_start_indexes) - 2:
+                continue
+            else:
+                raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)")
+        talker_input_embed = torch.cat([embed.to(self.talker.device) for embed in talker_input_embeds], dim=1)
+        talker_input_id = torch.cat([embed.to(self.talker.device) for embed in talker_input_ids], dim=1)
+        t2 = time.time()
+        talker_result = self.talker.generate(
+            inputs_embeds=talker_input_embed,
+            trailing_text_hidden=trailing_text_hidden,
+            tts_pad_embed=tts_pad_embed,
+            talker_input_ids=talker_input_id,  # Not use input_ids to prevent repetition penalty out of bound
+            **talker_kwargs,
+        )
+        t3 = time.time()
+        perf_stats["talker_time_s"] = max(0.0, t3 - t2)
+        talker_codes = (
+            torch.stack([hid[-1] for hid in talker_result.hidden_states if hid[-1] is not None], dim=1)
+            .transpose(1, 2)
+            .to(self.code2wav.device)
+        )
+        try:
+            # codes shape: (B, num_quantizers, T). We log T as token length.
+            perf_stats["talker_tokens"] = int(talker_codes.shape[-1])
+        except Exception:
+            perf_stats["talker_tokens"] = 0
+        perf_stats["talker_tps"] = (
+            (perf_stats["talker_tokens"] / perf_stats["talker_time_s"]) if perf_stats["talker_time_s"] > 0 else 0.0
+        )
+        t4 = time.time()
+        talker_wavs = self.code2wav.chunked_decode(talker_codes, chunk_size=300, left_context_size=25).float()
+        t5 = time.time()
+        perf_stats["code2wav_time_s"] = max(0.0, t5 - t4)
+        perf_stats["code2wav_tokens"] = perf_stats["talker_tokens"]  # same T, not times 16
+        perf_stats["code2wav_tps"] = (
+            (perf_stats["code2wav_tokens"] / perf_stats["code2wav_time_s"])
+            if perf_stats["code2wav_time_s"] > 0
+            else 0.0
+        )
+        perf_stats["total_tokens"] = perf_stats["thinker_tokens"] + perf_stats["talker_tokens"]
+        perf_stats["total_time_s"] = time.time() - total_t0
+        perf_stats["total_tps"] = (
+            (perf_stats["total_tokens"] / perf_stats["total_time_s"]) if perf_stats["total_time_s"] > 0 else 0.0
+        )
+        setattr(self, "_perf_stats_last", perf_stats)
+        if not hasattr(self, "_perf_stats_history"):
+            setattr(self, "_perf_stats_history", [])
+        self._perf_stats_history.append(perf_stats)
+        return thinker_result, talker_wavs.float()
+
+
+__all__ = [
+    "Qwen3OmniMoeForConditionalGenerationWithLogging",
+]

benchmarks/qwen3-omni/transformers/qwen3_omni_moe_transformers.py

@@ -0,0 +1,275 @@
+import argparse
+import json
+import os
+
+import soundfile as sf
+from qwen3_omni_moe_model import Qwen3OmniMoeForConditionalGenerationWithLogging
+from qwen_omni_utils import process_mm_info
+from tqdm import tqdm
+from transformers import Qwen3OmniMoeProcessor
+
+MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
+# MODEL_PATH = "Qwen/Qwen3-Omni-30B-A3B-Thinking"
+
+
+def load_prompts(prompts_file: str) -> list[str]:
+    """Load prompts from a text file, one prompt per line."""
+    prompts = []
+    with open(prompts_file, encoding="utf-8") as f:
+        for line in f:
+            line = line.strip()
+            if line:
+                prompts.append(line)
+    return prompts
+
+
+def run_benchmark(
+    model,
+    processor,
+    prompts: list[str],
+    output_dir: str = "benchmark_results",
+    speaker: str = "Ethan",
+    use_audio_in_video: bool = True,
+):
+    """
+    Run benchmark on a list of prompts and collect performance stats.
+
+    Args:
+        model: The Qwen3OmniMoe model
+        processor: The Qwen3OmniMoe processor
+        prompts: List of text prompts to process
+        output_dir: Directory to save results
+        speaker: Speaker voice for audio output
+        use_audio_in_video: Whether to use audio in video
+
+    Returns:
+        tuple: (aggregated_stats, results, audio_outputs)
+            - aggregated_stats: dict with aggregated performance statistics
+            - results: list of dicts with per-prompt results
+            - audio_outputs: list of audio tensors/arrays (or None if no audio)
+    """
+    os.makedirs(output_dir, exist_ok=True)
+    audio_dir = os.path.join(output_dir, "audio")
+    os.makedirs(audio_dir, exist_ok=True)
+
+    all_stats = []
+    results = []
+    audio_outputs = []
+
+    for idx, prompt in enumerate(tqdm(prompts, desc="Processing prompts")):
+        conversation = [
+            {
+                "role": "user",
+                "content": [{"type": "text", "text": prompt}],
+            },
+        ]
+
+        # Preparation for inference
+        text = processor.apply_chat_template(conversation, add_generation_prompt=True, tokenize=False)
+        audios, images, videos = process_mm_info(conversation, use_audio_in_video=use_audio_in_video)
+        inputs = processor(
+            text=text,
+            audio=audios,
+            images=images,
+            videos=videos,
+            return_tensors="pt",
+            padding=True,
+            use_audio_in_video=use_audio_in_video,
+        )
+        inputs = inputs.to(model.device).to(model.dtype)
+
+        # Inference: Generation of the output text and audio
+        text_ids, audio = model.generate(
+            **inputs, speaker=speaker, thinker_return_dict_in_generate=True, use_audio_in_video=use_audio_in_video
+        )
+
+        # Decode output text
+        output_text = processor.batch_decode(
+            text_ids.sequences[:, inputs["input_ids"].shape[1] :],
+            skip_special_tokens=True,
+            clean_up_tokenization_spaces=False,
+        )[0]
+
+        # Collect performance stats
+        perf_stats = None
+        if hasattr(model, "_perf_stats_last"):
+            perf_stats = model._perf_stats_last.copy()
+            perf_stats["prompt_idx"] = idx
+            perf_stats["prompt"] = prompt
+            all_stats.append(perf_stats)
+
+        # Save audio and collect audio output
+        audio_path = None
+        audio_data = None
+        if audio is not None:
+            audio_data = audio.reshape(-1).detach().cpu().numpy()
+            audio_path = os.path.join(audio_dir, f"output_{idx:04d}.wav")
+            sf.write(
+                audio_path,
+                audio_data,
+                samplerate=24000,
+            )
+            audio_outputs.append(audio_data)
+        else:
+            audio_outputs.append(None)
+
+        # Save result
+        result = {
+            "idx": idx,
+            "prompt": prompt,
+            "output": output_text,
+            "audio_path": audio_path,
+            "perf_stats": perf_stats,
+        }
+        results.append(result)
+
+    # Aggregate statistics
+    aggregated_stats = aggregate_stats(all_stats)
+
+    # Save all results
+    results_path = os.path.join(output_dir, "results.json")
+    with open(results_path, "w", encoding="utf-8") as f:
+        json.dump(results, f, ensure_ascii=False, indent=2)
+
+    # Save aggregated stats
+    stats_path = os.path.join(output_dir, "perf_stats.json")
+    with open(stats_path, "w", encoding="utf-8") as f:
+        json.dump({"aggregated": aggregated_stats, "per_prompt": all_stats}, f, ensure_ascii=False, indent=2)
+
+    # Count saved audio files
+    num_audio_saved = sum(1 for a in audio_outputs if a is not None)
+    print(f"\nSaved {num_audio_saved} audio files to {audio_dir}/")
+
+    return aggregated_stats, results, audio_outputs
+
+
+def aggregate_stats(all_stats: list[dict]) -> dict:
+    """Aggregate performance statistics from multiple runs."""
+    if not all_stats:
+        return {}
+
+    keys = [
+        "thinker_tokens",
+        "thinker_time_s",
+        "thinker_tps",
+        "talker_tokens",
+        "talker_time_s",
+        "talker_tps",
+        "code2wav_tokens",
+        "code2wav_time_s",
+        "code2wav_tps",
+        "total_tokens",
+        "total_time_s",
+        "total_tps",
+    ]
+
+    aggregated = {
+        "num_samples": len(all_stats),
+    }
+
+    for key in keys:
+        values = [s.get(key, 0) for s in all_stats if key in s]
+        if values:
+            aggregated[f"{key}_sum"] = sum(values)
+            aggregated[f"{key}_avg"] = sum(values) / len(values)
+            aggregated[f"{key}_min"] = min(values)
+            aggregated[f"{key}_max"] = max(values)
+
+    # Calculate overall throughput
+    total_tokens = aggregated.get("total_tokens_sum", 0)
+    total_time = aggregated.get("total_time_s_sum", 0)
+    if total_time > 0:
+        aggregated["overall_tps"] = total_tokens / total_time
+
+    return aggregated
+
+
+def print_stats(stats: dict):
+    """Print performance statistics in a formatted way."""
+    print("\n" + "=" * 60)
+    print("Performance Statistics Summary")
+    print("=" * 60)
+
+    print(f"\nNumber of samples: {stats.get('num_samples', 0)}")
+
+    print("\n--- Thinker ---")
+    print(f"  Total tokens:  {stats.get('thinker_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('thinker_time_s_sum', 0):.2f}s")
+    print(f"  Avg TPS:       {stats.get('thinker_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('thinker_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('thinker_tps_max', 0):.2f}")
+
+    print("\n--- Talker ---")
+    print(f"  Total tokens:  {stats.get('talker_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('talker_time_s_sum', 0):.2f}s")
+    print(f"  Avg TPS:       {stats.get('talker_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('talker_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('talker_tps_max', 0):.2f}")
+
+    print("\n--- Code2Wav ---")
+    print(f"  Total tokens:  {stats.get('code2wav_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('code2wav_time_s_sum', 0):.2f}s")
+    print(f"  Avg TPS:       {stats.get('code2wav_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('code2wav_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('code2wav_tps_max', 0):.2f}")
+
+    print("\n--- Overall ---")
+    print(f"  Total tokens:  {stats.get('total_tokens_sum', 0):.0f}")
+    print(f"  Total time:    {stats.get('total_time_s_sum', 0):.2f}s")
+    print(f"  Overall TPS:   {stats.get('overall_tps', 0):.2f}")
+    print(f"  Avg TPS:       {stats.get('total_tps_avg', 0):.2f}")
+    print(f"  Min TPS:       {stats.get('total_tps_min', 0):.2f}")
+    print(f"  Max TPS:       {stats.get('total_tps_max', 0):.2f}")
+
+    print("=" * 60 + "\n")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="Qwen3-Omni Benchmark Script")
+    parser.add_argument(
+        "--prompts-file",
+        type=str,
+        default="benchmark/build_dataset/top100.txt",
+        help="Path to the prompts file (one prompt per line)",
+    )
+    parser.add_argument(
+        "--output-dir", type=str, default="benchmark_results", help="Directory to save benchmark results"
+    )
+    parser.add_argument("--model-path", type=str, default=MODEL_PATH, help="Path to the model")
+    parser.add_argument("--speaker", type=str, default="Ethan", help="Speaker voice for audio output")
+    parser.add_argument("--num-prompts", type=int, default=None, help="Number of prompts to process (default: all)")
+    args = parser.parse_args()
+
+    # Load model and processor
+    print(f"Loading model from {args.model_path}...")
+    model = Qwen3OmniMoeForConditionalGenerationWithLogging.from_pretrained(
+        args.model_path,
+        dtype="auto",
+        device_map="auto",
+        attn_implementation="flash_attention_2",
+    )
+    processor = Qwen3OmniMoeProcessor.from_pretrained(args.model_path)
+
+    # Benchmark mode
+    print(f"Loading prompts from {args.prompts_file}...")
+    prompts = load_prompts(args.prompts_file)
+
+    if args.num_prompts:
+        prompts = prompts[: args.num_prompts]
+
+    print(f"Running benchmark on {len(prompts)} prompts...")
+
+    aggregated_stats, results, audio_outputs = run_benchmark(
+        model=model,
+        processor=processor,
+        prompts=prompts,
+        output_dir=args.output_dir,
+        speaker=args.speaker,
+    )
+
+    print_stats(aggregated_stats)
+    print(f"\nResults saved to {args.output_dir}/")
+
+
+if __name__ == "__main__":
+    main()

benchmarks/qwen3-omni/vllm_omni/eval_qwen3_moe_omni.sh

@@ -0,0 +1,37 @@
+#!/bin/bash
+# Qwen3-Omni Benchmark Evaluation Script
+# This script must be run from the vllm-omni root directory
+
+# Get the directory where this script is located
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+
+# Navigate to vllm-omni root directory (4 levels up from script location)
+VLLM_OMNI_ROOT="$(cd "$SCRIPT_DIR/../../.." && pwd)"
+cd "$VLLM_OMNI_ROOT" || { echo "Error: Failed to navigate to vllm-omni directory"; exit 1; }
+
+echo "Working directory: $(pwd)"
+
+# Verify we're in the correct directory and run benchmark
+if [[ ! -d "benchmarks/qwen3-omni/vllm_omni" ]]; then
+    echo "Error: Not in vllm-omni root directory. Please run from vllm-omni folder."
+else
+    log_dir=benchmarks/qwen3-omni/vllm_omni/logs
+    outputs_dir=benchmarks/qwen3-omni/vllm_omni/outputs
+    end2end_script_path=examples/offline_inference/qwen3_omni/end2end.py
+    build_dataset_path=benchmarks/build_dataset/top100.txt
+
+    python $end2end_script_path --output-wav $outputs_dir \
+                      --query-type text \
+                      --txt-prompts $build_dataset_path \
+                      --log-stats \
+                      --log-dir $log_dir
+    echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
+    echo "  - omni_pipeline_text                       run dir/base name"
+    echo "  - omni_pipeline_text.orchestrator.stats.jsonl  orchestrator-stage latency stats"
+    echo "  - omni_pipeline_text.overall.stats.jsonl       overall latency/TPS stats"
+    echo "  - omni_pipeline_text.stage0.log                per-stage detailed logs"
+    echo "  - omni_pipeline_text.stage1.log"
+    echo "  - omni_pipeline_text.stage2.log"
+    echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
+    echo "  - outputs/             Generated txt and wav files, there should be 100 text and wav files generated respectively"
+fi

benchmarks/qwen3-tts/README.md

@@ -0,0 +1,100 @@
+# Qwen3-TTS Benchmark
+
+Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
+
+## Prerequisites
+
+```bash
+pip install matplotlib aiohttp soundfile numpy tqdm
+pip install qwen_tts  # for HF baseline
+```
+
+## Quick Start
+
+Run the full benchmark (vllm-omni + HF baseline) with a single command:
+
+```bash
+cd benchmarks/qwen3-tts
+bash run_benchmark.sh
+```
+
+Results (JSON + PNG plots) are saved to `results/`.
+
+### Common options
+
+```bash
+# Only vllm-omni (skip HF baseline)
+bash run_benchmark.sh --async-only
+
+# Only HF baseline
+bash run_benchmark.sh --hf-only
+
+# Use a different model (e.g. 1.7B)
+MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
+
+# Use bs16 config for higher throughput
+STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs16.yaml bash run_benchmark.sh --async-only
+
+# Custom GPU, prompt count, concurrency levels
+GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
+```
+
+## Manual Steps
+
+### 1) Start the vLLM-Omni server
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
+    "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
+    --omni --host 127.0.0.1 --port 8000 \
+    --stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
+    --trust-remote-code
+```
+
+### 2) Run online serving benchmark
+
+```bash
+python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
+    --port 8000 \
+    --num-prompts 50 \
+    --max-concurrency 1 4 10 \
+    --config-name "async_chunk" \
+    --result-dir results/
+```
+
+### 3) Run HuggingFace baseline
+
+```bash
+python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
+    --model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
+    --num-prompts 50 \
+    --gpu-device 0 \
+    --result-dir results/
+```
+
+### 4) Generate comparison plots
+
+```bash
+python benchmarks/qwen3-tts/plot_results.py \
+    --results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
+    --labels "vllm-omni" "hf_transformers" \
+    --output results/comparison.png
+```
+
+## Stage Configs
+
+| Config | max_num_seqs | Description |
+|--------|:------------:|-------------|
+| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
+| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
+
+All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
+
+The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
+
+## Metrics
+
+- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
+- **E2E (End-to-End Latency)**: Total time from request to complete audio response
+- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
+- **Throughput**: Total audio seconds generated per wall-clock second

benchmarks/qwen3-tts/plot_results.py

@@ -0,0 +1,254 @@
+"""Plot Qwen3-TTS benchmark results.
+
+Generates comparison bar charts similar to the async_chunk design doc:
+- TTFP (Time-to-First-Packet) across concurrency levels
+- E2E latency across concurrency levels
+- RTF (Real-Time Factor) across concurrency levels
+
+Usage:
+    # Compare two configs (async_chunk vs no_async_chunk):
+    python plot_results.py \
+        --results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
+        --labels "async_chunk" "no_async_chunk" \
+        --output results/qwen3_tts_benchmark.png
+
+    # Single config:
+    python plot_results.py \
+        --results results/bench_async_chunk_*.json \
+        --labels "async_chunk" \
+        --output results/qwen3_tts_benchmark.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def load_results(result_files: list[str]) -> list[list[dict]]:
+    """Load benchmark results from JSON files."""
+    all_results = []
+    for f in result_files:
+        with open(f) as fh:
+            data = json.load(fh)
+        all_results.append(data)
+    return all_results
+
+
+def plot_comparison(
+    all_results: list[list[dict]],
+    labels: list[str],
+    output_path: str,
+    title_prefix: str = "Qwen3-TTS",
+):
+    """Generate comparison bar charts."""
+    n_configs = len(all_results)
+
+    # Collect concurrency levels present in ALL configs (skip missing data)
+    all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
+    concurrencies = sorted(set.union(*all_concurrencies))
+
+    # Build data arrays, using None for missing concurrency levels
+    ttfp_data = {label: [] for label in labels}
+    e2e_data = {label: [] for label in labels}
+    rtf_data = {label: [] for label in labels}
+    throughput_data = {label: [] for label in labels}
+
+    for results, label in zip(all_results, labels):
+        conc_map = {r["concurrency"]: r for r in results}
+        for c in concurrencies:
+            r = conc_map.get(c)
+            ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
+            e2e_data[label].append(r["mean_e2e_ms"] if r else None)
+            rtf_data[label].append(r["mean_rtf"] if r else None)
+            throughput_data[label].append(r["audio_throughput"] if r else None)
+
+    fig, axes = plt.subplots(2, 2, figsize=(14, 10))
+    fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
+
+    x = np.arange(len(concurrencies))
+    width = 0.35 if n_configs == 2 else 0.5
+    if n_configs > 1:
+        offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
+    else:
+        offsets = [0]
+
+    colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
+
+    def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
+        bars = []
+        for i, (label, values) in enumerate(data_dict.items()):
+            # Replace None with 0 for plotting, but track which are missing
+            plot_values = [v if v is not None else 0 for v in values]
+            color = colors[i % len(colors)]
+            bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
+            bars.append(bar)
+            # Add value labels on bars (skip None/missing data)
+            max_val = max((v for v in values if v is not None), default=1)
+            for rect, val in zip(bar, values):
+                if val is not None and val > 0:
+                    ax.text(
+                        rect.get_x() + rect.get_width() / 2,
+                        rect.get_height() + max_val * 0.02,
+                        f"{val:{fmt}}",
+                        ha="center",
+                        va="bottom",
+                        fontsize=9,
+                        fontweight="bold",
+                    )
+        ax.set_xlabel("Concurrency", fontsize=12)
+        ax.set_ylabel(ylabel, fontsize=12)
+        ax.set_title(title, fontsize=13, fontweight="bold")
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        ax.legend(fontsize=10)
+        ax.grid(axis="y", alpha=0.3)
+        ax.set_axisbelow(True)
+
+    plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
+    plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
+    plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
+    plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_single_summary(results: list[dict], label: str, output_path: str):
+    """Generate a single-config summary with percentile breakdown."""
+    concurrencies = [r["concurrency"] for r in results]
+
+    fig, axes = plt.subplots(1, 3, figsize=(16, 5))
+    fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
+
+    # TTFP breakdown
+    ax = axes[0]
+    means = [r["mean_ttfp_ms"] for r in results]
+    medians = [r["median_ttfp_ms"] for r in results]
+    p90s = [r["p90_ttfp_ms"] for r in results]
+    p99s = [r["p99_ttfp_ms"] for r in results]
+    x = np.arange(len(concurrencies))
+    w = 0.2
+    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
+    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
+    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
+    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("TTFP (ms)")
+    ax.set_title("Time to First Audio Packet")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    # E2E breakdown
+    ax = axes[1]
+    means = [r["mean_e2e_ms"] for r in results]
+    medians = [r["median_e2e_ms"] for r in results]
+    p90s = [r["p90_e2e_ms"] for r in results]
+    p99s = [r["p99_e2e_ms"] for r in results]
+    ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
+    ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
+    ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
+    ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("E2E Latency (ms)")
+    ax.set_title("End-to-End Latency")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    # RTF
+    ax = axes[2]
+    means = [r["mean_rtf"] for r in results]
+    medians = [r["median_rtf"] for r in results]
+    ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
+    ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_xlabel("Concurrency")
+    ax.set_ylabel("RTF")
+    ax.set_title("Real-Time Factor")
+    ax.legend(fontsize=9)
+    ax.grid(axis="y", alpha=0.3)
+
+    plt.tight_layout()
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
+    """Print a markdown-formatted comparison table."""
+    concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
+
+    print("\n## Benchmark Results\n")
+    header = "| Metric | Concurrency |"
+    sep = "| --- | --- |"
+    for label in labels:
+        header += f" {label} |"
+        sep += " --- |"
+    print(header)
+    print(sep)
+
+    for metric, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput (audio-s/s)", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            row = f"| {metric} | {c} |"
+            for results in all_results:
+                conc_map = {r["concurrency"]: r for r in results}
+                val = conc_map.get(c, {}).get(key, 0)
+                row += f" {val:{fmt}} |"
+            print(row)
+
+    # Improvement calculation (only if 2 configs)
+    if len(all_results) == 2:
+        print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
+        print("| Metric | Concurrency | Improvement |")
+        print("| --- | --- | --- |")
+        for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
+            for c in concurrencies:
+                m0 = {r["concurrency"]: r for r in all_results[0]}
+                m1 = {r["concurrency"]: r for r in all_results[1]}
+                v0 = m0.get(c, {}).get(key, 0)
+                v1 = m1.get(c, {}).get(key, 0)
+                if v1 > 0:
+                    pct = (v1 - v0) / v1 * 100
+                    print(f"| {metric} | {c} | {pct:+.1f}% |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
+    parser.add_argument(
+        "--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
+    )
+    parser.add_argument(
+        "--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
+    )
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
+    parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
+
+    all_results = load_results(args.results)
+    print_comparison_table(all_results, args.labels)
+
+    Path(args.output).parent.mkdir(parents=True, exist_ok=True)
+
+    if len(all_results) == 1:
+        plot_single_summary(all_results[0], args.labels[0], args.output)
+    else:
+        plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)

benchmarks/qwen3-tts/results/.gitignore

@@ -0,0 +1,3 @@
+# Benchmark results are machine-specific - do not commit
+*
+!.gitignore

benchmarks/qwen3-tts/run_benchmark.sh

@@ -0,0 +1,272 @@
+#!/bin/bash
+# Qwen3-TTS Benchmark Runner
+#
+# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
+# Produces JSON results and comparison plots.
+#
+# Usage:
+#   # Full comparison (vllm-omni + HF):
+#   bash run_benchmark.sh
+#
+#   # Only vllm-omni async_chunk config:
+#   bash run_benchmark.sh --async-only
+#
+#   # Only HuggingFace baseline:
+#   bash run_benchmark.sh --hf-only
+#
+#   # vllm-omni only (skip HF):
+#   bash run_benchmark.sh --skip-hf
+#
+#   # Custom settings:
+#   GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
+#
+#   # Use 1.7B model:
+#   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
+#
+#   # Use batch_size=4 config:
+#   STAGE_CONFIG=vllm_omni/configs/qwen3_tts_bs4.yaml bash run_benchmark.sh --async-only
+#
+# Environment variables:
+#   GPU_DEVICE       - GPU index to use (default: 0)
+#   NUM_PROMPTS      - Number of prompts per concurrency level (default: 50)
+#   CONCURRENCY      - Space-separated concurrency levels (default: "1 4 10")
+#   MODEL            - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
+#   PORT             - Server port (default: 8000)
+#   GPU_MEM_TALKER   - gpu_memory_utilization for talker stage (default: 0.3)
+#   GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
+#   STAGE_CONFIG     - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
+
+# Defaults
+GPU_DEVICE="${GPU_DEVICE:-0}"
+NUM_PROMPTS="${NUM_PROMPTS:-50}"
+CONCURRENCY="${CONCURRENCY:-1 4 10}"
+MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
+PORT="${PORT:-8000}"
+GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
+GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
+NUM_WARMUPS="${NUM_WARMUPS:-3}"
+STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
+RESULT_DIR="${SCRIPT_DIR}/results"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+# Parse args
+RUN_ASYNC=true
+RUN_HF=true
+for arg in "$@"; do
+    case "$arg" in
+        --async-only) RUN_HF=false ;;
+        --hf-only) RUN_ASYNC=false ;;
+        --skip-hf) RUN_HF=false ;;
+    esac
+done
+
+mkdir -p "${RESULT_DIR}"
+
+echo "============================================================"
+echo " Qwen3-TTS Benchmark"
+echo "============================================================"
+echo " GPU:          ${GPU_DEVICE}"
+echo " Model:        ${MODEL}"
+echo " Prompts:      ${NUM_PROMPTS}"
+echo " Concurrency:  ${CONCURRENCY}"
+echo " Port:         ${PORT}"
+echo " Stage config: ${STAGE_CONFIG}"
+echo " Results:      ${RESULT_DIR}"
+echo "============================================================"
+
+# Prepare stage config with correct GPU device and memory settings
+prepare_config() {
+    local config_template="$1"
+    local config_name="$2"
+    local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
+
+    # Use sed to patch GPU device and memory utilization
+    sed \
+        -e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
+        -e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
+        -e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
+        "${config_template}" > "${output_path}"
+
+    echo "${output_path}"
+}
+
+# Start server and wait for it to be ready
+start_server() {
+    local stage_config="$1"
+    local config_name="$2"
+    local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
+
+    echo ""
+    echo "Starting server with config: ${config_name}"
+    echo "  Stage config: ${stage_config}"
+    echo "  Log file: ${log_file}"
+
+    VLLM_WORKER_MULTIPROC_METHOD=spawn \
+    CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
+    python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
+        --omni \
+        --host 127.0.0.1 \
+        --port "${PORT}" \
+        --stage-configs-path "${stage_config}" \
+        --stage-init-timeout 120 \
+        --trust-remote-code \
+        --disable-log-stats \
+        > "${log_file}" 2>&1 &
+
+    SERVER_PID=$!
+    echo "  Server PID: ${SERVER_PID}"
+
+    # Wait for server to be ready
+    echo "  Waiting for server to be ready..."
+    local max_wait=300
+    local waited=0
+    while [ ${waited} -lt ${max_wait} ]; do
+        if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
+            echo "  Server is ready! (waited ${waited}s)"
+            return 0
+        fi
+        # Check if process is still alive
+        if ! kill -0 ${SERVER_PID} 2>/dev/null; then
+            echo "  ERROR: Server process died. Check log: ${log_file}"
+            tail -20 "${log_file}"
+            return 1
+        fi
+        sleep 2
+        waited=$((waited + 2))
+    done
+
+    echo "  ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
+    kill ${SERVER_PID} 2>/dev/null || true
+    return 1
+}
+
+# Stop the server
+stop_server() {
+    if [ -n "${SERVER_PID:-}" ]; then
+        echo "  Stopping server (PID: ${SERVER_PID})..."
+        kill ${SERVER_PID} 2>/dev/null || true
+        wait ${SERVER_PID} 2>/dev/null || true
+        # Kill any remaining child processes on the port
+        local pids
+        pids=$(lsof -ti:${PORT} 2>/dev/null || true)
+        if [ -n "${pids}" ]; then
+            echo "  Cleaning up remaining processes on port ${PORT}..."
+            echo "${pids}" | xargs kill -9 2>/dev/null || true
+        fi
+        echo "  Server stopped."
+        SERVER_PID=""
+    fi
+}
+
+# Cleanup on exit
+trap 'stop_server' EXIT
+
+# Run benchmark for a given config
+run_bench() {
+    local config_name="$1"
+    local config_template="$2"
+
+    echo ""
+    echo "============================================================"
+    echo " Benchmarking: ${config_name}"
+    echo "============================================================"
+
+    local stage_config
+    stage_config=$(prepare_config "${config_template}" "${config_name}")
+
+    start_server "${stage_config}" "${config_name}"
+
+    # Convert concurrency string to args
+    local conc_args=""
+    for c in ${CONCURRENCY}; do
+        conc_args="${conc_args} ${c}"
+    done
+
+    cd "${PROJECT_ROOT}"
+    python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
+        --host 127.0.0.1 \
+        --port "${PORT}" \
+        --num-prompts "${NUM_PROMPTS}" \
+        --max-concurrency ${conc_args} \
+        --num-warmups "${NUM_WARMUPS}" \
+        --config-name "${config_name}" \
+        --result-dir "${RESULT_DIR}"
+
+    stop_server
+
+    # Allow GPU memory to settle
+    sleep 5
+}
+
+# Run vllm-omni benchmark
+if [ "${RUN_ASYNC}" = true ]; then
+    run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
+fi
+
+# Run HuggingFace baseline benchmark
+if [ "${RUN_HF}" = true ]; then
+    echo ""
+    echo "============================================================"
+    echo " Benchmarking: HuggingFace transformers (offline)"
+    echo "============================================================"
+
+    cd "${PROJECT_ROOT}"
+    python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
+        --model "${MODEL}" \
+        --num-prompts "${NUM_PROMPTS}" \
+        --num-warmups "${NUM_WARMUPS}" \
+        --gpu-device "${GPU_DEVICE}" \
+        --config-name "hf_transformers" \
+        --result-dir "${RESULT_DIR}"
+
+    # Allow GPU memory to settle
+    sleep 5
+fi
+
+# Plot results
+echo ""
+echo "============================================================"
+echo " Generating plots..."
+echo "============================================================"
+
+RESULT_FILES=""
+LABELS=""
+
+if [ "${RUN_ASYNC}" = true ]; then
+    ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
+    if [ -n "${ASYNC_FILE}" ]; then
+        RESULT_FILES="${ASYNC_FILE}"
+        LABELS="async_chunk"
+    fi
+fi
+
+if [ "${RUN_HF}" = true ]; then
+    HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
+    if [ -n "${HF_FILE}" ]; then
+        if [ -n "${RESULT_FILES}" ]; then
+            RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
+            LABELS="${LABELS} hf_transformers"
+        else
+            RESULT_FILES="${HF_FILE}"
+            LABELS="hf_transformers"
+        fi
+    fi
+fi
+
+if [ -n "${RESULT_FILES}" ]; then
+    python "${SCRIPT_DIR}/plot_results.py" \
+        --results ${RESULT_FILES} \
+        --labels ${LABELS} \
+        --output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
+fi
+
+echo ""
+echo "============================================================"
+echo " Benchmark complete!"
+echo " Results: ${RESULT_DIR}"
+echo "============================================================"

benchmarks/qwen3-tts/transformers/bench_tts_hf.py

@@ -0,0 +1,281 @@
+"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
+
+Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
+Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
+
+Usage:
+    python bench_tts_hf.py \
+        --model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
+        --num-prompts 50 \
+        --num-warmups 3 \
+        --gpu-device 0 \
+        --result-dir results/
+"""
+
+import argparse
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import numpy as np
+import soundfile as sf
+import torch
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 1  # always 1 for offline
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    # TTFP stats - not applicable for HF offline, set to E2E for compatibility
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    # E2E stats (ms)
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    # RTF stats
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    p99_rtf: float = 0.0
+    # Audio stats
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    # Per-request details
+    per_request: list = field(default_factory=list)
+
+
+def run_benchmark(args):
+    from qwen_tts import Qwen3TTSModel
+
+    device = f"cuda:{args.gpu_device}"
+    print(f"Loading model: {args.model} on {device}")
+    model = Qwen3TTSModel.from_pretrained(
+        args.model,
+        device_map=device,
+        dtype=torch.bfloat16,
+    )
+    print("Model loaded.")
+
+    # Build prompt list
+    prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
+
+    # Warmup
+    if args.num_warmups > 0:
+        print(f"Warming up with {args.num_warmups} requests...")
+        for i in range(args.num_warmups):
+            p = PROMPTS[i % len(PROMPTS)]
+            wavs, sr = model.generate_custom_voice(
+                text=p,
+                language=args.language,
+                speaker=args.voice,
+            )
+        # Sync GPU
+        torch.cuda.synchronize(device)
+        print("Warmup done.")
+
+    # Benchmark
+    print(f"Running {args.num_prompts} requests sequentially...")
+    e2e_times = []
+    rtfs = []
+    audio_durations = []
+    per_request = []
+    failed = 0
+
+    audio_dir = None
+    if args.save_audio:
+        audio_dir = Path(args.result_dir) / "audio_hf"
+        audio_dir.mkdir(parents=True, exist_ok=True)
+
+    total_start = time.perf_counter()
+
+    for i, prompt in enumerate(prompts):
+        try:
+            torch.cuda.synchronize(device)
+            st = time.perf_counter()
+
+            wavs, sr = model.generate_custom_voice(
+                text=prompt,
+                language=args.language,
+                speaker=args.voice,
+            )
+
+            torch.cuda.synchronize(device)
+            elapsed = time.perf_counter() - st
+
+            # Compute audio duration
+            audio_samples = wavs[0]
+            if isinstance(audio_samples, torch.Tensor):
+                audio_samples = audio_samples.cpu().numpy()
+            audio_dur = len(audio_samples) / sr
+
+            rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
+
+            e2e_times.append(elapsed)
+            rtfs.append(rtf)
+            audio_durations.append(audio_dur)
+            per_request.append(
+                {
+                    "e2e_ms": elapsed * 1000,
+                    "ttfp_ms": elapsed * 1000,  # no streaming, TTFP = E2E
+                    "rtf": rtf,
+                    "audio_duration_s": audio_dur,
+                    "prompt": prompt,
+                }
+            )
+
+            if audio_dir:
+                sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
+
+            if (i + 1) % 10 == 0 or i == 0:
+                print(
+                    f"  [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms  rtf={rtf:.3f}  audio={audio_dur:.2f}s"
+                )
+
+        except Exception as e:
+            print(f"  [{i + 1}/{args.num_prompts}] FAILED: {e}")
+            failed += 1
+
+    total_duration = time.perf_counter() - total_start
+    completed = len(e2e_times)
+
+    # Compute stats
+    result = BenchmarkResult(
+        config_name=args.config_name,
+        concurrency=1,
+        num_prompts=args.num_prompts,
+        completed=completed,
+        failed=failed,
+        duration_s=total_duration,
+    )
+
+    if e2e_times:
+        e2e_ms = [t * 1000 for t in e2e_times]
+
+        result.mean_e2e_ms = float(np.mean(e2e_ms))
+        result.median_e2e_ms = float(np.median(e2e_ms))
+        result.std_e2e_ms = float(np.std(e2e_ms))
+        result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
+        result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
+        result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
+
+        # For HF offline, TTFP = E2E (no streaming)
+        result.mean_ttfp_ms = result.mean_e2e_ms
+        result.median_ttfp_ms = result.median_e2e_ms
+        result.std_ttfp_ms = result.std_e2e_ms
+        result.p90_ttfp_ms = result.p90_e2e_ms
+        result.p95_ttfp_ms = result.p95_e2e_ms
+        result.p99_ttfp_ms = result.p99_e2e_ms
+
+        result.mean_rtf = float(np.mean(rtfs))
+        result.median_rtf = float(np.median(rtfs))
+        result.std_rtf = float(np.std(rtfs))
+        result.p99_rtf = float(np.percentile(rtfs, 99))
+
+        result.mean_audio_duration_s = float(np.mean(audio_durations))
+        result.total_audio_duration_s = float(np.sum(audio_durations))
+        result.audio_throughput = result.total_audio_duration_s / total_duration
+        result.request_throughput = completed / total_duration
+        result.per_request = per_request
+
+    # Print summary in standardized performance template
+    W = 50
+    print("")
+    print(f"{'=' * W}")
+    print(f"{'Serving Benchmark Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Successful requests:':<40}{completed:<10}")
+    print(f"{'Failed requests:':<40}{failed:<10}")
+    print(f"{'Maximum request concurrency:':<40}{1:<10}")
+    print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
+    print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'End-to-end Latency':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
+    print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
+    print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
+    print(f"{'=' * W}")
+    print(f"{'Audio Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
+    print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Time to First Packet':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
+    print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
+    print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Real Time Factor':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
+    print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
+    print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
+    print(f"{'=' * W}")
+    print("")
+
+    # Save results (as a list with single concurrency=1 entry, matching serve format)
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump([asdict(result)], f, indent=2)
+    print(f"Results saved to {result_file}")
+
+    return result
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
+    parser.add_argument(
+        "--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
+    )
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--gpu-device", type=int, default=0)
+    parser.add_argument("--voice", type=str, default="Vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument(
+        "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
+    )
+    parser.add_argument("--result-dir", type=str, default="results")
+    parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    run_benchmark(args)

benchmarks/qwen3-tts/vllm_omni/bench_async_chunk.py

@@ -0,0 +1,301 @@
+"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
+concurrency levels for both async_chunk modes. Saves results as JSON.
+
+Usage:
+    # Run against a server already serving with a given config:
+    python bench_async_chunk.py \
+        --host 127.0.0.1 --port 8000 \
+        --config-name async_chunk_on \
+        --num-prompts 50 \
+        --max-concurrency 1 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0
+    e2e: float = 0.0
+    audio_bytes: int = 0
+    audio_duration: float = 0.0
+    rtf: float = 0.0
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0
+    request_throughput: float = 0.0
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    return num_bytes / sample_width / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": stream,
+        "response_format": "pcm",
+    }
+
+    result = RequestResult(prompt=prompt)
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+    stream: bool = True,
+) -> BenchmarkResult:
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
+    session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
+
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = [
+            send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
+            for i in range(num_warmups)
+        ]
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    print(f"\n{'=' * 60}")
+    print(f"  Concurrency: {max_concurrency}  |  Completed: {bench.completed}  |  Failed: {bench.failed}")
+    print(f"  Duration: {duration:.2f}s  |  Throughput: {bench.request_throughput:.2f} req/s")
+    print(
+        f"  TTFP (ms):  mean={bench.mean_ttfp_ms:.1f}  median={bench.median_ttfp_ms:.1f}"
+        f"  p90={bench.p90_ttfp_ms:.1f}  p99={bench.p99_ttfp_ms:.1f}"
+    )
+    print(
+        f"  E2E (ms):   mean={bench.mean_e2e_ms:.1f}  median={bench.median_e2e_ms:.1f}"
+        f"  p90={bench.p90_e2e_ms:.1f}  p99={bench.p99_e2e_ms:.1f}"
+    )
+    print(f"  RTF:        mean={bench.mean_rtf:.3f}  median={bench.median_rtf:.3f}")
+    print(f"  Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
+    print(f"{'=' * 60}\n")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+            stream=args.stream,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50)
+    parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument("--stream", action="store_true", default=True)
+    parser.add_argument("--no-stream", dest="stream", action="store_false")
+    parser.add_argument("--config-name", type=str, default="async_chunk_on")
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))

benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py

@@ -0,0 +1,348 @@
+"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
+
+Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
+across configurable concurrency levels. Saves results as JSON for plotting.
+
+Usage:
+    python bench_tts_serve.py \
+        --host 127.0.0.1 --port 8000 \
+        --num-prompts 50 \
+        --max-concurrency 1 4 10 \
+        --result-dir results/
+"""
+
+import argparse
+import asyncio
+import json
+import time
+from dataclasses import asdict, dataclass, field
+from datetime import datetime
+from pathlib import Path
+
+import aiohttp
+import numpy as np
+from tqdm.asyncio import tqdm
+
+PROMPTS = [
+    "Hello, welcome to the voice synthesis benchmark test.",
+    "She said she would be here by noon, but nobody showed up.",
+    "The quick brown fox jumps over the lazy dog near the riverbank.",
+    "I can't believe how beautiful the sunset looks from up here on the mountain.",
+    "Please remember to bring your identification documents to the appointment tomorrow morning.",
+    "Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
+    "The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
+    "After the meeting, we should discuss the quarterly results and plan for the next phase.",
+    "Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
+    "The train leaves at half past seven, so we need to arrive at the station before then.",
+    "Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
+    "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
+]
+
+
+@dataclass
+class RequestResult:
+    success: bool = False
+    ttfp: float = 0.0  # Time to first audio packet (seconds)
+    e2e: float = 0.0  # End-to-end latency (seconds)
+    audio_bytes: int = 0  # Total audio bytes received
+    audio_duration: float = 0.0  # Audio duration in seconds (estimated from PCM)
+    rtf: float = 0.0  # Real-time factor = e2e / audio_duration
+    prompt: str = ""
+    error: str = ""
+
+
+@dataclass
+class BenchmarkResult:
+    config_name: str = ""
+    concurrency: int = 0
+    num_prompts: int = 0
+    completed: int = 0
+    failed: int = 0
+    duration_s: float = 0.0
+    # TTFP stats (ms)
+    mean_ttfp_ms: float = 0.0
+    median_ttfp_ms: float = 0.0
+    std_ttfp_ms: float = 0.0
+    p90_ttfp_ms: float = 0.0
+    p95_ttfp_ms: float = 0.0
+    p99_ttfp_ms: float = 0.0
+    # E2E stats (ms)
+    mean_e2e_ms: float = 0.0
+    median_e2e_ms: float = 0.0
+    std_e2e_ms: float = 0.0
+    p90_e2e_ms: float = 0.0
+    p95_e2e_ms: float = 0.0
+    p99_e2e_ms: float = 0.0
+    # RTF stats
+    mean_rtf: float = 0.0
+    median_rtf: float = 0.0
+    std_rtf: float = 0.0
+    p99_rtf: float = 0.0
+    # Audio stats
+    mean_audio_duration_s: float = 0.0
+    total_audio_duration_s: float = 0.0
+    audio_throughput: float = 0.0  # audio_duration / wall_time
+    request_throughput: float = 0.0  # requests / second
+    # Per-request details
+    per_request: list = field(default_factory=list)
+
+
+def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
+    """Convert raw PCM byte count to duration in seconds."""
+    num_samples = num_bytes / sample_width
+    return num_samples / sample_rate
+
+
+async def send_tts_request(
+    session: aiohttp.ClientSession,
+    api_url: str,
+    prompt: str,
+    voice: str = "vivian",
+    language: str = "English",
+    pbar: tqdm | None = None,
+) -> RequestResult:
+    """Send a streaming TTS request and measure latency metrics."""
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": True,
+        "response_format": "pcm",
+    }
+
+    result = RequestResult(prompt=prompt)
+    st = time.perf_counter()
+
+    try:
+        async with session.post(api_url, json=payload) as response:
+            if response.status != 200:
+                result.error = f"HTTP {response.status}: {await response.text()}"
+                result.success = False
+                return result
+
+            first_chunk = True
+            total_bytes = 0
+
+            async for chunk in response.content.iter_any():
+                if first_chunk and len(chunk) > 0:
+                    result.ttfp = time.perf_counter() - st
+                    first_chunk = False
+                total_bytes += len(chunk)
+
+            result.e2e = time.perf_counter() - st
+            result.audio_bytes = total_bytes
+            result.audio_duration = pcm_bytes_to_duration(total_bytes)
+
+            if result.audio_duration > 0:
+                result.rtf = result.e2e / result.audio_duration
+            result.success = True
+
+    except Exception as e:
+        result.error = str(e)
+        result.success = False
+        result.e2e = time.perf_counter() - st
+
+    if pbar:
+        pbar.update(1)
+    return result
+
+
+async def run_benchmark(
+    host: str,
+    port: int,
+    num_prompts: int,
+    max_concurrency: int,
+    num_warmups: int = 3,
+    voice: str = "vivian",
+    language: str = "English",
+) -> BenchmarkResult:
+    """Run benchmark at a given concurrency level."""
+    api_url = f"http://{host}:{port}/v1/audio/speech"
+
+    connector = aiohttp.TCPConnector(
+        limit=max_concurrency,
+        limit_per_host=max_concurrency,
+        keepalive_timeout=60,
+    )
+    session = aiohttp.ClientSession(
+        connector=connector,
+        timeout=aiohttp.ClientTimeout(total=600),
+    )
+
+    # Warmup
+    if num_warmups > 0:
+        print(f"  Warming up with {num_warmups} requests...")
+        warmup_tasks = []
+        for i in range(num_warmups):
+            prompt = PROMPTS[i % len(PROMPTS)]
+            warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
+        await asyncio.gather(*warmup_tasks)
+        print("  Warmup done.")
+
+    # Build request list
+    request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
+
+    # Run benchmark
+    print(f"  Running {num_prompts} requests with concurrency={max_concurrency}...")
+    semaphore = asyncio.Semaphore(max_concurrency)
+    pbar = tqdm(total=num_prompts, desc=f"  concurrency={max_concurrency}")
+
+    async def limited_request(prompt):
+        async with semaphore:
+            return await send_tts_request(session, api_url, prompt, voice, language, pbar)
+
+    start_time = time.perf_counter()
+    tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
+    results: list[RequestResult] = await asyncio.gather(*tasks)
+    duration = time.perf_counter() - start_time
+    pbar.close()
+
+    await session.close()
+
+    # Compute stats
+    successful = [r for r in results if r.success]
+    failed = [r for r in results if not r.success]
+
+    bench = BenchmarkResult(
+        concurrency=max_concurrency,
+        num_prompts=num_prompts,
+        completed=len(successful),
+        failed=len(failed),
+        duration_s=duration,
+    )
+
+    if successful:
+        ttfps = [r.ttfp * 1000 for r in successful]  # convert to ms
+        e2es = [r.e2e * 1000 for r in successful]
+        rtfs = [r.rtf for r in successful]
+        audio_durs = [r.audio_duration for r in successful]
+
+        bench.mean_ttfp_ms = float(np.mean(ttfps))
+        bench.median_ttfp_ms = float(np.median(ttfps))
+        bench.std_ttfp_ms = float(np.std(ttfps))
+        bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
+        bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
+        bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
+
+        bench.mean_e2e_ms = float(np.mean(e2es))
+        bench.median_e2e_ms = float(np.median(e2es))
+        bench.std_e2e_ms = float(np.std(e2es))
+        bench.p90_e2e_ms = float(np.percentile(e2es, 90))
+        bench.p95_e2e_ms = float(np.percentile(e2es, 95))
+        bench.p99_e2e_ms = float(np.percentile(e2es, 99))
+
+        bench.mean_rtf = float(np.mean(rtfs))
+        bench.median_rtf = float(np.median(rtfs))
+        bench.std_rtf = float(np.std(rtfs))
+        bench.p99_rtf = float(np.percentile(rtfs, 99))
+
+        bench.mean_audio_duration_s = float(np.mean(audio_durs))
+        bench.total_audio_duration_s = float(np.sum(audio_durs))
+        bench.audio_throughput = bench.total_audio_duration_s / duration
+        bench.request_throughput = len(successful) / duration
+
+        bench.per_request = [
+            {
+                "ttfp_ms": r.ttfp * 1000,
+                "e2e_ms": r.e2e * 1000,
+                "rtf": r.rtf,
+                "audio_duration_s": r.audio_duration,
+                "prompt": r.prompt,
+            }
+            for r in successful
+        ]
+
+    # Print summary in standardized performance template
+    W = 50
+    print("")
+    print(f"{'=' * W}")
+    print(f"{'Serving Benchmark Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Successful requests:':<40}{bench.completed:<10}")
+    print(f"{'Failed requests:':<40}{bench.failed:<10}")
+    print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
+    print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
+    print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'End-to-end Latency':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
+    print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
+    print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
+    print(f"{'=' * W}")
+    print(f"{'Audio Result':^{W}}")
+    print(f"{'=' * W}")
+    print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
+    print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Time to First Packet':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
+    print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
+    print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
+    print(f"{'-' * W}")
+    print(f"{'Real Time Factor':^{W}}")
+    print(f"{'-' * W}")
+    print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
+    print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
+    print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
+    print(f"{'=' * W}")
+    print("")
+
+    if failed:
+        for r in failed[:3]:
+            print(f"  [ERROR] {r.error[:200]}")
+
+    return bench
+
+
+async def main(args):
+    all_results = []
+
+    for concurrency in args.max_concurrency:
+        result = await run_benchmark(
+            host=args.host,
+            port=args.port,
+            num_prompts=args.num_prompts,
+            max_concurrency=concurrency,
+            num_warmups=args.num_warmups,
+            voice=args.voice,
+            language=args.language,
+        )
+        result.config_name = args.config_name
+        all_results.append(asdict(result))
+
+    # Save results
+    result_dir = Path(args.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
+
+    with open(result_file, "w") as f:
+        json.dump(all_results, f, indent=2)
+    print(f"Results saved to {result_file}")
+
+    return all_results
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
+    parser.add_argument("--host", type=str, default="127.0.0.1")
+    parser.add_argument("--port", type=int, default=8000)
+    parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
+    parser.add_argument(  # noqa: E501
+        "--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
+    )
+    parser.add_argument("--num-warmups", type=int, default=3)
+    parser.add_argument("--voice", type=str, default="vivian")
+    parser.add_argument("--language", type=str, default="English")
+    parser.add_argument(
+        "--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
+    )
+    parser.add_argument("--result-dir", type=str, default="results")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    asyncio.run(main(args))

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

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

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

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

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

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

benchmarks/qwen3-tts/vllm_omni/plot_async_chunk.py

@@ -0,0 +1,249 @@
+"""Plot TTFP comparison: async_chunk off vs on.
+
+Generates a bar chart with improvement arrows, matching the Qwen3-Omni
+async_chunk benchmark figure style.
+
+Usage:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --output results/qwen3_tts_async_chunk_ttfp.png
+
+    # Also supports E2E and RTF metrics:
+    python plot_async_chunk.py \
+        --off results/bench_async_chunk_off_*.json \
+        --on results/bench_async_chunk_on_*.json \
+        --metric e2e \
+        --output results/qwen3_tts_async_chunk_e2e.png
+"""
+
+import argparse
+import json
+from pathlib import Path
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+METRIC_CONFIG = {
+    "ttfp": {
+        "key": "mean_ttfp_ms",
+        "ylabel": "TTFP (s)",
+        "title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "e2e": {
+        "key": "mean_e2e_ms",
+        "ylabel": "E2E (s)",
+        "title": "E2E Latency - Qwen3-TTS, by concurrency",
+        "to_seconds": True,
+    },
+    "rtf": {
+        "key": "mean_rtf",
+        "ylabel": "RTF",
+        "title": "Real-Time Factor - Qwen3-TTS, by concurrency",
+        "to_seconds": False,
+    },
+}
+
+
+def load_results(path: str) -> list[dict]:
+    with open(path) as f:
+        return json.load(f)
+
+
+def plot_ttfp_comparison(
+    off_results: list[dict],
+    on_results: list[dict],
+    metric: str,
+    output_path: str,
+    title_override: str | None = None,
+):
+    cfg = METRIC_CONFIG[metric]
+    key = cfg["key"]
+    to_seconds = cfg["to_seconds"]
+
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    off_vals = []
+    on_vals = []
+    for c in concurrencies:
+        v_off = off_map[c][key]
+        v_on = on_map[c][key]
+        if to_seconds:
+            v_off /= 1000.0
+            v_on /= 1000.0
+        off_vals.append(v_off)
+        on_vals.append(v_on)
+
+    fig, ax = plt.subplots(figsize=(8, 6))
+
+    x = np.arange(len(concurrencies))
+    width = 0.3
+
+    ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
+    ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+    # Draw improvement arrows and labels
+    for i in range(len(concurrencies)):
+        v_off = off_vals[i]
+        v_on = on_vals[i]
+        if v_on > 0:
+            improvement = v_off / v_on
+        else:
+            improvement = float("inf")
+
+        # Arrow from top of off-bar to top of on-bar
+        arrow_start_x = x[i] - width / 2
+        arrow_start_y = v_off * 0.95
+        arrow_end_x = x[i] + width / 2
+        arrow_end_y = v_on * 1.05
+
+        ax.annotate(
+            "",
+            xy=(arrow_end_x, arrow_end_y),
+            xytext=(arrow_start_x, arrow_start_y),
+            arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+        )
+
+        # Improvement label
+        label_x = (arrow_start_x + arrow_end_x) / 2
+        label_y = arrow_start_y + (v_off - v_on) * 0.15
+        ax.text(
+            label_x,
+            label_y,
+            f"{improvement:.1f}x improvement",
+            ha="center",
+            va="bottom",
+            fontsize=10,
+            color="red",
+            fontweight="bold",
+        )
+
+    title = title_override or cfg["title"]
+    ax.set_title(title, fontsize=13, fontweight="bold")
+    ax.set_ylabel(cfg["ylabel"], fontsize=12)
+    ax.set_xlabel("Max concurrency", fontsize=12)
+    ax.set_xticks(x)
+    ax.set_xticklabels([str(c) for c in concurrencies])
+    ax.set_yscale("log")
+    ax.legend(loc="upper left", fontsize=11)
+    ax.grid(axis="y", alpha=0.3, linestyle="--")
+    ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
+    """Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    fig, axes = plt.subplots(1, 3, figsize=(18, 6))
+    fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
+
+    for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
+        cfg = METRIC_CONFIG[metric]
+        key = cfg["key"]
+        to_seconds = cfg["to_seconds"]
+
+        off_vals = []
+        on_vals = []
+        for c in concurrencies:
+            v_off = off_map[c][key]
+            v_on = on_map[c][key]
+            if to_seconds:
+                v_off /= 1000.0
+                v_on /= 1000.0
+            off_vals.append(v_off)
+            on_vals.append(v_on)
+
+        x = np.arange(len(concurrencies))
+        width = 0.3
+        ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
+        ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
+
+        for i in range(len(concurrencies)):
+            if on_vals[i] > 0:
+                improvement = off_vals[i] / on_vals[i]
+                ax.annotate(
+                    "",
+                    xy=(x[i] + width / 2, on_vals[i] * 1.05),
+                    xytext=(x[i] - width / 2, off_vals[i] * 0.95),
+                    arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
+                )
+                label_y = off_vals[i] * 0.85
+                ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
+
+        ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
+        ax.set_ylabel(cfg["ylabel"], fontsize=11)
+        ax.set_xlabel("Max concurrency", fontsize=11)
+        ax.set_xticks(x)
+        ax.set_xticklabels([str(c) for c in concurrencies])
+        if metric != "rtf":
+            ax.set_yscale("log")
+        ax.legend(fontsize=9)
+        ax.grid(axis="y", alpha=0.3, linestyle="--")
+        ax.set_axisbelow(True)
+
+    plt.tight_layout()
+    Path(output_path).parent.mkdir(parents=True, exist_ok=True)
+    plt.savefig(output_path, dpi=150, bbox_inches="tight")
+    print(f"Plot saved to {output_path}")
+    plt.close()
+
+
+def print_table(off_results: list[dict], on_results: list[dict]):
+    off_map = {r["concurrency"]: r for r in off_results}
+    on_map = {r["concurrency"]: r for r in on_results}
+    concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
+
+    print("\n## Benchmark Results: async_chunk off vs on\n")
+    print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
+    print("| --- | --- | --- | --- | --- |")
+
+    for name, key, fmt in [
+        ("TTFP (ms)", "mean_ttfp_ms", ".1f"),
+        ("E2E (ms)", "mean_e2e_ms", ".1f"),
+        ("RTF", "mean_rtf", ".3f"),
+        ("Throughput", "audio_throughput", ".2f"),
+    ]:
+        for c in concurrencies:
+            v_off = off_map[c].get(key, 0)
+            v_on = on_map[c].get(key, 0)
+            if v_on > 0 and key != "audio_throughput":
+                ratio = f"{v_off / v_on:.1f}x"
+            elif v_off > 0 and key == "audio_throughput":
+                ratio = f"{v_on / v_off:.1f}x"
+            else:
+                ratio = "N/A"
+            print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
+    parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
+    parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
+    parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
+    parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
+    parser.add_argument("--title", type=str, default=None, help="Custom title override")
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    off_results = load_results(args.off)
+    on_results = load_results(args.on)
+
+    print_table(off_results, on_results)
+
+    if args.metric == "all":
+        plot_all_metrics(off_results, on_results, args.output)
+    else:
+        plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)

benchmarks/qwen3-tts/vllm_omni/run_async_chunk_benchmark.sh

@@ -0,0 +1,167 @@
+#!/bin/bash
+# Qwen3-TTS async_chunk on vs off Benchmark
+#
+# Starts two servers (async_chunk on and off), benchmarks both,
+# and generates comparison plots.
+#
+# Usage:
+#   bash run_async_chunk_benchmark.sh
+#
+# Environment variables:
+#   GPU_DEVICE       - GPU index (default: 0)
+#   MODEL            - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
+#   NUM_PROMPTS      - Prompts per concurrency level (default: 50)
+#   CONCURRENCY      - Space-separated concurrency levels (default: "1 10")
+#   PORT_ON          - Port for async_chunk on server (default: 8000)
+#   PORT_OFF         - Port for async_chunk off server (default: 8001)
+
+set -euo pipefail
+
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
+cd "$PROJECT_ROOT"
+
+GPU_DEVICE="${GPU_DEVICE:-0}"
+MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
+NUM_PROMPTS="${NUM_PROMPTS:-50}"
+CONCURRENCY="${CONCURRENCY:-1 10}"
+NUM_WARMUPS="${NUM_WARMUPS:-3}"
+PORT_ON="${PORT_ON:-8000}"
+PORT_OFF="${PORT_OFF:-8001}"
+RESULT_DIR="${SCRIPT_DIR}/results"
+TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
+
+STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
+STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
+
+mkdir -p "${RESULT_DIR}"
+
+echo "============================================================"
+echo " Qwen3-TTS async_chunk Benchmark"
+echo "============================================================"
+echo " GPU:            ${GPU_DEVICE}"
+echo " Model:          ${MODEL}"
+echo " Prompts:        ${NUM_PROMPTS}"
+echo " Concurrency:    ${CONCURRENCY}"
+echo " Port (on/off):  ${PORT_ON} / ${PORT_OFF}"
+echo " Results:        ${RESULT_DIR}"
+echo "============================================================"
+
+cleanup() {
+    echo "Cleaning up servers..."
+    kill "$PID_ON" 2>/dev/null || true
+    kill "$PID_OFF" 2>/dev/null || true
+    wait "$PID_ON" 2>/dev/null || true
+    wait "$PID_OFF" 2>/dev/null || true
+}
+trap cleanup EXIT
+
+wait_for_server() {
+    local port=$1
+    local name=$2
+    local max_wait=300
+    local elapsed=0
+    echo "Waiting for ${name} server on port ${port}..."
+    while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
+        sleep 5
+        elapsed=$((elapsed + 5))
+        if [ $elapsed -ge $max_wait ]; then
+            echo "ERROR: ${name} server failed to start within ${max_wait}s"
+            exit 1
+        fi
+    done
+    echo "${name} server ready (${elapsed}s)"
+}
+
+# ---- Phase 1: Start async_chunk ON server ----
+echo ""
+echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_ON}" \
+    --host 0.0.0.0 --port "${PORT_ON}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
+PID_ON=$!
+
+wait_for_server "${PORT_ON}" "async_chunk_on"
+
+echo "[Phase 1] Benchmarking async_chunk ON..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_ON}" \
+    --config-name "async_chunk_on" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 1] Stopping async_chunk ON server..."
+kill "$PID_ON" 2>/dev/null || true
+wait "$PID_ON" 2>/dev/null || true
+sleep 5
+
+# ---- Phase 2: Start async_chunk OFF server ----
+echo ""
+echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
+CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
+    --stage-configs-path "${STAGE_CONFIG_OFF}" \
+    --host 0.0.0.0 --port "${PORT_OFF}" \
+    --trust-remote-code --enforce-eager --omni \
+    > "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
+PID_OFF=$!
+
+wait_for_server "${PORT_OFF}" "async_chunk_off"
+
+echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
+# shellcheck disable=SC2086
+python "${SCRIPT_DIR}/bench_async_chunk.py" \
+    --host 127.0.0.1 --port "${PORT_OFF}" \
+    --config-name "async_chunk_off" \
+    --num-prompts "${NUM_PROMPTS}" \
+    --max-concurrency ${CONCURRENCY} \
+    --num-warmups "${NUM_WARMUPS}" \
+    --no-stream \
+    --result-dir "${RESULT_DIR}"
+
+echo "[Phase 2] Stopping async_chunk OFF server..."
+kill "$PID_OFF" 2>/dev/null || true
+wait "$PID_OFF" 2>/dev/null || true
+
+# ---- Phase 3: Plot results ----
+echo ""
+echo "[Phase 3] Generating plots..."
+
+# Find the latest result files
+RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
+RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
+
+if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
+    echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
+    exit 1
+fi
+
+echo "  ON results:  ${RESULT_ON}"
+echo "  OFF results: ${RESULT_OFF}"
+
+# TTFP comparison (main figure)
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric ttfp \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+
+# All metrics comparison
+python "${SCRIPT_DIR}/plot_async_chunk.py" \
+    --off "${RESULT_OFF}" \
+    --on "${RESULT_ON}" \
+    --metric all \
+    --output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+
+echo ""
+echo "============================================================"
+echo " Benchmark complete!"
+echo " Results: ${RESULT_DIR}/"
+echo " Plots:"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
+echo "   - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
+echo "============================================================"

collect_env.py

@@ -0,0 +1,760 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+# ruff: noqa
+# code borrowed from https://github.com/pytorch/pytorch/blob/main/torch/utils/collect_env.py
+
+import datetime
+import locale
+import os
+import subprocess
+import sys
+
+# Unlike the rest of the PyTorch this file must be python2 compliant.
+# This script outputs relevant system environment info
+# Run it with `python collect_env.py` or `python -m torch.utils.collect_env`
+from collections import namedtuple
+
+import regex as re
+
+from vllm.envs import environment_variables
+
+try:
+    import torch
+
+    TORCH_AVAILABLE = True
+except (ImportError, NameError, AttributeError, OSError):
+    TORCH_AVAILABLE = False
+
+# System Environment Information
+SystemEnv = namedtuple(
+    "SystemEnv",
+    [
+        "torch_version",
+        "is_debug_build",
+        "cuda_compiled_version",
+        "gcc_version",
+        "clang_version",
+        "cmake_version",
+        "os",
+        "libc_version",
+        "python_version",
+        "python_platform",
+        "is_cuda_available",
+        "cuda_runtime_version",
+        "cuda_module_loading",
+        "nvidia_driver_version",
+        "nvidia_gpu_models",
+        "cudnn_version",
+        "pip_version",  # 'pip' or 'pip3'
+        "pip_packages",
+        "conda_packages",
+        "hip_compiled_version",
+        "hip_runtime_version",
+        "miopen_runtime_version",
+        "caching_allocator_config",
+        "is_xnnpack_available",
+        "cpu_info",
+        "rocm_version",  # vllm specific field
+        "vllm_version",  # vllm specific field
+        "vllm_omni_version",  # vllm-omni specific field
+        "vllm_build_flags",  # vllm specific field
+        "gpu_topo",  # vllm specific field
+        "env_vars",
+    ],
+)
+
+DEFAULT_CONDA_PATTERNS = {
+    "torch",
+    "numpy",
+    "cudatoolkit",
+    "soumith",
+    "mkl",
+    "magma",
+    "triton",
+    "optree",
+    "nccl",
+    "transformers",
+    "zmq",
+    "nvidia",
+    "pynvml",
+    "flashinfer-python",
+}
+
+DEFAULT_PIP_PATTERNS = {
+    "torch",
+    "numpy",
+    "mypy",
+    "flake8",
+    "triton",
+    "optree",
+    "onnx",
+    "nccl",
+    "transformers",
+    "zmq",
+    "nvidia",
+    "pynvml",
+    "flashinfer-python",
+}
+
+
+def run(command):
+    """Return (return-code, stdout, stderr)."""
+    shell = True if type(command) is str else False
+    try:
+        p = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=shell)
+        raw_output, raw_err = p.communicate()
+        rc = p.returncode
+        if get_platform() == "win32":
+            enc = "oem"
+        else:
+            enc = locale.getpreferredencoding()
+        output = raw_output.decode(enc)
+        if command == "nvidia-smi topo -m":
+            # don't remove the leading whitespace of `nvidia-smi topo -m`
+            #   because they are meaningful
+            output = output.rstrip()
+        else:
+            output = output.strip()
+        err = raw_err.decode(enc)
+        return rc, output, err.strip()
+
+    except FileNotFoundError:
+        cmd_str = command if isinstance(command, str) else command[0]
+        return 127, "", f"Command not found: {cmd_str}"
+
+
+def run_and_read_all(run_lambda, command):
+    """Run command using run_lambda; reads and returns entire output if rc is 0."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    return out
+
+
+def run_and_parse_first_match(run_lambda, command, regex):
+    """Run command using run_lambda, returns the first regex match if it exists."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    match = re.search(regex, out)
+    if match is None:
+        return None
+    return match.group(1)
+
+
+def run_and_return_first_line(run_lambda, command):
+    """Run command using run_lambda and returns first line if output is not empty."""
+    rc, out, _ = run_lambda(command)
+    if rc != 0:
+        return None
+    return out.split("\n")[0]
+
+
+def get_conda_packages(run_lambda, patterns=None):
+    if patterns is None:
+        patterns = DEFAULT_CONDA_PATTERNS
+    conda = os.environ.get("CONDA_EXE", "conda")
+    out = run_and_read_all(run_lambda, [conda, "list"])
+    if out is None:
+        return out
+
+    return "\n".join(
+        line for line in out.splitlines() if not line.startswith("#") and any(name in line for name in patterns)
+    )
+
+
+def get_gcc_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, "gcc --version", r"gcc (.*)")
+
+
+def get_clang_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, "clang --version", r"clang version (.*)")
+
+
+def get_cmake_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, "cmake --version", r"cmake (.*)")
+
+
+def get_nvidia_driver_version(run_lambda):
+    if get_platform() == "darwin":
+        cmd = "kextstat | grep -i cuda"
+        return run_and_parse_first_match(run_lambda, cmd, r"com[.]nvidia[.]CUDA [(](.*?)[)]")
+    smi = get_nvidia_smi()
+    return run_and_parse_first_match(run_lambda, smi, r"Driver Version: (.*?) ")
+
+
+def get_gpu_info(run_lambda):
+    if get_platform() == "darwin" or (
+        TORCH_AVAILABLE and hasattr(torch.version, "hip") and torch.version.hip is not None
+    ):
+        if TORCH_AVAILABLE and torch.cuda.is_available():
+            if torch.version.hip is not None:
+                prop = torch.cuda.get_device_properties(0)
+                if hasattr(prop, "gcnArchName"):
+                    gcnArch = " ({})".format(prop.gcnArchName)
+                else:
+                    gcnArch = "NoGCNArchNameOnOldPyTorch"
+            else:
+                gcnArch = ""
+            return torch.cuda.get_device_name(None) + gcnArch
+        return None
+    smi = get_nvidia_smi()
+    uuid_regex = re.compile(r" \(UUID: .+?\)")
+    rc, out, _ = run_lambda(smi + " -L")
+    if rc != 0:
+        return None
+    # Anonymize GPUs by removing their UUID
+    return re.sub(uuid_regex, "", out)
+
+
+def get_running_cuda_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, "nvcc --version", r"release .+ V(.*)")
+
+
+def get_cudnn_version(run_lambda):
+    """Return a list of libcudnn.so; it's hard to tell which one is being used."""
+    if get_platform() == "win32":
+        system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
+        cuda_path = os.environ.get("CUDA_PATH", "%CUDA_PATH%")
+        where_cmd = os.path.join(system_root, "System32", "where")
+        cudnn_cmd = '{} /R "{}\\bin" cudnn*.dll'.format(where_cmd, cuda_path)
+    elif get_platform() == "darwin":
+        # CUDA libraries and drivers can be found in /usr/local/cuda/. See
+        # https://docs.nvidia.com/cuda/cuda-installation-guide-mac-os-x/index.html#install
+        # https://docs.nvidia.com/deeplearning/sdk/cudnn-install/index.html#installmac
+        # Use CUDNN_LIBRARY when cudnn library is installed elsewhere.
+        cudnn_cmd = "ls /usr/local/cuda/lib/libcudnn*"
+    else:
+        cudnn_cmd = 'ldconfig -p | grep libcudnn | rev | cut -d" " -f1 | rev'
+    rc, out, _ = run_lambda(cudnn_cmd)
+    # find will return 1 if there are permission errors or if not found
+    if len(out) == 0 or (rc != 1 and rc != 0):
+        l = os.environ.get("CUDNN_LIBRARY")
+        if l is not None and os.path.isfile(l):
+            return os.path.realpath(l)
+        return None
+    files_set = set()
+    for fn in out.split("\n"):
+        fn = os.path.realpath(fn)  # eliminate symbolic links
+        if os.path.isfile(fn):
+            files_set.add(fn)
+    if not files_set:
+        return None
+    # Alphabetize the result because the order is non-deterministic otherwise
+    files = sorted(files_set)
+    if len(files) == 1:
+        return files[0]
+    result = "\n".join(files)
+    return "Probably one of the following:\n{}".format(result)
+
+
+def get_nvidia_smi():
+    # Note: nvidia-smi is currently available only on Windows and Linux
+    smi = "nvidia-smi"
+    if get_platform() == "win32":
+        system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
+        program_files_root = os.environ.get("PROGRAMFILES", "C:\\Program Files")
+        legacy_path = os.path.join(program_files_root, "NVIDIA Corporation", "NVSMI", smi)
+        new_path = os.path.join(system_root, "System32", smi)
+        smis = [new_path, legacy_path]
+        for candidate_smi in smis:
+            if os.path.exists(candidate_smi):
+                smi = '"{}"'.format(candidate_smi)
+                break
+    return smi
+
+
+def get_rocm_version(run_lambda):
+    """Returns the ROCm version if available, otherwise 'N/A'."""
+    return run_and_parse_first_match(run_lambda, "hipcc --version", r"HIP version: (\S+)")
+
+
+def get_vllm_version():
+    from vllm import __version__, __version_tuple__
+
+    if __version__ == "dev":
+        return "N/A (dev)"
+    version_str = __version_tuple__[-1]
+    if isinstance(version_str, str) and version_str.startswith("g"):
+        # it's a dev build
+        if "." in version_str:
+            # it's a dev build containing local changes
+            git_sha = version_str.split(".")[0][1:]
+            date = version_str.split(".")[-1][1:]
+            return f"{__version__} (git sha: {git_sha}, date: {date})"
+        else:
+            # it's a dev build without local changes
+            git_sha = version_str[1:]  # type: ignore
+            return f"{__version__} (git sha: {git_sha})"
+    return __version__
+
+
+def get_vllm_omni_version(run_lambda):
+    try:
+        import vllm_omni
+        from vllm_omni import __version__, __version_tuple__
+
+        version_str = __version_tuple__[-1]
+        if isinstance(version_str, str) and version_str.startswith("g"):
+            if "." in version_str:
+                git_sha = version_str.split(".")[0][1:]
+                date = version_str.split(".")[-1][1:]
+                return f"{__version__} (git sha: {git_sha}, date: {date})"
+            else:
+                git_sha = version_str[1:]
+                return f"{__version__} (git sha: {git_sha})"
+
+        package_dir = os.path.dirname(os.path.abspath(vllm_omni.__file__))
+        git_sha = run_and_read_all(run_lambda, f"git -C {package_dir} rev-parse --short HEAD")
+        if git_sha:
+            return f"{__version__} (git sha: {git_sha})"
+
+        return __version__
+    except ImportError:
+        return "N/A (vllm_omni not installed)"
+
+
+def summarize_vllm_build_flags():
+    # This could be a static method if the flags are constant, or dynamic if you need to check environment variables, etc.
+    return "CUDA Archs: {}; ROCm: {}".format(
+        os.environ.get("TORCH_CUDA_ARCH_LIST", "Not Set"),
+        "Enabled" if os.environ.get("ROCM_HOME") else "Disabled",
+    )
+
+
+def get_gpu_topo(run_lambda):
+    output = None
+
+    if get_platform() == "linux":
+        output = run_and_read_all(run_lambda, "nvidia-smi topo -m")
+        if output is None:
+            output = run_and_read_all(run_lambda, "rocm-smi --showtopo")
+
+    return output
+
+
+def get_cpu_info(run_lambda):
+    rc, out, err = 0, "", ""
+    if get_platform() == "linux":
+        rc, out, err = run_lambda("lscpu")
+    elif get_platform() == "win32":
+        rc, out, err = run_lambda(
+            "wmic cpu get Name,Manufacturer,Family,Architecture,ProcessorType,DeviceID, \
+        CurrentClockSpeed,MaxClockSpeed,L2CacheSize,L2CacheSpeed,Revision /VALUE"
+        )
+    elif get_platform() == "darwin":
+        rc, out, err = run_lambda("sysctl -n machdep.cpu.brand_string")
+    cpu_info = "None"
+    if rc == 0:
+        cpu_info = out
+    else:
+        cpu_info = err
+    return cpu_info
+
+
+def get_platform():
+    if sys.platform.startswith("linux"):
+        return "linux"
+    elif sys.platform.startswith("win32"):
+        return "win32"
+    elif sys.platform.startswith("cygwin"):
+        return "cygwin"
+    elif sys.platform.startswith("darwin"):
+        return "darwin"
+    else:
+        return sys.platform
+
+
+def get_mac_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, "sw_vers -productVersion", r"(.*)")
+
+
+def get_windows_version(run_lambda):
+    system_root = os.environ.get("SYSTEMROOT", "C:\\Windows")
+    wmic_cmd = os.path.join(system_root, "System32", "Wbem", "wmic")
+    findstr_cmd = os.path.join(system_root, "System32", "findstr")
+    return run_and_read_all(run_lambda, "{} os get Caption | {} /v Caption".format(wmic_cmd, findstr_cmd))
+
+
+def get_lsb_version(run_lambda):
+    return run_and_parse_first_match(run_lambda, "lsb_release -a", r"Description:\t(.*)")
+
+
+def check_release_file(run_lambda):
+    return run_and_parse_first_match(run_lambda, "cat /etc/*-release", r'PRETTY_NAME="(.*)"')
+
+
+def get_os(run_lambda):
+    from platform import machine
+
+    platform = get_platform()
+
+    if platform == "win32" or platform == "cygwin":
+        return get_windows_version(run_lambda)
+
+    if platform == "darwin":
+        version = get_mac_version(run_lambda)
+        if version is None:
+            return None
+        return "macOS {} ({})".format(version, machine())
+
+    if platform == "linux":
+        # Ubuntu/Debian based
+        desc = get_lsb_version(run_lambda)
+        if desc is not None:
+            return "{} ({})".format(desc, machine())
+
+        # Try reading /etc/*-release
+        desc = check_release_file(run_lambda)
+        if desc is not None:
+            return "{} ({})".format(desc, machine())
+
+        return "{} ({})".format(platform, machine())
+
+    # Unknown platform
+    return platform
+
+
+def get_python_platform():
+    import platform
+
+    return platform.platform()
+
+
+def get_libc_version():
+    import platform
+
+    if get_platform() != "linux":
+        return "N/A"
+    return "-".join(platform.libc_ver())
+
+
+def is_uv_venv():
+    if os.environ.get("UV"):
+        return True
+    pyvenv_cfg_path = os.path.join(sys.prefix, "pyvenv.cfg")
+    if os.path.exists(pyvenv_cfg_path):
+        with open(pyvenv_cfg_path, "r") as f:
+            return any(line.startswith("uv = ") for line in f)
+    return False
+
+
+def get_pip_packages(run_lambda, patterns=None):
+    """Return `pip list` output. Note: will also find conda-installed pytorch and numpy packages."""
+    if patterns is None:
+        patterns = DEFAULT_PIP_PATTERNS
+
+    def run_with_pip():
+        try:
+            import importlib.util
+
+            pip_spec = importlib.util.find_spec("pip")
+            pip_available = pip_spec is not None
+        except ImportError:
+            pip_available = False
+
+        if pip_available:
+            cmd = [sys.executable, "-mpip", "list", "--format=freeze"]
+        elif is_uv_venv():
+            print("uv is set")
+            cmd = ["uv", "pip", "list", "--format=freeze"]
+        else:
+            raise RuntimeError("Could not collect pip list output (pip or uv module not available)")
+
+        out = run_and_read_all(run_lambda, cmd)
+        return "\n".join(line for line in out.splitlines() if any(name in line for name in patterns))
+
+    pip_version = "pip3" if sys.version[0] == "3" else "pip"
+    out = run_with_pip()
+    return pip_version, out
+
+
+def get_cachingallocator_config():
+    ca_config = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
+    return ca_config
+
+
+def get_cuda_module_loading_config():
+    if TORCH_AVAILABLE and torch.cuda.is_available():
+        torch.cuda.init()
+        config = os.environ.get("CUDA_MODULE_LOADING", "")
+        return config
+    else:
+        return "N/A"
+
+
+def is_xnnpack_available():
+    if TORCH_AVAILABLE:
+        import torch.backends.xnnpack
+
+        return str(torch.backends.xnnpack.enabled)  # type: ignore[attr-defined]
+    else:
+        return "N/A"
+
+
+def get_env_vars():
+    env_vars = ""
+    secret_terms = ("secret", "token", "api", "access", "password")
+    report_prefix = (
+        "TORCH",
+        "NCCL",
+        "PYTORCH",
+        "CUDA",
+        "CUBLAS",
+        "CUDNN",
+        "OMP_",
+        "MKL_",
+        "NVIDIA",
+    )
+    for k, v in os.environ.items():
+        if any(term in k.lower() for term in secret_terms):
+            continue
+        if k in environment_variables:
+            env_vars = env_vars + "{}={}".format(k, v) + "\n"
+        if k.startswith(report_prefix):
+            env_vars = env_vars + "{}={}".format(k, v) + "\n"
+
+    return env_vars
+
+
+def get_env_info():
+    run_lambda = run
+    pip_version, pip_list_output = get_pip_packages(run_lambda)
+
+    if TORCH_AVAILABLE:
+        version_str = torch.__version__
+        debug_mode_str = str(torch.version.debug)
+        cuda_available_str = str(torch.cuda.is_available())
+        cuda_version_str = torch.version.cuda
+        if not hasattr(torch.version, "hip") or torch.version.hip is None:  # cuda version
+            hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A"
+        else:  # HIP version
+
+            def get_version_or_na(cfg, prefix):
+                _lst = [s.rsplit(None, 1)[-1] for s in cfg if prefix in s]
+                return _lst[0] if _lst else "N/A"
+
+            cfg = torch._C._show_config().split("\n")
+            hip_runtime_version = get_version_or_na(cfg, "HIP Runtime")
+            miopen_runtime_version = get_version_or_na(cfg, "MIOpen")
+            cuda_version_str = "N/A"
+            hip_compiled_version = torch.version.hip
+    else:
+        version_str = debug_mode_str = cuda_available_str = cuda_version_str = "N/A"
+        hip_compiled_version = hip_runtime_version = miopen_runtime_version = "N/A"
+
+    sys_version = sys.version.replace("\n", " ")
+
+    conda_packages = get_conda_packages(run_lambda)
+
+    rocm_version = get_rocm_version(run_lambda)
+    vllm_version = get_vllm_version()
+    vllm_omni_version = get_vllm_omni_version(run_lambda)
+    vllm_build_flags = summarize_vllm_build_flags()
+    gpu_topo = get_gpu_topo(run_lambda)
+
+    return SystemEnv(
+        torch_version=version_str,
+        is_debug_build=debug_mode_str,
+        python_version="{} ({}-bit runtime)".format(sys_version, sys.maxsize.bit_length() + 1),
+        python_platform=get_python_platform(),
+        is_cuda_available=cuda_available_str,
+        cuda_compiled_version=cuda_version_str,
+        cuda_runtime_version=get_running_cuda_version(run_lambda),
+        cuda_module_loading=get_cuda_module_loading_config(),
+        nvidia_gpu_models=get_gpu_info(run_lambda),
+        nvidia_driver_version=get_nvidia_driver_version(run_lambda),
+        cudnn_version=get_cudnn_version(run_lambda),
+        hip_compiled_version=hip_compiled_version,
+        hip_runtime_version=hip_runtime_version,
+        miopen_runtime_version=miopen_runtime_version,
+        pip_version=pip_version,
+        pip_packages=pip_list_output,
+        conda_packages=conda_packages,
+        os=get_os(run_lambda),
+        libc_version=get_libc_version(),
+        gcc_version=get_gcc_version(run_lambda),
+        clang_version=get_clang_version(run_lambda),
+        cmake_version=get_cmake_version(run_lambda),
+        caching_allocator_config=get_cachingallocator_config(),
+        is_xnnpack_available=is_xnnpack_available(),
+        cpu_info=get_cpu_info(run_lambda),
+        rocm_version=rocm_version,
+        vllm_version=vllm_version,
+        vllm_omni_version=vllm_omni_version,
+        vllm_build_flags=vllm_build_flags,
+        gpu_topo=gpu_topo,
+        env_vars=get_env_vars(),
+    )
+
+
+env_info_fmt = """
+==============================
+        System Info
+==============================
+OS                           : {os}
+GCC version                  : {gcc_version}
+Clang version                : {clang_version}
+CMake version                : {cmake_version}
+Libc version                 : {libc_version}
+
+==============================
+       PyTorch Info
+==============================
+PyTorch version              : {torch_version}
+Is debug build               : {is_debug_build}
+CUDA used to build PyTorch   : {cuda_compiled_version}
+ROCM used to build PyTorch   : {hip_compiled_version}
+
+==============================
+      Python Environment
+==============================
+Python version               : {python_version}
+Python platform              : {python_platform}
+
+==============================
+       CUDA / GPU Info
+==============================
+Is CUDA available            : {is_cuda_available}
+CUDA runtime version         : {cuda_runtime_version}
+CUDA_MODULE_LOADING set to   : {cuda_module_loading}
+GPU models and configuration : {nvidia_gpu_models}
+Nvidia driver version        : {nvidia_driver_version}
+cuDNN version                : {cudnn_version}
+HIP runtime version          : {hip_runtime_version}
+MIOpen runtime version       : {miopen_runtime_version}
+Is XNNPACK available         : {is_xnnpack_available}
+
+==============================
+          CPU Info
+==============================
+{cpu_info}
+
+==============================
+Versions of relevant libraries
+==============================
+{pip_packages}
+{conda_packages}
+""".strip()
+
+# both the above code and the following code use `strip()` to
+# remove leading/trailing whitespaces, so we need to add a newline
+# in between to separate the two sections
+env_info_fmt += "\n\n"
+
+env_info_fmt += """
+==============================
+         vLLM Info
+==============================
+ROCM Version                 : {rocm_version}
+vLLM Version                 : {vllm_version}
+vLLM-Omni Version            : {vllm_omni_version}
+vLLM Build Flags:
+  {vllm_build_flags}
+GPU Topology:
+  {gpu_topo}
+
+==============================
+     Environment Variables
+==============================
+{env_vars}
+""".strip()
+
+
+def pretty_str(envinfo):
+    def replace_nones(dct, replacement="Could not collect"):
+        for key in dct.keys():
+            if dct[key] is not None:
+                continue
+            dct[key] = replacement
+        return dct
+
+    def replace_bools(dct, true="Yes", false="No"):
+        for key in dct.keys():
+            if dct[key] is True:
+                dct[key] = true
+            elif dct[key] is False:
+                dct[key] = false
+        return dct
+
+    def prepend(text, tag="[prepend]"):
+        lines = text.split("\n")
+        updated_lines = [tag + line for line in lines]
+        return "\n".join(updated_lines)
+
+    def replace_if_empty(text, replacement="No relevant packages"):
+        if text is not None and len(text) == 0:
+            return replacement
+        return text
+
+    def maybe_start_on_next_line(string):
+        # If `string` is multiline, prepend a \n to it.
+        if string is not None and len(string.split("\n")) > 1:
+            return "\n{}\n".format(string)
+        return string
+
+    mutable_dict = envinfo._asdict()
+
+    # If nvidia_gpu_models is multiline, start on the next line
+    mutable_dict["nvidia_gpu_models"] = maybe_start_on_next_line(envinfo.nvidia_gpu_models)
+
+    # If the machine doesn't have CUDA, report some fields as 'No CUDA'
+    dynamic_cuda_fields = [
+        "cuda_runtime_version",
+        "nvidia_gpu_models",
+        "nvidia_driver_version",
+    ]
+    all_cuda_fields = dynamic_cuda_fields + ["cudnn_version"]
+    all_dynamic_cuda_fields_missing = all(mutable_dict[field] is None for field in dynamic_cuda_fields)
+    if TORCH_AVAILABLE and not torch.cuda.is_available() and all_dynamic_cuda_fields_missing:
+        for field in all_cuda_fields:
+            mutable_dict[field] = "No CUDA"
+        if envinfo.cuda_compiled_version is None:
+            mutable_dict["cuda_compiled_version"] = "None"
+
+    # Replace True with Yes, False with No
+    mutable_dict = replace_bools(mutable_dict)
+
+    # Replace all None objects with 'Could not collect'
+    mutable_dict = replace_nones(mutable_dict)
+
+    # If either of these are '', replace with 'No relevant packages'
+    mutable_dict["pip_packages"] = replace_if_empty(mutable_dict["pip_packages"])
+    mutable_dict["conda_packages"] = replace_if_empty(mutable_dict["conda_packages"])
+
+    # Tag conda and pip packages with a prefix
+    # If they were previously None, they'll show up as ie '[conda] Could not collect'
+    if mutable_dict["pip_packages"]:
+        mutable_dict["pip_packages"] = prepend(mutable_dict["pip_packages"], "[{}] ".format(envinfo.pip_version))
+    if mutable_dict["conda_packages"]:
+        mutable_dict["conda_packages"] = prepend(mutable_dict["conda_packages"], "[conda] ")
+    mutable_dict["cpu_info"] = envinfo.cpu_info
+    return env_info_fmt.format(**mutable_dict)
+
+
+def get_pretty_env_info():
+    return pretty_str(get_env_info())
+
+
+def main():
+    print("Collecting environment information...")
+    output = get_pretty_env_info()
+    print(output)
+
+    if TORCH_AVAILABLE and hasattr(torch, "utils") and hasattr(torch.utils, "_crash_handler"):
+        minidump_dir = torch.utils._crash_handler.DEFAULT_MINIDUMP_DIR
+        if sys.platform == "linux" and os.path.exists(minidump_dir):
+            dumps = [os.path.join(minidump_dir, dump) for dump in os.listdir(minidump_dir)]
+            latest = max(dumps, key=os.path.getctime)
+            ctime = os.path.getctime(latest)
+            creation_time = datetime.datetime.fromtimestamp(ctime).strftime("%Y-%m-%d %H:%M:%S")
+            msg = (
+                "\n*** Detected a minidump at {} created on {}, ".format(latest, creation_time)
+                + "if this is related to your bug please include it when you file a report ***"
+            )
+            print(msg, file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

docker/Dockerfile.ci

@@ -0,0 +1,20 @@
+ARG VLLM_BASE_IMAGE=vllm/vllm-openai
+ARG VLLM_BASE_TAG=v0.18.0
+FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
+ARG APP_DIR=/workspace/vllm-omni
+WORKDIR ${APP_DIR}
+COPY . .
+
+# Install system dependencies
+RUN apt-get update && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Install vllm-omni into the same uv-managed Python environment used by the base image.
+# Use bash -c so that $(python3 -c ...) is expanded inside the container.
+RUN uv pip install --system --no-cache-dir ".[dev]"
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENTRYPOINT []

docker/Dockerfile.npu

@@ -0,0 +1,31 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.17.0rc1
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+WORKDIR /vllm-workspace/vllm
+RUN git fetch origin --tags && git checkout v0.18.0
+
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# Install vllm-ascend
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip cache purge
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+COPY . .
+
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []

docker/Dockerfile.npu.a3

@@ -0,0 +1,31 @@
+ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
+ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
+FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
+
+WORKDIR /vllm-workspace/vllm
+RUN git fetch origin --tags && git checkout v0.18.0
+
+WORKDIR /vllm-workspace/vllm-ascend
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+# Install vllm-ascend
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
+    python3 -m pip cache purge
+
+ARG APP_DIR=/vllm-workspace/vllm-omni
+WORKDIR ${APP_DIR}
+COPY . .
+
+RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
+    source /usr/local/Ascend/nnal/atb/set_env.sh && \
+    export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
+    python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT []

docker/Dockerfile.rocm

@@ -0,0 +1,53 @@
+ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
+FROM ${BASE_IMAGE} AS base
+
+# Declare a variable to know if we want to use the nightly build or the stable build.
+# NOTE: REMINDER to vLLM-Omni rebase maintainer
+# Remember to set `USE_NIGHTLY_BUILD` to 0, when switching back to stable vLLM docker image
+# 1. If vLLM-omni maintainer is forced to use custom commits
+#    during rebasing, they can change the two variables
+# 2. Whenever vLLM upstream has released a stable version,
+#    we should swap over to use stable release ASAP.
+#    We should avoid relying on custom commits.
+ARG USE_NIGHTLY_BUILD=0
+ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
+ARG ARG_PYTORCH_ROCM_ARCH
+ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
+
+ARG COMMON_WORKDIR=/app
+WORKDIR ${COMMON_WORKDIR}
+
+# Step 1: Setup - Install system dependencies
+RUN apt-get update && \
+    apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
+    apt-get clean && \
+    rm -rf /var/lib/apt/lists/*
+
+# Step 2: Conditionally reinstall vllm from source for nightly builds
+RUN if [ "${USE_NIGHTLY_BUILD}" = "1" ]; then \
+        python3 -m pip uninstall -y vllm && rm -rf vllm && \
+        git clone https://github.com/vllm-project/vllm.git && \
+        cd vllm && \
+        git checkout ${VLLM_VERSION_OR_COMMIT_HASH} && \
+        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; \
+    fi
+
+# Step 3: Copy vllm-omni code and install without uv
+RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
+COPY . ${COMMON_WORKDIR}/vllm-omni
+RUN cd ${COMMON_WORKDIR}/vllm-omni && uv pip install --python "$(python3 -c 'import sys; print(sys.executable)')" --no-cache-dir ".[dev]" --no-build-isolation
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+FROM base AS test
+
+CMD ["/bin/bash"]
+ENTRYPOINT []
+
+# Set entrypoint for vllm-openai official images
+FROM base AS vllm-openai
+ENTRYPOINT ["vllm", "serve", "--omni"]

docker/Dockerfile.xpu

@@ -0,0 +1,152 @@
+# Argument to configure vllm base image if pre-built
+ARG VLLM_BASE=vllm-base
+
+FROM intel/deep-learning-essentials:2025.3.2-0-devel-ubuntu24.04 AS vllm-base
+
+WORKDIR /workspace/
+
+ARG PYTHON_VERSION=3.12
+ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
+
+RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
+    echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
+
+RUN apt clean && apt-get update -y && \
+    apt-get install -y --no-install-recommends --fix-missing \
+    curl \
+    espeak-ng \
+    ffmpeg \
+    git \
+    libsndfile1 \
+    libsm6 \
+    libxext6 \
+    libgl1 \
+    lsb-release \
+    libaio-dev \
+    numactl \
+    wget \
+    vim \
+    python3.12 \
+    python3.12-dev \
+    python3-pip
+
+RUN apt update && apt upgrade -y && \
+    apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
+
+# Install UMD with fixed version
+RUN mkdir neo && \
+    cd neo && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
+    wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
+    wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
+    dpkg -i *.deb && \
+    cd .. && \
+    rm -rf neo
+
+ENV PATH="/root/.local/bin:$PATH"
+ENV VIRTUAL_ENV="/opt/venv"
+ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
+RUN curl -LsSf https://astral.sh/uv/install.sh | sh
+RUN uv venv --python ${PYTHON_VERSION} --seed ${VIRTUAL_ENV}
+ENV PATH="$VIRTUAL_ENV/bin:$PATH"
+
+# This oneccl contains the BMG support which is not the case for default version of oneapi 2025.2.
+ARG ONECCL_INSTALLER="intel-oneccl-2021.15.7.8_offline.sh"
+RUN wget "https://github.com/uxlfoundation/oneCCL/releases/download/2021.15.7/${ONECCL_INSTALLER}" && \
+    bash "${ONECCL_INSTALLER}" -a --silent --eula accept && \
+    rm "${ONECCL_INSTALLER}" && \
+    echo "source /opt/intel/oneapi/setvars.sh --force" >> /root/.bashrc && \
+    echo "source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force" >> /root/.bashrc
+RUN rm -f /opt/intel/oneapi/ccl/latest && \
+    ln -s /opt/intel/oneapi/ccl/2021.15 /opt/intel/oneapi/ccl/latest
+
+SHELL ["bash", "-c"]
+CMD ["bash", "-c", "source /root/.bashrc && exec bash"]
+
+WORKDIR /workspace/
+ENV UV_HTTP_TIMEOUT=500
+
+# Configure package index for XPU
+ENV PIP_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
+ENV UV_INDEX_STRATEGY="unsafe-best-match"
+ENV UV_LINK_MODE="copy"
+
+ARG VLLM_VERSION=v0.18.0
+RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
+WORKDIR /workspace/vllm
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --upgrade pip && \
+    uv pip install -r requirements/xpu.txt
+
+ # used for suffix method speculative decoding
+ # build deps for proto + nanobind-based extensions to set up the build environment
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install grpcio-tools protobuf nanobind
+ # arctic-inference is built from source which needs torch-xpu properly installed first
+RUN --mount=type=cache,target=/root/.cache/uv \
+    source /opt/intel/oneapi/setvars.sh --force && \
+    source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
+    export CMAKE_PREFIX_PATH="$(python -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}" && \
+    uv pip install --no-build-isolation arctic-inference==0.1.1
+
+ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
+
+ENV VLLM_TARGET_DEVICE=xpu
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install --no-build-isolation .
+
+CMD ["/bin/bash"]
+
+FROM vllm-base AS vllm-openai
+
+# install additional dependencies for openai api server
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip install accelerate hf_transfer pytest pytest_asyncio lm_eval[api] modelscope
+
+# install development dependencies (for testing)
+RUN uv pip install -e tests/vllm_test_utils
+
+# install nixl from source code
+ENV NIXL_VERSION=0.7.0
+RUN python /workspace/vllm/tools/install_nixl_from_source_ubuntu.py
+
+# ensure vllm is properly installed
+RUN python -c "import vllm, inspect; print(vllm.__file__)"
+RUN uv pip show vllm
+
+CMD ["/bin/bash"]
+
+ENTRYPOINT []
+
+FROM ${VLLM_BASE} AS vllm-omni
+
+WORKDIR /workspace/vllm-omni
+COPY . .
+
+ENV VLLM_OMNI_TARGET_DEVICE=xpu
+RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
+
+# FIX triton
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall triton triton-xpu && \
+    uv pip install triton-xpu==3.6.0
+
+# remove torch bundled oneccl to avoid conflicts
+RUN --mount=type=cache,target=/root/.cache/uv \
+    uv pip uninstall oneccl oneccl-devel
+
+FROM vllm-omni AS vllm-omni-openai
+
+RUN ln -sf /usr/bin/python3 /usr/bin/python
+
+ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
+
+ENTRYPOINT ["vllm", "serve", "--omni"]

docs/.nav.yml

@@ -3,16 +3,47 @@ nav:
 - User Guide:
   - Getting Started:
     - getting_started/quickstart.md
-    - getting_started/installation
+    - getting_started/installation/*
+  - Serving:
+    - OpenAI-Compatible API:
+      - Diffusion Chat API: serving/diffusion_chat_api.md
+      - Image Generation: serving/image_generation_api.md
+      - Image Edit: serving/image_edit_api.md
+      - Text to Speech: serving/speech_api.md
   - Examples:
     - examples/README.md
     - Offline Inference:
-      - Offline Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
-      - Offline Example of vLLM-Omni for Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
-      - Qwen-Image Offline Inference: user_guide/examples/offline_inference/qwen_image.md
+      - BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
+      - CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
+      - Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
+      - GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
+      - Helios Video Generation: user_guide/examples/offline_inference/helios.md
+      - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
+      - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
+      - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
+      - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
+      - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
+      - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
+      - Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
+      - Qwen3-TTS: 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
+      - Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
+      - X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
     - Online Serving:
-      - Online serving Example of vLLM-Omni for Qwen2.5-Omni: user_guide/examples/online_serving/qwen2_5_omni.md
-      - Online serving Example of vLLM-Omni for Qwen3-omni: user_guide/examples/online_serving/qwen3_omni.md
+      - BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
+      - vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
+      - Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
+      - GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
+      - Image-To-Image: user_guide/examples/online_serving/image_to_image.md
+      - Image-To-Video: user_guide/examples/online_serving/image_to_video.md
+      - Online serving Example of vLLM-Omni for MiMo-Audio: user_guide/examples/online_serving/mimo_audio.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: user_guide/examples/online_serving/qwen3_tts.md
+      - Text-To-Image: user_guide/examples/online_serving/text_to_image.md
+      - Text-To-Video: user_guide/examples/online_serving/text_to_video.md
   - General:
     - usage/*
   - Configuration:
@@ -20,25 +51,65 @@ nav:
     - configuration/*
   - Models:
     - models/supported_models.md
+  - Features:
+    - Sleep Mode: features/sleep_mode.md
+    - Diffusion Features:
+      - Acceleration Overview: user_guide/diffusion_acceleration.md
+      - TeaCache: user_guide/diffusion/teacache.md
+      - Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
+      - Quantization:
+        - Overview: user_guide/diffusion/quantization/overview.md
+        - FP8: user_guide/diffusion/quantization/fp8.md
+        - Int8: user_guide/diffusion/quantization/int8.md
+        - GGUF: user_guide/diffusion/quantization/gguf.md
+      - Frame Interpolation: user_guide/diffusion/frame_interpolation.md
+      - Step Execution: user_guide/diffusion/step_execution.md
+      - Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
+      - CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
+      - LoRA: user_guide/diffusion/lora.md
+      - Hybrid Sharded Data Parallel: design/feature/hsdp.md
+      - Custom Pipeline: features/custom_pipeline.md
+    - ComfyUI: features/comfyui.md
 - Developer Guide:
   - General:
     - contributing/README.md
+    - pr_reviewer.md
     - glob: contributing/*
       flatten_single_child_sections: true
   - Model Implementation:
     - contributing/model/README.md
+    - contributing/model/adding_omni_model.md
+    - contributing/model/adding_tts_model.md
+    - contributing/model/adding_diffusion_model.md
+    - contributing/model/adding_quantization_model.md
   - CI: contributing/ci
   - Design Documents:
     - design/index.md
     - design/architecture_overview.md
-    - design/vllm_omni_design.md
-    - design/mrs_design.md
+    - Feature Design:
+      - design/feature/disaggregated_inference.md
+      - design/feature/ray_based_execution.md
+      - design/feature/omni_connectors/
+      - design/feature/cfg_parallel.md
+      - design/feature/sequence_parallel.md
+      - design/feature/tensor_parallel.md
+      - design/feature/cache_dit.md
+      - design/feature/teacache.md
+      - design/feature/async_chunk_design.md
+      - design/feature/vae_parallel.md
+      - design/feature/diffusion_step_execution.md
+    - Module Design:
+      - design/module/ar_module.md
+      - design/module/dit_module.md
+      - design/module/entrypoint_module.md
+      - design/module/async_omni_architecture.md
   - Docs Guide: contributing/DOCS_GUIDE.md
 - API Reference:
   - api/README.md
   - api/vllm_omni
 - CLI Reference: cli
 - Community:
+  - Governance: community/governance.md
   - community/*
   - Slack: https://slack.vllm.ai
   - Blog: https://blog.vllm.ai

docs/README.md

@@ -59,4 +59,6 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
 - Multi-modality generation models (e.g. Qwen-Image)
 
 For more information, checkout the following:
+
+- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
 - [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)