add

.buildkite/pipeline-intel.yaml

@@ -10,7 +10,7 @@ steps:
           DOCKER_BUILDKIT: "1"
           # Buildkite will automatically replace this with the actual commit hash
           VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
-          VLLM_VERSION: "v0.19.0"
+          VLLM_VERSION: "v0.18.0"
         priority: 100
         timeout_in_minutes: 60
         soft_fail: true

.buildkite/pipeline.yml

@@ -1,21 +1,6 @@
-# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
-# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs if nightly-test
-# / main NIGHTLY so upload-nightly is not skipped together with test-ready/test-merge.
-#
-# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite).
-steps:
-  - label: ":github: Resolve skip-ci & upload pipeline"
-    key: upload-ci-pipeline
-    commands:
-      - "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
-    agents:
-      queue: "cpu_queue_premerge"
-
----
 steps:
   - label: ":docker: Build image"
     key: image-build
-    if: __IMAGE_BUILD_IF__
     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 ."
@@ -28,7 +13,7 @@ steps:
   - label: "Upload Ready Pipeline"
     depends_on: image-build
     key: upload-ready-pipeline
-    if: __UPLOAD_READY_IF__
+    if: build.branch != "main" && build.pull_request.labels includes "ready"
     commands:
       - buildkite-agent pipeline upload .buildkite/test-ready.yml
     agents:
@@ -38,7 +23,7 @@ steps:
   - label: "Upload Merge Pipeline"
     depends_on: image-build
     key: upload-merge-pipeline
-    if: __UPLOAD_MERGE_IF__
+    if: build.branch == "main" && build.env("NIGHTLY") != "1"
     commands:
       - buildkite-agent pipeline upload .buildkite/test-merge.yml
     agents:

.buildkite/scripts/generate-and-upload-nightly-index.sh

@@ -19,7 +19,7 @@ has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12)
 if [[ "$has_new_python" -eq 0 ]]; then
     # use new python from docker
     docker pull python:3-slim
-    PYTHON="docker run --rm --user $(id -u):$(id -g) -v $(pwd):/app -w /app python:3-slim python3"
+    PYTHON="docker run --rm -v $(pwd):/app -w /app python:3-slim python3"
 fi
 
 echo "Using python interpreter: $PYTHON"
@@ -36,7 +36,7 @@ mkdir -p "$INDICES_OUTPUT_DIR"
 
 # HACK: we do not need regex module here, but it is required by pre-commit hook
 # To avoid any external dependency, we simply replace it back to the stdlib re module
-sed -i.bak 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py && rm -f .buildkite/scripts/generate-nightly-index.py.bak
+sed -i 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
 
 # Generate indices -- the version is just the commit hash (not omni/{commit})
 # because relative paths are computed between the index and wheel directories,
@@ -73,16 +73,15 @@ echo "Pure version (without variant): $pure_version"
 
 # re-generate and copy to /omni/{version}/ only if it does not have "dev" in the version
 if [[ "$version" != *"dev"* ]]; then
-    s3_version="v$pure_version"
+    echo "Re-generating indices for /omni/$pure_version/"
-    echo "Re-generating indices for /omni/$s3_version/"
     rm -rf "${INDICES_OUTPUT_DIR:?}"
     mkdir -p "$INDICES_OUTPUT_DIR"
     # wheel-dir is overridden to be the commit directory, so that the indices point to the correct wheel path
     $PYTHON .buildkite/scripts/generate-nightly-index.py \
-        --version "$s3_version" \
+        --version "$pure_version" \
         --wheel-dir "$BUILDKITE_COMMIT" \
         --current-objects "$obj_json" \
         --output-dir "$INDICES_OUTPUT_DIR" \
         --comment "version $pure_version"
-    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$s3_version/"
+    aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
 fi

.buildkite/scripts/generate-nightly-index.py

@@ -4,7 +4,6 @@
 
 import argparse
 import json
-import re
 import sys
 from dataclasses import asdict, dataclass
 from datetime import datetime
@@ -12,6 +11,8 @@ from pathlib import Path
 from typing import Any
 from urllib.parse import quote
 
+import regex as re
+
 
 def normalize_package_name(name: str) -> str:
     """Normalize package name per PEP 503."""

.buildkite/scripts/upload_pipeline_with_skip_ci.sh

@@ -1,137 +0,0 @@
-#!/usr/bin/env bash
-# Evaluate docs-only skip-ci and upload continuation steps from the same `.buildkite/pipeline.yml`
-# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
-set -euo pipefail
-
-ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
-PIPELINE_YML="${ROOT}/.buildkite/pipeline.yml"
-
-# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
-is_docs_only_change() {
-  local file_path
-  local has_any=0
-
-  while IFS= read -r file_path; do
-    [[ -z "${file_path}" ]] && continue
-    has_any=1
-
-    if [[ "${file_path}" == docs/* ]]; then
-      continue
-    fi
-    if [[ "${file_path}" == *.md ]]; then
-      continue
-    fi
-    if [[ "${file_path}" == "mkdocs.yaml" ]]; then
-      continue
-    fi
-    return 1
-  done
-
-  [[ "${has_any}" -eq 1 ]]
-}
-
-resolve_skip_ci() {
-  local is_pr_build=0
-  local files
-  local base_branch base_ref
-
-  if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
-    is_pr_build=1
-  fi
-
-  if [[ "${is_pr_build}" -eq 1 ]]; then
-    base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
-    if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-      echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
-      git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
-    fi
-
-    base_ref=""
-    if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
-      base_ref="origin/${base_branch}"
-    elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
-      base_ref="${base_branch}"
-    else
-      echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-
-    if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-      echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-  elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
-    if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
-      echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-    if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
-      echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
-      echo -n 0
-      return 0
-    fi
-  else
-    echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
-    echo -n 0
-    return 0
-  fi
-
-  if is_docs_only_change <<< "${files}"; then
-    echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
-    echo -n 1
-    return 0
-  fi
-
-  echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
-  echo -n 0
-}
-
-SKIP_CI="$(resolve_skip_ci)"
-
-if [[ ! -f "${PIPELINE_YML}" ]]; then
-  echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
-  exit 1
-fi
-
-export ROOT SKIP_CI PIPELINE_YML
-python3 <<'PY' | buildkite-agent pipeline upload
-import os
-import pathlib
-
-path = pathlib.Path(os.environ["PIPELINE_YML"])
-text = path.read_text(encoding="utf-8")
-sep = "\n---\n"
-if sep not in text:
-    raise SystemExit(
-        "upload_pipeline_with_skip_ci: .buildkite/pipeline.yml must contain a '\\n---\\n' separator "
-        "(document 1 = bootstrap, document 2 = uploaded steps)"
-    )
-_, continuation = text.split(sep, 1)
-
-skip = os.environ.get("SKIP_CI") == "1"
-# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
-# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
-nightly_only = (
-    '(build.pull_request.labels includes "nightly-test") '
-    '|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
-)
-# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
-if skip:
-    rep = f"'{nightly_only}'"
-    ready_rep = "'false'"
-    merge_rep = "'false'"
-else:
-    rep = "'true'"
-    ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
-    merge_rep = "'build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\"'"
-rendered = (
-    continuation
-    .replace("__IMAGE_BUILD_IF__", rep)
-    .replace("__UPLOAD_READY_IF__", ready_rep)
-    .replace("__UPLOAD_MERGE_IF__", merge_rep)
-)
-print(rendered, end="")
-PY

.buildkite/test-merge.yml

@@ -56,8 +56,8 @@ steps:
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
-      # Single pytest session for one combined summary at end of log.
+      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
-      - pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py tests/e2e/offline_inference/test_diffusion_layerwise_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:
@@ -111,7 +111,8 @@ steps:
     timeout_in_minutes: 20
     depends_on: upload-merge-pipeline
     commands:
-      - pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py tests/diffusion/distributed/test_ulysses_uaa_perf.py
+      - 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:
@@ -173,7 +174,7 @@ steps:
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
         '
     agents:
-      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      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
@@ -192,7 +193,8 @@ steps:
     commands:
       - export VLLM_LOGGING_LEVEL=DEBUG
       - export VLLM_WORKER_MULTIPROC_METHOD=spawn
-      - pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
+      - 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:
@@ -214,7 +216,7 @@ steps:
           export VLLM_LOGGING_LEVEL=DEBUG
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
         '
     agents:
       queue: "gpu_1_queue"
@@ -237,7 +239,7 @@ steps:
           export VLLM_LOGGING_LEVEL=DEBUG
           export VLLM_WORKER_MULTIPROC_METHOD=spawn
           export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
-          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
+          pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
         '
     agents:
       queue: "gpu_1_queue"
@@ -257,8 +259,9 @@ steps:
     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/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
+      - 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:
@@ -344,7 +347,8 @@ steps:
           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 tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
+          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"
@@ -388,7 +392,8 @@ steps:
         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 tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
+          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"
@@ -408,11 +413,6 @@ steps:
                 env:
                   - name: HF_HOME
                     value: /root/.cache/huggingface
-                  - name: HF_TOKEN
-                    valueFrom:
-                      secretKeyRef:
-                        name: hf-token-secret
-                        key: token
             nodeSelector:
               node.kubernetes.io/instance-type: gpu-h100-sxm
             volumes:

.buildkite/test-nightly-diffusion.yml

@@ -1,367 +0,0 @@
-# Nightly diffusion GPU tests — appended to the main nightly build via
-#   buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
-# from test-nightly.yml (step key: nightly-diffusion-model-test). Top-level groups are
-# foldable in the Buildkite UI (Other / Wan / Qwen-Image).
-steps:
-  - group: ":card_index_dividers: Other Model Test"
-    key: nightly-other-model-test-group
-    steps:
-      - label: ":full_moon: Diffusion · Other · Function Test with H100"
-        timeout_in_minutes: 120
-        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 and not test_wan_2_1_vace_expansion and not test_qwen_image" -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 · Other · Function Test with L4"
-        timeout_in_minutes: 60
-        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: Diffusion · Other · Doc Test"
-        timeout_in_minutes: 60
-        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
-
-  - group: ":card_index_dividers: Wan Series Model Test"
-    key: nightly-wan-model-test-group
-    steps:
-      - label: ":full_moon: Diffusion · Wan · Function Test"
-        timeout_in_minutes: 90
-        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 tests/e2e/online_serving/test_wan_2_1_vace_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 · Wan · Accuracy Test"
-        key: nightly-wan22-i2v-accuracy
-        timeout_in_minutes: 180
-        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/wan22_i2v/test_wan22_i2v_video_similarity.py --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
-
-  - group: ":card_index_dividers: Qwen-Image Series Model Test"
-    key: nightly-qwen-image-edit-group
-    steps:
-      - label: ":full_moon: Diffusion · Qwen-Image · Function Test with H100"
-        timeout_in_minutes: 120
-        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_qwen_image*_expansion.py -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 · Qwen-Image · GEBench Accuracy Test"
-        key: nightly-gebench-accuracy
-        timeout_in_minutes: 60
-        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: Diffusion · Qwen-Image · GEdit-Bench Accuracy Test"
-        key: nightly-gedit-bench-accuracy
-        timeout_in_minutes: 60
-        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: Diffusion · Qwen-Image · Perf Test"
-        key: nightly-qwen-image-performance
-        timeout_in_minutes: 180
-        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

.buildkite/test-nightly.yml

@@ -1,15 +1,21 @@
 steps:
-  # Group: collapses under one heading in the Buildkite UI; child steps still run in parallel.
+  - label: ":full_moon: Omni Model Test with H100"
-  - group: ":card_index_dividers: Omni Model Test"
-    key: nightly-omni-test-group
-    steps:
-      - label: ":full_moon: Omni · Function 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_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model"
+      - |
+        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:
@@ -44,13 +50,14 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-      - label: ":full_moon: Omni · Function Test with L4"
+  - label: ":full_moon: Omni Model Test with L4"
     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
       - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+      - pytest -s -v tests/examples/ -m "advanced_model and L4 and omni" --run-level "advanced_model"
       - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
     agents:
       queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
@@ -66,14 +73,95 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-      - label: ":full_moon: Omni · Doc Test with L4"
+  - 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
-          - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
+      - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
-          - pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --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:
@@ -88,13 +176,15 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-      - label: ":full_moon: Omni · Doc Test with H100"
+
-        timeout_in_minutes: 90
+  - 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
-          - pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model"
+      - 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:
@@ -129,7 +219,7 @@ steps:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
 
-      - label: ":full_moon: Omni · Perf Test"
+  - label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
     key: nightly-omni-performance
     timeout_in_minutes: 180
     depends_on: upload-nightly-pipeline
@@ -140,60 +230,6 @@ steps:
       - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
       - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py
       - buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
-          - 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
-
-  # Dynamically appends steps from test-nightly-diffusion.yml into this build (same mechanism as
-  # pipeline.yml → test-ready.yml / test-merge.yml / test-nightly.yml). Foldable groups stay in the
-  # uploaded YAML (Other / Wan / Qwen-Image).
-  - label: ":card_index_dividers: Diffusion Model Test"
-    key: nightly-diffusion-model-test
-    depends_on: upload-nightly-pipeline
-    if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
-    commands:
-      - buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
-    agents:
-      queue: "cpu_queue_premerge"
-
-  - label: ":bar_chart: Testcase Statistics"
-    key: nightly-testcase-statistics
-    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
       - 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:
@@ -230,13 +266,189 @@ steps:
                   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: Wan22 I2V Accuracy Test with H100"
+    key: nightly-wan22-i2v-accuracy
+    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
+      - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --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:c392ce21e9cf9ea65c52b866447793db10e0261c
+                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: 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
-      - nightly-testcase-statistics
     if: build.env("NIGHTLY") == "1"
     commands:
       - pip install openpyxl

.buildkite/test-ready.yml

@@ -16,10 +16,11 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "CUDA Unit Test with single card"
+  - label: "Voxtral TTS CUDA Unit Test"
+    timeout_in_minutes: 10
     depends_on: upload-ready-pipeline
     commands:
-      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and not distributed_cuda' --ignore=tests/e2e --ignore=tests/engine/test_async_omni_engine_abort.py --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
+      - "timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py"
     agents:
       queue: "gpu_1_queue"
     plugins:
@@ -32,22 +33,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "CUDA Unit Test with multi cards"
-    depends_on: upload-ready-pipeline
-    commands:
-      - timeout 10m pytest -v -s -m 'core_model and cuda and L4 and distributed_cuda' --ignore=tests/e2e --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml
-    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
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
   - label: "Diffusion Model Test"
     depends_on: upload-ready-pipeline
     commands:
@@ -167,6 +152,24 @@ steps:
           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
@@ -177,7 +180,7 @@ steps:
                 pytest -s -v tests/engine/test_async_omni_engine_abort.py
         '
     agents:
-      queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
+      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
@@ -268,7 +271,7 @@ steps:
               - 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
+                    nvidia.com/gpu: 2
                 volumeMounts:
                   - name: devshm
                     mountPath: /dev/shm
@@ -317,29 +320,6 @@ steps:
           volumes:
             - "/fsx/hf_cache:/fsx/hf_cache"
 
-  - label: "OmniVoice 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_omnivoice.py -m "core_model" --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
-          shm-size: "8gb"
-          environment:
-            - "HF_HOME=/fsx/hf_cache"
-          volumes:
-            - "/fsx/hf_cache:/fsx/hf_cache"
-
   - label: "Voxtral-TTS E2E Test"
     timeout_in_minutes: 20
     depends_on: upload-ready-pipeline
@@ -348,7 +328,8 @@ steps:
         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 "core_model" --run-level "core_model"
+          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"
@@ -545,46 +526,3 @@ steps:
                 hostPath:
                   path: /mnt/hf-cache
                   type: DirectoryOrCreate
-
-  - label: "CosyVoice3-TTS E2E Test"
-    timeout_in_minutes: 20
-    depends_on: upload-ready-pipeline
-    commands:
-      - |
-        timeout 20m bash -c '
-          export VLLM_WORKER_MULTIPROC_METHOD=spawn
-          pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.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

.claude/skills/add-diffusion-model/SKILL.md

@@ -1,534 +0,0 @@
----
-name: add-diffusion-model
-description: Add a new diffusion model (text-to-image, text-to-video, image-to-video, text-to-audio, image editing) to vLLM-Omni, including Cache-DiT acceleration and parallelism support (TP, SP/USP, CFG-Parallel, HSDP). Use when integrating a new diffusion model, porting a diffusers pipeline or a custom model repo to vllm-omni, creating a new DiT transformer adapter, adding diffusion model support, or enabling multi-GPU parallelism and cache acceleration for an existing model.
----
-
-# Adding a Diffusion Model to vLLM-Omni
-
-## Overview
-
-This skill guides you through adding a new diffusion model to vLLM-Omni. The model may come from HuggingFace Diffusers (structured pipeline) or from a private/custom repo. The workflow differs significantly depending on the source.
-
-## Prerequisites
-
-Before starting, determine:
-
-1. **Model category**: Text-to-Image, Text-to-Video, Image-to-Video, Image Editing, Text-to-Audio, or Omni
-2. **Reference source**: Diffusers pipeline, custom repo, or a combination
-3. **Model HuggingFace ID** or local checkpoint path
-4. **Architecture**: Scheduler, text encoder, VAE, transformer/backbone
-
-## Step 0: Classify the Migration Path
-
-Check the model's HF repo for `model_index.json`. This determines your path:
-
-| Scenario | How to identify | Migration path |
-|----------|----------------|----------------|
-| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
-| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
-| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
-| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
-
-## Path A: Diffusers-Based Model
-
-For models with a standard diffusers layout. See [references/transformer-adaptation.md](references/transformer-adaptation.md) for detailed code patterns.
-
-### A1. Analyze `model_index.json`
-
-Identify components: `transformer`, `scheduler`, `vae`, `text_encoder`, `tokenizer`.
-
-### A2. Create model directory
-
-```
-vllm_omni/diffusion/models/your_model_name/
-├── __init__.py
-├── pipeline_your_model.py
-└── your_model_transformer.py
-```
-
-### A3. Adapt transformer
-
-1. Copy from diffusers source. Remove mixins (`ModelMixin`, `ConfigMixin`, `AttentionModuleMixin`).
-2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
-3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
-4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
-5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
-
-### A4. Adapt pipeline
-
-Inherit from `nn.Module`. The key contract:
-
-```python
-class YourPipeline(nn.Module):
-    def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""):
-        # Load VAE, text encoder, tokenizer via from_pretrained()
-        # Instantiate transformer (weights loaded later via weights_sources)
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model, subfolder="transformer",
-                prefix="transformer.", fall_back_to_pt=True)]
-
-    def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
-        # Encode prompt → prepare latents → denoise loop → VAE decode
-        return DiffusionOutput(output=output)
-
-    def load_weights(self, weights):
-        return AutoWeightsLoader(self).load_weights(weights)
-```
-
-Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
-
-### A5. Register, test, docs → continue at Step 4 below.
-
----
-
-## Path B: Custom/Private Repo Model
-
-For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
-
-### B1. Understand the reference repo
-
-Study the original model's code to identify:
-- **Model architecture files** (transformers, fusion modules, embeddings)
-- **Weight format** (safetensors, `.pth`, custom checkpoint structure)
-- **Weight loading helpers** (custom init functions, checkpoint loaders)
-- **Pre/post-processing** (image/audio transforms, tokenization, VAE encode/decode)
-- **External dependencies** (packages not on PyPI)
-- **Config format** (JSON config files, hardcoded dicts)
-
-### B2. Decide what lives WHERE
-
-This is the key design decision for custom models. Follow these placement rules:
-
-| Code type | Where to place | Example |
-|-----------|---------------|---------|
-| **Pipeline orchestration** (init, forward, denoise loop) | `vllm_omni/diffusion/models/<name>/pipeline_<name>.py` | Always required |
-| **Custom transformer/backbone** (ported and adapted to vllm-omni) | `vllm_omni/diffusion/models/<name>/<name>_transformer.py` or similar | `wan2_2.py`, `fusion.py`, `bagel_transformer.py` |
-| **Custom sub-models** (VAE, fusion, autoencoder) | `vllm_omni/diffusion/models/<name>/` as separate files | `autoencoder.py`, `fusion.py` |
-| **External dependency code** (original repo utilities) | **External repo**, installed via download script or pip | `dreamid_omni` package via git clone |
-| **Hardcoded model configs** | Module-level dicts in pipeline file | `VIDEO_CONFIG`, `AUDIO_CONFIG` dicts |
-| **Download/setup script** | `examples/offline_inference/<name>/download_<name>.py` | `download_dreamid_omni.py` |
-| **Custom `model_index.json`** | Generated by download script, placed at model root | Minimal: `{"_class_name": "YourPipeline", ...}` |
-
-### B3. Handle external dependencies
-
-If the model's code lives in a separate git repo:
-
-**Option 1: Import with graceful fallback** (recommended for models with external utils)
-
-```python
-try:
-    from external_model.utils import init_vae, load_checkpoint
-except ImportError:
-    raise ImportError(
-        "Failed to import from dependency 'external_model'. "
-        "Please run the download script first."
-    )
-```
-
-**Option 2: Port the code directly** (preferred when feasible)
-
-Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and adapt them. This avoids external dependencies. BAGEL does this — `autoencoder.py` and `bagel_transformer.py` are ported directly.
-
-**Decision criteria**: Port if the code is self-contained and won't diverge. Use external deps if the model repo is actively maintained and the code is complex.
-
-### B4. Handle custom weight loading
-
-Custom models have two patterns for weight loading:
-
-**Pattern 1: Bypass standard loader** (DreamID-Omni style)
-
-When the original model has complex custom init functions that load weights in `__init__`:
-
-```python
-class CustomPipeline(nn.Module):
-    def __init__(self, *, od_config, prefix=""):
-        super().__init__()
-        model = od_config.model
-        # Load everything eagerly in __init__ using custom helpers
-        self.vae = custom_init_vae(model, device=self.device)
-        self.text_encoder = custom_init_text_encoder(model, device=self.device)
-        self.transformer = CustomFusionModel(CONFIG)
-        load_custom_checkpoint(self.transformer,
-            checkpoint_path=os.path.join(model, "model.safetensors"))
-        # NO weights_sources defined — bypasses standard loader
-
-    def load_weights(self, weights):
-        pass  # No-op — all weights loaded in __init__
-```
-
-**Pattern 2: Use standard loader with custom `load_weights`** (BAGEL style)
-
-When weights are in safetensors format but need name remapping:
-
-```python
-class CustomPipeline(nn.Module):
-    def __init__(self, *, od_config, prefix=""):
-        super().__init__()
-        # Instantiate model architecture without weights
-        self.bagel = BagelModel(config)
-        self.vae = AutoEncoder(ae_params)
-
-        # Point loader at the safetensors in the model root
-        self.weights_sources = [
-            DiffusersPipelineLoader.ComponentSource(
-                model_or_path=od_config.model,
-                subfolder=None,  # weights at root, not in subfolder
-                prefix="",
-                fall_back_to_pt=False,
-            )
-        ]
-
-    def load_weights(self, weights):
-        # Custom name remapping for non-diffusers weight names
-        params = dict(self.named_parameters())
-        loaded = set()
-        for name, tensor in weights:
-            # Remap original weight names to vllm-omni module names
-            name = self._remap_weight_name(name)
-            if name in params:
-                default_weight_loader(params[name], tensor)
-                loaded.add(name)
-        return loaded
-```
-
-### B5. Create the `model_index.json`
-
-Custom models need a `model_index.json` at the model root for vllm-omni to discover them. For custom models, this is minimal:
-
-```json
-{
-    "_class_name": "YourModelPipeline",
-    "custom_key": "path/to/custom_weights.safetensors"
-}
-```
-
-The `_class_name` must match a key in `_DIFFUSION_MODELS` in `registry.py`. Additional keys are model-specific (accessed via `od_config.model_config`).
-
-If the model's weights come from multiple HF repos, write a **download script** that:
-1. Downloads from each repo
-2. Assembles into a single directory
-3. Generates `model_index.json`
-4. Installs any external dependencies (git clone + `.pth` file)
-
-Place at: `examples/offline_inference/<name>/download_<name>.py`
-
-### B6. Handle multi-modal inputs
-
-If the model accepts images, audio, or other multi-modal inputs, implement the protocol classes from `vllm_omni/diffusion/models/interface.py`:
-
-```python
-from vllm_omni.diffusion.models.interface import SupportImageInput, SupportAudioInput
-
-class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput):
-    # Protocol markers — the engine uses these to enable proper input routing
-    pass
-```
-
-Preprocessing for custom models is typically done **inside `forward()`** rather than via registered pre-process functions, since the logic is often tightly coupled to the model.
-
-### B7. Continue at Step 4 below.
-
----
-
-## Common Steps (Both Paths)
-
-### Step 4: Register Model in registry.py
-
-Edit `vllm_omni/diffusion/registry.py`:
-
-```python
-_DIFFUSION_MODELS = {
-    "YourModelPipeline": ("your_model_name", "pipeline_your_model", "YourModelPipeline"),
-}
-_DIFFUSION_POST_PROCESS_FUNCS = {
-    "YourModelPipeline": "get_your_model_post_process_func",  # if applicable
-}
-_DIFFUSION_PRE_PROCESS_FUNCS = {
-    "YourModelPipeline": "get_your_model_pre_process_func",  # if applicable
-}
-```
-
-The registry key is the `_class_name` from `model_index.json`. The tuple is `(folder_name, module_file, class_name)`.
-
-Create `__init__.py` exporting the pipeline class and any factory functions.
-
-### Step 5: Run, Test, Debug
-
-Use the appropriate existing example script:
-
-| Category | Script |
-|----------|--------|
-| Text-to-Image | `examples/offline_inference/text_to_image/text_to_image.py` |
-| Text-to-Video | `examples/offline_inference/text_to_video/text_to_video.py` |
-| Image-to-Video | `examples/offline_inference/image_to_video/image_to_video.py` |
-| Image-to-Image | `examples/offline_inference/image_to_image/image_edit.py` |
-| Text-to-Audio | `examples/offline_inference/text_to_audio/text_to_audio.py` |
-
-For custom/Omni models that don't fit these categories, create a dedicated example script.
-
-**Validation**: No errors, output is meaningful, quality matches reference implementation.
-
-See [references/troubleshooting.md](references/troubleshooting.md) for common errors.
-
-### Step 6: Add Example Scripts
-
-For Omni or custom models, create:
-- `examples/offline_inference/your_model_name/` — offline script + README
-- `examples/online_serving/your_model_name/` — server script + client
-- Download script if weights require assembly from multiple sources
-
-### Step 7: Update Documentation
-
-Required updates:
-1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
-2. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
-3. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
-4. `examples/offline_inference/xxx/README.md` — offline example docs
-5. `examples/online_serve/xxx/README.md` — online serve docs
-
-### Step 8: Add E2E Tests (Recommended)
-
-Create `tests/e2e/online_serving/test_your_model_expansion.py`.
-
-### Step 9: Add Cache-DiT Acceleration
-
-Cache-DiT accelerates inference by caching intermediate computation results across denoising steps. After your model is working correctly on a single GPU, add cache-dit support.
-
-See [references/cache-dit-patterns.md](references/cache-dit-patterns.md) for detailed code patterns.
-
-#### 9a. Determine your model type
-
-| Model Type | Description | Action |
-|------------|-------------|--------|
-| **Standard single-transformer** | One transformer with one `ModuleList` of blocks | No code needed — `CacheDiTBackend` auto-detects via `enable_cache_for_dit()` |
-| **Multi-block-list** | One transformer with multiple block lists (e.g., `transformer_blocks` + `single_transformer_blocks`) | Write custom enabler with `BlockAdapter` |
-| **Dual-transformer** | Two transformers (e.g., high-noise + low-noise) | Write custom enabler with `BlockAdapter` wrapping both |
-
-#### 9b. Standard models — verify automatic support
-
-For standard single-transformer models, test directly:
-
-```python
-omni = Omni(
-    model="your-model-name",
-    cache_backend="cache_dit",
-    cache_config={
-        "Fn_compute_blocks": 1,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4,
-    }
-)
-```
-
-Check logs for "Cache-dit enabled successfully on xxx". If it works, skip to Step 9e.
-
-#### 9c. Custom architectures — write a custom enabler
-
-For multi-block-list or dual-transformer models, write a custom enabler function:
-
-```python
-from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
-
-def enable_cache_for_your_model(pipeline, cache_config):
-    db_cache_config = DBCacheConfig(
-        num_inference_steps=None,
-        Fn_compute_blocks=cache_config.Fn_compute_blocks,
-        Bn_compute_blocks=cache_config.Bn_compute_blocks,
-        max_warmup_steps=cache_config.max_warmup_steps,
-        max_cached_steps=cache_config.max_cached_steps,
-        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
-        residual_diff_threshold=cache_config.residual_diff_threshold,
-    )
-
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=[
-                pipeline.transformer.transformer_blocks,
-                pipeline.transformer.single_transformer_blocks,
-            ],
-            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
-            params_modifiers=[ParamsModifier(...)],
-        ),
-        cache_config=db_cache_config,
-    )
-
-    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
-        cache_dit.refresh_context(
-            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
-        )
-    return refresh_cache_context
-```
-
-#### 9d. Register the custom enabler
-
-Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
-
-```python
-CUSTOM_DIT_ENABLERS = {
-    "Wan22Pipeline": enable_cache_for_wan22,
-    "LongCatImagePipeline": enable_cache_for_longcat_image,
-    "YourModelPipeline": enable_cache_for_your_model,  # Add here
-}
-```
-
-#### 9e. Test Cache-DiT
-
-```python
-omni = Omni(
-    model="your-model-name",
-    cache_backend="cache_dit",
-    cache_config={
-        "Fn_compute_blocks": 1, "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4, "residual_diff_threshold": 0.24,
-    }
-)
-images = omni.generate("a beautiful landscape",
-    OmniDiffusionSamplingParams(num_inference_steps=50))
-```
-
-**Verify**: 1) logs show cache enabled, 2) 1.5-2x speedup, 3) output quality acceptable vs baseline.
-
-If quality degrades, lower `residual_diff_threshold` (try 0.12-0.18) or increase `max_warmup_steps` (try 6-8).
-
----
-
-### Step 10: Add Parallelism Support
-
-After the model works on a single GPU, add multi-GPU parallelism. Add each type incrementally, testing after each addition.
-
-See [references/parallelism-patterns.md](references/parallelism-patterns.md) for detailed code patterns and API reference.
-
-**Recommended order**: TP → SP/USP → CFG Parallel → HSDP
-
-#### 10a. Tensor Parallelism (TP)
-
-Shards DiT linear layers across GPUs. Requires code changes in the transformer.
-
-**What to change in the transformer**:
-1. Replace `nn.Linear` with `ColumnParallelLinear` / `RowParallelLinear` / `QKVParallelLinear`
-2. Update `load_weights()` to handle QKV fusion with `stacked_params_mapping`
-3. Use `self.to_qkv.num_heads` (local heads) instead of total heads for split sizes
-
-```python
-from vllm.model_executor.layers.linear import (
-    QKVParallelLinear, RowParallelLinear, ColumnParallelLinear,
-)
-
-# Attention: QKV → RowParallel output
-self.to_qkv = QKVParallelLinear(dim, head_dim, num_heads, num_kv_heads)
-self.to_out = RowParallelLinear(dim, dim, input_is_parallel=True)
-
-# FFN: ColumnParallel → RowParallel
-self.w1 = ColumnParallelLinear(dim, ffn_dim)
-self.w2 = RowParallelLinear(ffn_dim, dim, input_is_parallel=True)
-```
-
-**Constraints**: `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`.
-
-**Test**: `--tensor-parallel-size 2`
-
-#### 10b. Sequence Parallelism (SP / USP)
-
-Splits sequence tokens across GPUs. Non-intrusive via `_sp_plan` on the transformer class — no changes to `forward()`.
-
-**What to change in the transformer**:
-
-Add `_sp_plan` class attribute:
-
-```python
-from vllm_omni.diffusion.distributed.sp_plan import (
-    SequenceParallelInput, SequenceParallelOutput,
-)
-
-class YourTransformer(nn.Module):
-    _sp_plan = {
-        "blocks.0": {
-            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-        },
-        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-    }
-```
-
-If inline tensor ops (e.g., `torch.cat`) exist between shard/gather points, extract them into `nn.Module` submodules so hooks can intercept them.
-
-For RoPE that needs splitting, add an entry for the RoPE module with `split_output=True`.
-
-**Test**: `--ulysses-degree 2` (offline) or `--usp 2` (online serving)
-
-#### 10c. CFG Parallel
-
-Distributes positive/negative CFG branches across 2 GPUs. Requires the pipeline to inherit `CFGParallelMixin`.
-
-**What to change in the pipeline**:
-
-```python
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-
-class YourPipeline(nn.Module, CFGParallelMixin):
-    def diffuse(self, ...) -> torch.Tensor:
-        for i, t in enumerate(timesteps):
-            positive_kwargs = {...}
-            negative_kwargs = {...} if do_true_cfg else None
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                do_true_cfg=do_true_cfg, true_cfg_scale=cfg_scale,
-                positive_kwargs=positive_kwargs, negative_kwargs=negative_kwargs,
-            )
-            latents = self.scheduler_step_maybe_with_cfg(
-                noise_pred, t, latents, do_true_cfg
-            )
-        return latents
-```
-
-Override `predict_noise()` if your transformer call is non-standard. Override `combine_cfg_noise()` for multi-output models (e.g., video + audio).
-
-**Constraint**: Exactly 2 GPUs. Only for models using classifier-free guidance.
-
-**Test**: `--cfg-parallel-size 2`
-
-#### 10d. HSDP (Hybrid Sharded Data Parallel)
-
-Shards transformer weights via PyTorch FSDP2 to reduce per-GPU VRAM. No code changes to the forward pass — just add a class attribute.
-
-**What to change in the transformer**:
-
-```python
-class YourTransformer(nn.Module):
-    @staticmethod
-    def _is_transformer_block(name: str, module) -> bool:
-        return "blocks" in name and name.split(".")[-1].isdigit()
-
-    _hsdp_shard_conditions = [_is_transformer_block]
-```
-
-**Constraint**: Cannot combine with TP. For standalone HSDP, set `hsdp_shard_size` explicitly.
-
-**Test**: `--use-hsdp` or `DiffusionParallelConfig(use_hsdp=True)`
-
-#### 10e. Update parallelism documentation
-
-After adding parallelism support, update:
-1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
-2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
-
----
-
-## Iterative Development Tips
-
-1. **Start minimal**: Basic generation first, no parallelism/caching
-2. **Use `--enforce-eager`**: Disable torch.compile during debugging
-3. **Use small models**: Test with smaller variants first
-4. **Check tensor shapes**: Most errors are reshape mismatches in attention
-5. **Add features incrementally**: Single GPU → TP → SP → CFG → HSDP → Cache-DiT
-6. **For custom models**: Get the model running with the original code first, then progressively replace components with vllm-omni equivalents
-7. **Cache-DiT before parallelism tuning**: Cache-DiT is lossy — verify quality at baseline before combining with parallelism
-8. **Combine lossless + lossy**: e.g., TP + SP + Cache-DiT for maximum throughput
-
-## Reference Files
-
-- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
-- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
-- [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details
-- [Cache-DiT Patterns](references/cache-dit-patterns.md) — cache-dit acceleration for standard and custom architectures
-- [Troubleshooting](references/troubleshooting.md) — common errors and fixes

.claude/skills/add-diffusion-model/references/cache-dit-patterns.md

@@ -1,254 +0,0 @@
-# Cache-DiT Patterns Reference
-
-## Overview
-
-Cache-DiT accelerates Diffusion Transformers by caching intermediate computation results across denoising steps. Adjacent steps produce similar features, so redundant computations can be skipped.
-
-Three caching strategies:
-- **DBCache**: Dynamic block-level caching — selectively computes or caches transformer blocks based on residual differences
-- **TaylorSeer**: Calibration-based prediction using Taylor expansion to estimate block outputs
-- **SCM** (Step Computation Masking): Dynamic step skipping based on configurable policies
-
-**Typical speedup**: 1.5-2.5x depending on model and configuration.
-
-**Official docs**: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cache_dit
-
-## Architecture
-
-vLLM-Omni integrates cache-dit through `CacheDiTBackend`:
-
-| Component | Purpose |
-|-----------|---------|
-| `CacheDiTBackend` | Unified backend — auto-selects enabler (standard or custom) |
-| `enable_cache_for_dit()` | Default enabler for standard single-transformer models |
-| `CUSTOM_DIT_ENABLERS` dict | Registry of custom enablers keyed by pipeline class name |
-| `BlockAdapter` | Wraps complex architectures (multi-block-list or multi-transformer) |
-| `ForwardPattern` | Specifies block forward signature: `Pattern_0`, `Pattern_1`, `Pattern_2` |
-| `ParamsModifier` | Per-transformer or per-block-list config customization |
-| `DBCacheConfig` | Configuration for DBCache parameters |
-| `cache_dit.refresh_context()` | Updates cache context when `num_inference_steps` changes |
-
-**Source files**:
-- `vllm_omni/diffusion/cache/cache_dit_backend.py` — `CacheDiTBackend`, enablers, `CUSTOM_DIT_ENABLERS`
-- `vllm_omni/diffusion/cache/` — cache backend implementations
-
-## Standard Models: Automatic Support
-
-Most DiT models follow this pattern:
-- Single transformer with one `nn.ModuleList` of blocks
-- Standard forward signature
-- Compatible with cache-dit's automatic detection
-
-**Examples**: Qwen-Image, Z-Image, FLUX
-
-No code changes needed. `CacheDiTBackend` automatically uses `enable_cache_for_dit()`:
-
-```python
-from vllm_omni import Omni
-
-omni = Omni(
-    model="Qwen/Qwen-Image",
-    cache_backend="cache_dit",
-    cache_config={
-        "Fn_compute_blocks": 1,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4,
-    }
-)
-```
-
-What happens automatically:
-
-```python
-def enable_cache_for_dit(pipeline, cache_config):
-    db_cache_config = DBCacheConfig(
-        num_inference_steps=None,
-        Fn_compute_blocks=cache_config.Fn_compute_blocks,
-        Bn_compute_blocks=cache_config.Bn_compute_blocks,
-        max_warmup_steps=cache_config.max_warmup_steps,
-        max_cached_steps=cache_config.max_cached_steps,
-        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
-        residual_diff_threshold=cache_config.residual_diff_threshold,
-    )
-
-    cache_dit.enable_cache(pipeline.transformer, cache_config=db_cache_config)
-
-    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
-        cache_dit.refresh_context(
-            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
-        )
-    return refresh_cache_context
-```
-
-## Custom Architectures: Writing Custom Enablers
-
-### When you need a custom enabler
-
-- Model has multiple block lists in one transformer (e.g., `transformer_blocks` + `single_transformer_blocks`)
-- Model has two transformers (e.g., high-noise + low-noise like Wan2.2)
-- Model uses non-standard block forward signature
-
-### Pattern 1: Multi-Block-List (LongCat-Image style)
-
-Single transformer with two block lists:
-
-```python
-import cache_dit
-from cache_dit import BlockAdapter, ForwardPattern, ParamsModifier, DBCacheConfig
-
-def enable_cache_for_your_model(pipeline, cache_config):
-    db_cache_config = DBCacheConfig(
-        num_inference_steps=None,
-        Fn_compute_blocks=cache_config.Fn_compute_blocks,
-        Bn_compute_blocks=cache_config.Bn_compute_blocks,
-        max_warmup_steps=cache_config.max_warmup_steps,
-        max_cached_steps=cache_config.max_cached_steps,
-        max_continuous_cached_steps=cache_config.max_continuous_cached_steps,
-        residual_diff_threshold=cache_config.residual_diff_threshold,
-    )
-
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=pipeline.transformer,
-            blocks=[
-                pipeline.transformer.transformer_blocks,
-                pipeline.transformer.single_transformer_blocks,
-            ],
-            forward_pattern=[ForwardPattern.Pattern_1, ForwardPattern.Pattern_1],
-            params_modifiers=[ParamsModifier(...)],
-        ),
-        cache_config=db_cache_config,
-    )
-
-    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
-        cache_dit.refresh_context(
-            pipeline.transformer, num_inference_steps=num_inference_steps, verbose=verbose
-        )
-    return refresh_cache_context
-```
-
-For single transformer with multiple block lists, `refresh_context` works the same as standard models — call it once on the transformer.
-
-### Pattern 2: Dual-Transformer (Wan2.2 style)
-
-Two transformers with separate configs:
-
-```python
-def enable_cache_for_dual_transformer(pipeline, cache_config):
-    db_cache_config = DBCacheConfig(...)
-
-    cache_dit.enable_cache(
-        BlockAdapter(
-            transformer=[pipeline.transformer, pipeline.transformer_2],
-            blocks=[pipeline.transformer.blocks, pipeline.transformer_2.blocks],
-            forward_pattern=[ForwardPattern.Pattern_2, ForwardPattern.Pattern_2],
-            params_modifiers=[
-                ParamsModifier(...),  # Config for transformer 1
-                ParamsModifier(...),  # Config for transformer 2
-            ],
-        ),
-        cache_config=db_cache_config,
-    )
-
-    def refresh_cache_context(pipeline, num_inference_steps, verbose=True):
-        high_steps, low_steps = _split_inference_steps(num_inference_steps)
-        cache_dit.refresh_context(
-            pipeline.transformer, num_inference_steps=high_steps, verbose=verbose
-        )
-        cache_dit.refresh_context(
-            pipeline.transformer_2, num_inference_steps=low_steps, verbose=verbose
-        )
-    return refresh_cache_context
-```
-
-Key difference: `refresh_context` must be called on **each transformer separately** with its own step count.
-
-### Choosing the ForwardPattern
-
-| Pattern | Block forward signature | Example models |
-|---------|------------------------|----------------|
-| `Pattern_0` | `block(hidden_states, **kwargs)` → residual added inside block | Default |
-| `Pattern_1` | `block(hidden_states, **kwargs)` → returns `(hidden_states, ...)` tuple | FLUX-style single blocks |
-| `Pattern_2` | `block(hidden_states, **kwargs)` → `(hidden_states, ...)` with different residual pattern | Wan2.2 blocks |
-
-Inspect your block's `forward()` return type and residual connection pattern to choose the right one. See [Cache-DiT API Reference](https://cache-dit.readthedocs.io/en/latest/user_guide/CACHE_API/) for details.
-
-## Registering Custom Enablers
-
-Add your enabler to `CUSTOM_DIT_ENABLERS` in `vllm_omni/diffusion/cache/cache_dit_backend.py`:
-
-```python
-CUSTOM_DIT_ENABLERS = {
-    "Wan22Pipeline": enable_cache_for_wan22,
-    "LongCatImagePipeline": enable_cache_for_longcat_image,
-    "YourModelPipeline": enable_cache_for_your_model,
-}
-```
-
-The key must match `pipeline.__class__.__name__`.
-
-## Configuration Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `Fn_compute_blocks` | 1 | Number of blocks to always compute at the front |
-| `Bn_compute_blocks` | 0 | Number of blocks to always compute at the back |
-| `max_warmup_steps` | 4 | Steps to run without caching at the beginning |
-| `max_cached_steps` | — | Max total cached steps |
-| `max_continuous_cached_steps` | — | Max consecutive cached steps |
-| `residual_diff_threshold` | 0.24 | Threshold for deciding whether to cache a block |
-
-### Tuning for quality vs speed
-
-| Goal | Adjustments |
-|------|-------------|
-| **More speed, acceptable quality loss** | Higher `residual_diff_threshold` (0.24-0.4), lower `max_warmup_steps` (2-4) |
-| **Better quality, less speed** | Lower `residual_diff_threshold` (0.12-0.18), higher `max_warmup_steps` (6-8), lower `max_continuous_cached_steps` (2) |
-
-## Testing
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(
-    model="your-model-name",
-    cache_backend="cache_dit",
-    cache_config={
-        "Fn_compute_blocks": 1,
-        "Bn_compute_blocks": 0,
-        "max_warmup_steps": 4,
-        "residual_diff_threshold": 0.24,
-    }
-)
-images = omni.generate(
-    "a beautiful landscape",
-    OmniDiffusionSamplingParams(num_inference_steps=50),
-)
-```
-
-CLI (online serving):
-
-```bash
-vllm serve your-model --omni --port 8098 \
-  --cache-backend cache_dit \
-  --cache-config '{"Fn_compute_blocks": 1, "Bn_compute_blocks": 0, "max_warmup_steps": 4}'
-```
-
-**Verification checklist**:
-1. Logs show "Cache-dit enabled successfully on xxx"
-2. Performance: 1.5-2x speedup vs no cache
-3. Quality: compare output with `cache_backend=None`
-
-## Excluded Models
-
-Models listed in `_NO_CACHE_ACCELERATION` in `vllm_omni/diffusion/registry.py` do not support cache-dit (e.g., `NextStep11Pipeline`, `StableDiffusionPipeline`). Check this set before attempting to enable cache-dit.
-
-## Reference Implementations
-
-| Model | Path | Notes |
-|-------|------|-------|
-| Standard DiT | `cache_dit_backend.py::enable_cache_for_dit` | Default enabler, automatic |
-| Wan2.2 | `cache_dit_backend.py::enable_cache_for_wan22` | Dual-transformer, auto-detects mode |
-| LongCat | `cache_dit_backend.py::enable_cache_for_longcat_image` | Multi-block-list |
-| BAGEL | `cache_dit_backend.py::enable_cache_for_bagel` | Complex omni model |

.claude/skills/add-diffusion-model/references/custom-model-patterns.md

@@ -1,273 +0,0 @@
-# Custom Model Patterns Reference
-
-Patterns for adding models that don't come from the standard diffusers pipeline format.
-
-## Directory Structure Comparison
-
-### Diffusers-based model (e.g., Wan2.2)
-
-```
-vllm_omni/diffusion/models/wan2_2/
-├── __init__.py                    # Exports pipeline + transformer + helpers
-├── pipeline_wan2_2.py             # Pipeline: loads components via from_pretrained()
-├── pipeline_wan2_2_i2v.py         # Variant pipeline for image-to-video
-└── wan2_2_transformer.py          # Transformer: ported from diffusers, uses Attention layer
-```
-
-The transformer is loaded separately via `weights_sources` + `load_weights()`. Non-transformer components (VAE, text encoder) are loaded in `__init__` via `from_pretrained()`.
-
-### Custom model with external deps (e.g., DreamID-Omni)
-
-```
-vllm_omni/diffusion/models/dreamid_omni/
-├── __init__.py                    # Exports pipeline only
-├── pipeline_dreamid_omni.py       # Pipeline: loads ALL weights in __init__ via custom helpers
-├── fusion.py                      # Custom fusion architecture (video + audio cross-attention)
-└── wan2_2.py                      # Re-implemented Wan backbone with split API
-
-examples/offline_inference/x_to_video_audio/
-└── download_dreamid_omni.py       # Downloads weights from 3 HF repos + clones code repo
-```
-
-All weights loaded eagerly in `__init__`. `load_weights()` is a no-op. External dependency (`dreamid_omni` package) imported with try/except.
-
-### Custom model with ported code (e.g., BAGEL)
-
-```
-vllm_omni/diffusion/models/bagel/
-├── __init__.py
-├── pipeline_bagel.py              # Pipeline: instantiates models, uses weights_sources
-├── bagel_transformer.py           # Full LLM backbone (Qwen2-MoT) ported into vllm-omni
-└── autoencoder.py                 # Custom VAE ported from original repo
-```
-
-Model code is fully ported (no external dependency). Uses `weights_sources` and `load_weights()` with custom name remapping to handle non-diffusers safetensors format.
-
-## Weight Loading Patterns
-
-### Pattern 1: Standard diffusers flow (Wan2.2, Z-Image, FLUX)
-
-```
-init → create transformer (empty) → set weights_sources → [loader calls load_weights()]
-```
-
-- `weights_sources` points to safetensors in HF subfolder (e.g., `transformer/`)
-- `load_weights()` receives `(name, tensor)` pairs from the loader
-- Name remapping handles diffusers→vllm-omni differences (QKV fusion, Sequential index removal)
-
-### Pattern 2: Custom safetensors at root (BAGEL)
-
-```
-init → create all models (empty) → set weights_sources(subfolder=None) → [loader calls load_weights()]
-```
-
-- `weights_sources` points to **root** of model directory, not a subfolder
-- Weights have non-diffusers names (e.g., `bagel.language_model.model.layers.0.self_attn.q_proj.weight`)
-- `load_weights()` does heavy name normalization
-
-```python
-self.weights_sources = [
-    DiffusersPipelineLoader.ComponentSource(
-        model_or_path=od_config.model,
-        subfolder=None,      # root directory
-        prefix="",           # no prefix stripping
-        fall_back_to_pt=False,
-    )
-]
-```
-
-### Pattern 3: Fully custom loading (DreamID-Omni)
-
-```
-init → load ALL weights eagerly via custom helpers → load_weights() = no-op
-```
-
-- No `weights_sources` attribute — standard loader finds nothing to iterate
-- Custom init functions (e.g., `init_wan_vae_2_2()`, `load_fusion_checkpoint()`) handle downloading and loading
-- `load_weights()` is `pass`
-- Weights may come from multiple HF repos in different formats (`.pth`, `.safetensors`)
-
-Use this when:
-- The original model has complex, well-tested loading code you don't want to rewrite
-- Weights span multiple HF repos
-- Weight format is non-standard (e.g., a single `.pth` file, not sharded safetensors)
-
-## model_index.json for Custom Models
-
-Standard diffusers `model_index.json`:
-```json
-{
-    "_class_name": "WanPipeline",
-    "_diffusers_version": "0.35.0.dev0",
-    "scheduler": ["diffusers", "UniPCMultistepScheduler"],
-    "transformer": ["diffusers", "WanTransformer3DModel"],
-    "vae": ["diffusers", "AutoencoderKLWan"]
-}
-```
-
-Custom model `model_index.json` (minimal):
-```json
-{
-    "_class_name": "DreamIDOmniPipeline",
-    "fusion": "DreamID-Omni/dreamid_omni.safetensors"
-}
-```
-
-The only **required** field is `_class_name` — it must match a key in `_DIFFUSION_MODELS` in `registry.py`. Other fields are model-specific and accessible via `od_config.model_config` dict.
-
-## External Dependency Management
-
-### Git clone + .pth injection (DreamID-Omni pattern)
-
-```python
-def download_dependency():
-    CACHE_DIR.mkdir(parents=True, exist_ok=True)
-    with open(LOCK_FILE, "w") as f:
-        fcntl.flock(f, fcntl.LOCK_EX)
-        if not DEPENDENCY_DIR.exists():
-            subprocess.run([
-                "git", "clone", "--depth", "1",
-                REPO_URL, "--branch", BRANCH,
-                str(DEPENDENCY_DIR)
-            ], check=True)
-        fcntl.flock(f, fcntl.LOCK_UN)
-
-    # Add to Python path via .pth file
-    site_packages = Path(site.getsitepackages()[0])
-    pth_file = site_packages / "vllm_omni_dependency.pth"
-    pth_file.write_text(str(DEPENDENCY_DIR))
-```
-
-### Direct port (BAGEL pattern)
-
-Copy essential files from the original repo into `vllm_omni/diffusion/models/<name>/`. Adapt imports to use vllm-omni utilities. Benefits: no external dependency, no git clone step. Drawback: must maintain the ported code.
-
-## Multi-Modal Input/Output Protocols
-
-Custom models that handle images, audio, or video I/O should implement protocol classes:
-
-```python
-from vllm_omni.diffusion.models.interface import (
-    SupportImageInput,    # Model accepts image input
-    SupportAudioInput,    # Model accepts audio input
-    SupportAudioOutput,   # Model produces audio output
-)
-
-class MyPipeline(nn.Module, SupportImageInput, SupportAudioInput, SupportAudioOutput):
-    pass  # Protocol markers enable proper engine routing
-```
-
-The engine checks `isinstance(pipeline, SupportImageInput)` at startup to configure input validation and warmup behavior.
-
-## Hardcoded Config vs Config Files
-
-Diffusers models use `config.json` in each subfolder. Custom models often use:
-
-**Module-level config dicts** (DreamID-Omni):
-```python
-VIDEO_CONFIG = {
-    "patch_size": [1, 2, 2], "model_type": "ti2v",
-    "dim": 3072, "ffn_dim": 14336, "num_heads": 24, "num_layers": 30, ...
-}
-```
-
-**Loaded from custom JSON** (BAGEL):
-```python
-cfg_path = os.path.join(model_path, "config.json")
-with open(cfg_path) as f:
-    bagel_cfg = json.load(f)
-vae_cfg = bagel_cfg.get("vae_config", {})
-```
-
-## Custom Architecture Patterns
-
-### Split forward API (DreamID-Omni)
-
-When a fusion model needs to interleave blocks from two backbones:
-
-```python
-class WanModel(nn.Module):
-    def prepare_transformer_block_kwargs(self, x, t, context, ...):
-        # Patch embed, time embed, text embed, RoPE
-        return x, e, kwargs
-
-    def post_transformer_block_out(self, x, grid_sizes, e):
-        # Output projection, unpatchify
-        return output
-
-    def forward(self, *args, **kwargs):
-        raise NotImplementedError  # Fusion model handles block iteration
-```
-
-The `FusionModel` then iterates blocks in lock-step:
-```python
-for video_block, audio_block in zip(self.video_model.blocks, self.audio_model.blocks):
-    video_out = video_block(video_hidden, ...)
-    audio_out = audio_block(audio_hidden, ...)
-    # Cross-attend between modalities
-    video_out = cross_attention(video_out, audio_out)
-    audio_out = cross_attention(audio_out, video_out)
-```
-
-### LLM-as-denoiser (BAGEL)
-
-When the backbone is a language model that also does diffusion:
-
-```python
-class BagelModel(nn.Module):
-    def __init__(self):
-        self.language_model = Qwen2MoTForCausalLM(config)
-        self.vit_model = SiglipVisionModel(vit_config)
-```
-
-The LLM processes both text tokens and latent image tokens in a single forward pass, using KV caching for the text portion.
-
-## Pre/Post Processing for Custom Models
-
-Custom models typically handle pre/post processing **inside `forward()`** rather than via registered functions, because the logic is tightly coupled:
-
-```python
-def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
-    # Inline preprocessing
-    image = self._load_and_resize_image(req.prompts[0].get("multi_modal_data", {}).get("image"))
-    image_latent = self._vae_encode(image)
-
-    # ... denoising loop ...
-
-    # Inline postprocessing
-    pil_image = self._decode_to_pil(latents)
-    return DiffusionOutput(output=[pil_image])
-```
-
-If pre/post functions are not registered in `_DIFFUSION_PRE_PROCESS_FUNCS` / `_DIFFUSION_POST_PROCESS_FUNCS`, the engine simply skips those steps.
-
-## Download Script Template
-
-```python
-# examples/offline_inference/<name>/download_<name>.py
-from huggingface_hub import snapshot_download
-import json, os
-
-def main(output_dir):
-    # Download model weights from HF
-    snapshot_download(repo_id="org/model-weights", local_dir=os.path.join(output_dir, "weights"))
-
-    # Download additional components if from separate repos
-    snapshot_download(repo_id="org/vae-weights", local_dir=os.path.join(output_dir, "vae"),
-        allow_patterns=["*.safetensors"])
-
-    # Generate model_index.json
-    config = {"_class_name": "YourPipeline", "custom_key": "weights/model.safetensors"}
-    with open(os.path.join(output_dir, "model_index.json"), "w") as f:
-        json.dump(config, f, indent=2)
-
-    # Install external code dependency (if needed)
-    download_dependency()
-
-if __name__ == "__main__":
-    import argparse
-    parser = argparse.ArgumentParser()
-    parser.add_argument("--output-dir", default="./your_model")
-    args = parser.parse_args()
-    main(args.output_dir)
-```

.claude/skills/add-diffusion-model/references/parallelism-patterns.md

@@ -1,571 +0,0 @@
-# Parallelism Patterns Reference
-
-## Overview
-
-vLLM-Omni supports multiple parallelism strategies for diffusion models. Each targets a different bottleneck:
-
-| Strategy | Splits | Best For | Constraint |
-|----------|--------|----------|------------|
-| Tensor Parallel (TP) | Model layers across GPUs | Latency reduction, large models | Requires fast GPU interconnect, `num_heads % tp == 0` |
-| Sequence Parallel (SP/USP) | Sequence tokens across GPUs | Long sequences (video, high-res) | Near-linear scaling |
-| CFG Parallel | Positive/negative CFG branches | Models using classifier-free guidance | Exactly 2 GPUs |
-| HSDP | Weight shards via FSDP2 | VRAM reduction | Cannot combine with TP |
-| VAE Patch Parallel | VAE decode spatial tiles | Large VAE outputs | Auto-enables tiling |
-
-**Recommended integration order**: TP → SP → CFG Parallel → HSDP
-
-**Official design docs**:
-- TP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/tensor_parallel
-- SP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/sequence_parallel
-- CFG: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/cfg_parallel
-- HSDP: https://docs.vllm.ai/projects/vllm-omni/en/latest/design/feature/hsdp
-
----
-
-## Tensor Parallelism (TP)
-
-Replace standard `nn.Linear` with vLLM's parallel linear layers. This is the most invasive change but provides direct VRAM savings and compute speedup.
-
-### Layer replacement rules
-
-| Pattern | vLLM Layer | When to Use |
-|---------|-----------|-------------|
-| Fan-out (first in FFN) | `ColumnParallelLinear` | Projection that splits output across ranks |
-| Fan-in (second in FFN) | `RowParallelLinear` | Projection that gathers across ranks |
-| QKV projection | `QKVParallelLinear` | Fused Q/K/V for self-attention |
-| Single Q or K or V | `ColumnParallelLinear` | Separate projections (cross-attention) |
-| Attention output | `RowParallelLinear` | Output projection after attention |
-| Must not shard | `ReplicatedLinear` | Layers that must stay replicated |
-
-### MLP Block (Up-Down Pattern)
-
-```python
-from vllm.model_executor.layers.linear import (
-    ColumnParallelLinear, RowParallelLinear,
-)
-
-class TPFeedForward(nn.Module):
-    def __init__(self, dim, ffn_dim):
-        super().__init__()
-        self.fc1 = ColumnParallelLinear(dim, ffn_dim, bias=False, return_bias=False)
-        self.fc2 = RowParallelLinear(
-            ffn_dim, dim, bias=False,
-            input_is_parallel=True,  # Input already sharded from fc1
-            return_bias=False,
-        )
-
-    def forward(self, x):
-        x, _ = self.fc1(x)
-        x = torch.nn.functional.gelu(x)
-        x, _ = self.fc2(x)
-        return x
-```
-
-### Attention Block (QKV-Out Pattern)
-
-```python
-from vllm.model_executor.layers.linear import QKVParallelLinear, RowParallelLinear
-from vllm_omni.diffusion.attention.layer import Attention
-
-class TPSelfAttention(nn.Module):
-    def __init__(self, dim, num_heads, num_kv_heads=None):
-        super().__init__()
-        num_kv_heads = num_kv_heads or num_heads
-        self.head_dim = dim // num_heads
-
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=self.head_dim,
-            total_num_heads=num_heads,
-            total_num_kv_heads=num_kv_heads,
-            bias=False,
-            return_bias=False,
-        )
-        self.to_out = RowParallelLinear(
-            dim, dim, bias=False,
-            input_is_parallel=True,
-            return_bias=False,
-        )
-        self.attn = Attention(
-            num_heads=self.to_qkv.num_heads,      # Local heads per GPU
-            head_size=self.head_dim,
-            softmax_scale=1.0 / (self.head_dim ** 0.5),
-            causal=False,
-            num_kv_heads=self.to_qkv.num_kv_heads,  # Local KV heads per GPU
-        )
-
-    def forward(self, x):
-        qkv, _ = self.to_qkv(x)
-        q, k, v = qkv.split(
-            [self.to_qkv.num_heads * self.head_dim,
-             self.to_qkv.num_kv_heads * self.head_dim,
-             self.to_qkv.num_kv_heads * self.head_dim],
-            dim=-1,
-        )
-        B, S, _ = x.shape
-        q = q.view(B, S, self.to_qkv.num_heads, self.head_dim)
-        k = k.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
-        v = v.view(B, S, self.to_qkv.num_kv_heads, self.head_dim)
-        out = self.attn(q, k, v)
-        out = out.reshape(B, S, -1)
-        out, _ = self.to_out(out)
-        return out
-```
-
-### QKV Fusion in load_weights
-
-When you fuse separate Q/K/V into `QKVParallelLinear`, map diffusers' separate weight names:
-
-```python
-stacked_params_mapping = [
-    ("to_qkv", "to_q", "q"),
-    ("to_qkv", "to_k", "k"),
-    ("to_qkv", "to_v", "v"),
-]
-
-def load_weights(self, weights):
-    params = dict(self.named_parameters())
-    loaded = set()
-    for name, tensor in weights:
-        for fused_name, orig_name, shard_id in stacked_params_mapping:
-            if orig_name in name:
-                name = name.replace(orig_name, fused_name)
-                param = params[name]
-                param.weight_loader(param, tensor, shard_id)
-                loaded.add(name)
-                break
-        else:
-            if name in params:
-                param = params[name]
-                if hasattr(param, "weight_loader"):
-                    param.weight_loader(param, tensor)
-                else:
-                    default_weight_loader(param, tensor)
-                loaded.add(name)
-    return loaded
-```
-
-### RMSNorm with TP
-
-When RMSNorm sits between TP-sharded dimensions, use `DistributedRMSNorm` — it computes global RMS via all-reduce across TP ranks. See the Wan2.2 implementation for the pattern.
-
-### TP Constraints
-
-- `num_heads % tp_size == 0`
-- `num_kv_heads % tp_size == 0`
-- Use `self.to_qkv.num_heads` (local per-GPU count), not total heads, for split sizes
-
-### Testing TP
-
-```bash
-python text_to_image.py --model Your-org/your-model \
-  --tensor-parallel-size 2 --output "tp_test.png"
-```
-
-**Verify**: speedup, memory reduction proportional to TP size, quality matches single-GPU.
-
-### Reference implementations
-
-| Model | Path |
-|-------|------|
-| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
-| FLUX | `vllm_omni/diffusion/models/flux/flux_transformer.py` |
-| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
-
----
-
-## Sequence Parallelism (SP / USP)
-
-SP splits sequence tokens across GPUs using Ulysses (all-to-all) or Ring (P2P) communication. It is applied non-intrusively via the `_sp_plan` dict — no changes to `forward()` logic.
-
-### Approach 1: Non-Intrusive `_sp_plan` (Recommended)
-
-The framework automatically registers hooks to shard inputs and gather outputs at `nn.Module` boundaries.
-
-#### Step 1: Identify module boundaries
-
-Find where tensors need sharding/gathering:
-
-```python
-class MyTransformer(nn.Module):
-    def __init__(self):
-        self.patch_embed = PatchEmbed()    # Before blocks
-        self.pos_embed = RoPE()            # RoPE may need splitting
-        self.blocks = nn.ModuleList([...]) # Blocks process sharded x
-        self.norm_out = LayerNorm()
-        self.proj_out = Linear()           # Gather after this
-
-    def forward(self, x):
-        x = self.patch_embed(x)
-        pos = self.pos_embed(x)
-        for block in self.blocks:
-            x = block(x, pos)
-        x = self.norm_out(x)
-        return self.proj_out(x)
-```
-
-#### Step 2: Handle inline operations
-
-`_sp_plan` hooks only work at `nn.Module` boundaries. Inline ops like `torch.cat()` must be extracted into submodules:
-
-```python
-# BAD: Inline — hooks can't intercept
-unified = torch.cat([x, cap_feats], dim=1)
-
-# GOOD: Extract into submodule
-class UnifiedPrepare(nn.Module):
-    def forward(self, x, cap_feats):
-        return torch.cat([x, cap_feats], dim=1)
-
-self.unified_prepare = UnifiedPrepare()
-unified = self.unified_prepare(x, cap_feats)
-```
-
-Common cases: `torch.cat()`, `pad_sequence()`, `tensor.reshape()`, complex preprocessing.
-
-#### Step 3: Write `_sp_plan`
-
-**Pattern 1: Shard at first block, gather at output** (most common)
-
-```python
-from vllm_omni.diffusion.distributed.sp_plan import (
-    SequenceParallelInput, SequenceParallelOutput,
-)
-
-class StandardTransformer(nn.Module):
-    _sp_plan = {
-        "blocks.0": {
-            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-        },
-        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-    }
-```
-
-**Pattern 2: Shard RoPE outputs separately**
-
-```python
-class TransformerWithRoPE(nn.Module):
-    _sp_plan = {
-        "rope": {
-            0: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
-            1: SequenceParallelInput(split_dim=1, expected_dims=4, split_output=True),
-        },
-        "blocks.0": {
-            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-        },
-        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-    }
-```
-
-**Pattern 3: Dual-stream (shard image, replicate text)**
-
-```python
-class DualStreamTransformer(nn.Module):
-    _sp_plan = {
-        "rope_preparer": {
-            2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
-            3: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True),
-        },
-        "transformer_blocks.0": {
-            "hidden_states": SequenceParallelInput(split_dim=1, expected_dims=3),
-        },
-        "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3),
-    }
-```
-
-### API Reference
-
-**SequenceParallelInput**:
-
-| Parameter | Type | Description |
-|-----------|------|-------------|
-| `split_dim` | int | Dimension to split (usually 1 for sequence) |
-| `expected_dims` | int/None | Expected tensor rank for validation |
-| `split_output` | bool | `False`: shard input params; `True`: shard output tensors |
-| `auto_pad` | bool | Auto-pad if sequence not divisible by world_size |
-
-**SequenceParallelOutput**:
-
-| Parameter | Type | Description |
-|-----------|------|-------------|
-| `gather_dim` | int | Dimension to gather (usually 1 for sequence) |
-| `expected_dims` | int/None | Expected tensor rank for validation |
-
-**Module naming**:
-
-| Key | Meaning |
-|-----|---------|
-| `"blocks.0"` | First element of ModuleList |
-| `"blocks.*"` | All elements of ModuleList |
-| `"rope"` | Named submodule |
-
-**Dictionary value types**:
-
-| Key type | split_output | Description |
-|----------|-------------|-------------|
-| `"param_name"` (str) | False | Shard input parameter by name |
-| `0, 1, ...` (int) | True | Shard output tuple by index |
-
-### Approach 2: Intrusive Modification (Complex Cases)
-
-For dynamic sharding logic that can't be expressed via `_sp_plan`:
-
-```python
-from vllm_omni.diffusion.distributed.sp_sharding import sp_shard, sp_gather
-
-def forward(self, hidden_states, ...):
-    if self.parallel_config.sequence_parallel_size > 1:
-        hidden_states = sp_shard(hidden_states, dim=1)
-    for block in self.blocks:
-        hidden_states = block(hidden_states)
-    if self.parallel_config.sequence_parallel_size > 1:
-        hidden_states = sp_gather(hidden_states, dim=1)
-    return hidden_states
-```
-
-Use intrusive modification as a last resort — `_sp_plan` is preferred for maintainability.
-
-### UAA Mode (Experimental)
-
-`ulysses_mode="advanced_uaa"` handles arbitrary sequence lengths and head counts that aren't divisible by `ulysses_degree`. Uses variable all-to-all split sizes and temporary head padding.
-
-### Combining SP methods
-
-Ulysses and Ring can be combined: `ulysses_degree × ring_degree = total SP GPUs`.
-
-```python
-DiffusionParallelConfig(ulysses_degree=2, ring_degree=2)  # 4 GPUs total
-```
-
-### Testing SP
-
-```bash
-# Offline
-python text_to_image.py --model Your-model --ulysses-degree 2
-
-# Online serving
-vllm serve Your-model --omni --usp 2
-```
-
-### Reference implementations
-
-| Model | Path |
-|-------|------|
-| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py` |
-| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
-| Z-Image | `vllm_omni/diffusion/models/z_image/z_image_transformer.py` |
-
----
-
-## CFG Parallelism
-
-Distributes positive/negative Classifier-Free Guidance branches across 2 GPUs.
-
-### Implementation
-
-Inherit `CFGParallelMixin` and implement `diffuse()`:
-
-```python
-from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
-
-class YourPipeline(nn.Module, CFGParallelMixin):
-    def diffuse(self, latents, timesteps, prompt_embeds, negative_embeds,
-                do_true_cfg, true_cfg_scale, **kwargs):
-        for i, t in enumerate(timesteps):
-            positive_kwargs = {
-                "hidden_states": latents,
-                "encoder_hidden_states": prompt_embeds,
-                "timestep": t,
-            }
-            negative_kwargs = {
-                "hidden_states": latents,
-                "encoder_hidden_states": negative_embeds,
-                "timestep": t,
-            } if do_true_cfg else None
-
-            noise_pred = self.predict_noise_maybe_with_cfg(
-                do_true_cfg=do_true_cfg,
-                true_cfg_scale=true_cfg_scale,
-                positive_kwargs=positive_kwargs,
-                negative_kwargs=negative_kwargs,
-            )
-            latents = self.scheduler_step_maybe_with_cfg(
-                noise_pred, t, latents, do_true_cfg
-            )
-        return latents
-```
-
-### Customization hooks
-
-| Method | Override when |
-|--------|-------------|
-| `predict_noise()` | Non-standard transformer call (e.g., dual-transformer like Wan2.2) |
-| `cfg_normalize_function()` | Custom normalization (e.g., LongCat with clamping) |
-| `combine_cfg_noise()` | Multi-output models (e.g., video + audio: CFG on video, positive-only on audio) |
-
-**Custom predict_noise** (Wan2.2 — selects active transformer):
-
-```python
-def predict_noise(self, current_model=None, **kwargs):
-    if current_model is None:
-        current_model = self.transformer
-    return current_model(**kwargs)[0]
-```
-
-**Custom combine_cfg_noise** (multi-output):
-
-```python
-def combine_cfg_noise(self, positive_pred, negative_pred, scale, normalize):
-    video_pos, audio_pos = positive_pred
-    video_neg, audio_neg = negative_pred
-    video_combined = super().combine_cfg_noise(video_pos, video_neg, scale, normalize)
-    return (video_combined, audio_pos)
-```
-
-### Composite scheduler for multi-output
-
-When each output has its own schedule:
-
-```python
-class VideoAudioScheduler:
-    def __init__(self, video_scheduler, audio_scheduler):
-        self.video_scheduler = video_scheduler
-        self.audio_scheduler = audio_scheduler
-
-    def step(self, noise_pred, t, latents, return_dict=False, generator=None):
-        video_out = self.video_scheduler.step(
-            noise_pred[0], t[0], latents[0], return_dict=False, generator=generator
-        )[0]
-        audio_out = self.audio_scheduler.step(
-            noise_pred[1], t[1], latents[1], return_dict=False, generator=generator
-        )[0]
-        return ((video_out, audio_out),)
-```
-
-### Testing CFG Parallel
-
-```bash
-python text_to_image.py --model Your-model \
-  --cfg-parallel-size 2 --cfg-scale 4.0 \
-  --negative-prompt "ugly, unclear"
-```
-
-**Constraint**: `guidance_scale > 1.0` and negative prompt must be provided.
-
-### Reference implementations
-
-| Model | Path |
-|-------|------|
-| Qwen-Image | `vllm_omni/diffusion/models/qwen_image/cfg_parallel.py` |
-| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py` |
-| Mixin base | `vllm_omni/diffusion/distributed/cfg_parallel.py` |
-
----
-
-## HSDP (Hybrid Sharded Data Parallel)
-
-Shards model weights across GPUs using PyTorch FSDP2. Reduces per-GPU VRAM without changing computation.
-
-### Implementation
-
-Add `_hsdp_shard_conditions` to the transformer class:
-
-```python
-class YourTransformer(nn.Module):
-    @staticmethod
-    def _is_transformer_block(name: str, module) -> bool:
-        return "blocks" in name and name.split(".")[-1].isdigit()
-
-    _hsdp_shard_conditions = [_is_transformer_block]
-```
-
-For MoE models, add additional conditions:
-
-```python
-class MoETransformer(nn.Module):
-    @staticmethod
-    def _is_transformer_block(name, module):
-        return "blocks" in name and name.split(".")[-1].isdigit()
-
-    @staticmethod
-    def _is_moe_expert(name, module):
-        return "experts" in name and name.split(".")[-1].isdigit()
-
-    _hsdp_shard_conditions = [_is_transformer_block, _is_moe_expert]
-```
-
-A module is sharded if **any** condition returns `True`.
-
-### Constraints
-
-- Cannot combine with Tensor Parallelism
-- For standalone HSDP (no other parallelism), `hsdp_shard_size` must be specified explicitly
-- Can combine with SP: HSDP reduces memory while SP distributes sequence
-
-### Testing HSDP
-
-```python
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-
-parallel_config = DiffusionParallelConfig(use_hsdp=True, hsdp_shard_size=8)
-omni = Omni(model="your-model", parallel_config=parallel_config)
-```
-
-Or CLI:
-
-```bash
-vllm serve Your-model --omni --use-hsdp
-```
-
-**Verify**: logs show "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root". Check VRAM reduction.
-
-### Reference implementations
-
-| Model | Path |
-|-------|------|
-| Wan2.2 | `vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py` |
-| HSDP Core | `vllm_omni/diffusion/distributed/hsdp.py` |
-
----
-
-## VAE Patch Parallelism
-
-Shards VAE decode spatially across ranks using tiling:
-
-```bash
-python text_to_image.py --model Your-model --vae-patch-parallel-size 4
-```
-
-Auto-enables `--vae-use-tiling`. Uses `DistributedAutoencoderKLWan` or similar distributed VAE. Set `vae_patch_parallel_size` in `DiffusionParallelConfig`.
-
----
-
-## Combining Parallelism Methods
-
-Common multi-GPU recipes:
-
-```bash
-# 4 GPUs: CFG (2) × Ulysses (2)
-python text_to_image.py --model Qwen/Qwen-Image \
-  --cfg-parallel-size 2 --ulysses-degree 2
-
-# 8 GPUs: Ulysses (4) × Ring (2) + VAE patch (8)
-python text_to_video.py --model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
-  --ulysses-degree 4 --ring-degree 2 --vae-patch-parallel-size 8
-
-# 2 GPUs: HSDP + Ulysses (cannot combine HSDP with TP)
-vllm serve Your-model --omni --use-hsdp --usp 2
-```
-
-## Discovering Parallelism Support
-
-Check which parallelism methods a model supports:
-
-| Check | How |
-|-------|-----|
-| **Ulysses / Ring SP** | Transformer defines `_sp_plan`. Search: `grep -r '_sp_plan' vllm_omni/diffusion/models/` |
-| **CFG Parallel** | Pipeline inherits `CFGParallelMixin`. Search: `grep -r 'CFGParallelMixin' vllm_omni/diffusion/models/` |
-| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
-| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
-
-The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.

.claude/skills/add-diffusion-model/references/transformer-adaptation.md

@@ -1,218 +0,0 @@
-# Transformer Adaptation Reference
-
-## Adapting a Diffusers Transformer to vLLM-Omni
-
-### Step-by-step Checklist
-
-1. Copy the transformer class from diffusers source
-2. Remove all mixin classes — inherit only from `nn.Module`
-3. Replace attention dispatch with `vllm_omni.diffusion.attention.layer.Attention`
-4. Replace logger with `vllm.logger.init_logger`
-5. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`
-6. Remove training-only code (gradient checkpointing, dropout)
-7. Add `load_weights()` method for weight loading from safetensors
-8. Add class-level attributes for acceleration features
-
-### Mixin Removal
-
-Remove these diffusers mixins (and their imports):
-
-```python
-# Remove all of these:
-from diffusers.models.modeling_utils import ModelMixin
-from diffusers.configuration_utils import ConfigMixin, register_to_config
-from diffusers.models.attention_processor import AttentionModuleMixin
-from diffusers.loaders import PeftAdapterMixin, FromOriginalModelMixin
-
-# Replace:
-class MyTransformer(ModelMixin, ConfigMixin, AttentionModuleMixin):
-# With:
-class MyTransformer(nn.Module):
-```
-
-Also remove `@register_to_config` decorators from `__init__`.
-
-### Attention Replacement
-
-The vLLM-Omni `Attention` layer wraps backend selection (FlashAttention, SDPA, SageAttn, etc.) and supports sequence parallelism hooks.
-
-**QKV tensor shape must be `[batch, seq_len, num_heads, head_dim]`.**
-
-#### Self-Attention Pattern
-
-```python
-from vllm_omni.diffusion.attention.layer import Attention
-from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata
-
-class SelfAttentionBlock(nn.Module):
-    def __init__(self, dim, num_heads):
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-
-        self.to_q = nn.Linear(dim, dim)
-        self.to_k = nn.Linear(dim, dim)
-        self.to_v = nn.Linear(dim, dim)
-        self.to_out = nn.Linear(dim, dim)
-
-        self.attn = Attention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            softmax_scale=1.0 / (self.head_dim ** 0.5),
-            causal=False,
-            num_kv_heads=num_heads,
-        )
-
-    def forward(self, x, attn_mask=None):
-        B, S, _ = x.shape
-        q = self.to_q(x).view(B, S, self.num_heads, self.head_dim)
-        k = self.to_k(x).view(B, S, self.num_heads, self.head_dim)
-        v = self.to_v(x).view(B, S, self.num_heads, self.head_dim)
-
-        attn_metadata = AttentionMetadata(attn_mask=attn_mask)
-        out = self.attn(q, k, v, attn_metadata=attn_metadata)
-        out = out.reshape(B, S, -1)
-        return self.to_out(out)
-```
-
-#### Fused QKV with TP (Advanced)
-
-For tensor parallelism, use vLLM's parallel linear layers:
-
-```python
-from vllm.model_executor.layers.linear import (
-    QKVParallelLinear, RowParallelLinear
-)
-
-class TPSelfAttention(nn.Module):
-    def __init__(self, dim, num_heads):
-        super().__init__()
-        self.num_heads = num_heads
-        self.head_dim = dim // num_heads
-
-        self.to_qkv = QKVParallelLinear(
-            hidden_size=dim,
-            head_size=self.head_dim,
-            total_num_heads=num_heads,
-            total_num_kv_heads=num_heads,
-        )
-        self.to_out = RowParallelLinear(dim, dim)
-
-        self.attn = Attention(
-            num_heads=num_heads,
-            head_size=self.head_dim,
-            softmax_scale=1.0 / (self.head_dim ** 0.5),
-            causal=False,
-            num_kv_heads=num_heads,
-        )
-```
-
-### Logger Replacement
-
-```python
-# Replace:
-from diffusers.utils import logging
-logger = logging.get_logger(__name__)
-
-# With:
-from vllm.logger import init_logger
-logger = init_logger(__name__)
-```
-
-### Custom Layers from vLLM-Omni
-
-Available utility layers:
-
-```python
-from vllm.model_executor.layers.layernorm import RMSNorm
-from vllm_omni.diffusion.layers.rope import RotaryEmbedding
-from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm
-```
-
-### Config Support
-
-```python
-from vllm_omni.diffusion.data import OmniDiffusionConfig
-
-class MyTransformer(nn.Module):
-    def __init__(self, *, od_config=None, num_layers=28, hidden_size=3072, **kwargs):
-        super().__init__()
-        self.od_config = od_config
-        self.parallel_config = od_config.parallel_config if od_config else None
-        # ... build layers
-```
-
-The transformer config values come from `model_index.json` → `config.json` in the transformer subfolder. The pipeline uses `get_transformer_config_kwargs(od_config.tf_model_config, TransformerClass)` to filter config keys to match the `__init__` signature.
-
-### Weight Loading
-
-The `load_weights` method receives an iterable of `(name, tensor)` from safetensors files, with the prefix (e.g., `"transformer."`) already stripped by the loader.
-
-```python
-from vllm.model_executor.model_loader.weight_utils import default_weight_loader
-
-class MyTransformer(nn.Module):
-    def load_weights(self, weights):
-        params = dict(self.named_parameters())
-        loaded = set()
-        for name, tensor in weights:
-            # Optional: remap names from diffusers to vllm-omni naming
-            # e.g., "ff.net.0.proj" -> "ff.net_0.proj"
-
-            if name in params:
-                param = params[name]
-                if hasattr(param, "weight_loader"):
-                    param.weight_loader(param, tensor)
-                else:
-                    default_weight_loader(param, tensor)
-                loaded.add(name)
-        return loaded
-```
-
-#### QKV Fusion in load_weights
-
-If you fused separate Q/K/V into a `QKVParallelLinear`, you need to map diffusers' separate weight names:
-
-```python
-stacked_params_mapping = [
-    ("to_qkv", "to_q", "q"),
-    ("to_qkv", "to_k", "k"),
-    ("to_qkv", "to_v", "v"),
-]
-
-def load_weights(self, weights):
-    params = dict(self.named_parameters())
-    loaded = set()
-    for name, tensor in weights:
-        for fused_name, orig_name, shard_id in stacked_params_mapping:
-            if orig_name in name:
-                name = name.replace(orig_name, fused_name)
-                param = params[name]
-                param.weight_loader(param, tensor, shard_id)
-                loaded.add(name)
-                break
-        else:
-            # Normal loading
-            ...
-    return loaded
-```
-
-### Class-Level Attributes for Features
-
-```python
-class MyTransformer(nn.Module):
-    # torch.compile: list block class names that repeat and can be compiled
-    _repeated_blocks = ["MyTransformerBlock"]
-
-    # CPU offload: attribute name of the nn.ModuleList containing blocks
-    _layerwise_offload_blocks_attr = "blocks"
-
-    # LoRA: mapping of fused param names to original param names
-    packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
-
-    # Sequence parallelism plan (advanced — add after basic impl works)
-    _sp_plan = {
-        "blocks.0": SequenceParallelInput(split_dim=1),
-        "proj_out": SequenceParallelOutput(gather_dim=1),
-    }
-```

.claude/skills/add-diffusion-model/references/troubleshooting.md

@@ -1,178 +0,0 @@
-# Troubleshooting Reference
-
-## Common Errors When Adding a Diffusion Model
-
-### ImportError / ModuleNotFoundError
-
-**Cause**: Missing or incorrect registration.
-
-**Fix checklist**:
-1. Model registered in `vllm_omni/diffusion/registry.py` `_DIFFUSION_MODELS` dict
-2. `__init__.py` exports the pipeline class
-3. Pipeline file exists at the correct path: `vllm_omni/diffusion/models/{folder}/{file}.py`
-4. Class name in registry matches the actual class name in the file
-
-### Shape Mismatch in Attention
-
-**Symptom**: `RuntimeError: shape mismatch` or `expected 4D tensor`
-
-**Cause**: QKV tensors not reshaped to `[batch, seq_len, num_heads, head_dim]`.
-
-**Fix**: Before calling `self.attn(q, k, v, ...)`, ensure:
-```python
-q = q.view(batch, seq_len, self.num_heads, self.head_dim)
-k = k.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
-v = v.view(batch, kv_seq_len, self.num_kv_heads, self.head_dim)
-```
-
-After attention, reshape back:
-```python
-out = out.reshape(batch, seq_len, -1)
-```
-
-### Weight Loading Failures
-
-**Symptom**: `RuntimeError: size mismatch for parameter ...` or missing keys
-
-**Debugging**:
-1. Print diffusers weight names: `safetensors.safe_open(path, "pt").keys()`
-2. Print model parameter names: `dict(model.named_parameters()).keys()`
-3. Compare and add name remappings in `load_weights()`
-
-**Common remappings needed**:
-- `ff.net.0.proj` → `ff.net_0.proj` (PyTorch Sequential indexing)
-- `.to_out.0.` → `.to_out.` (Sequential unwrapping)
-- `scale_shift_table` → moved to a wrapper module
-
-### Black/Blank/Noisy Output
-
-**Possible causes**:
-1. **Wrong latent normalization**: Check VAE expects latents scaled by `vae.config.scaling_factor`
-2. **Wrong scheduler**: Using the wrong scheduler class or wrong `flow_shift`
-3. **Missing CFG**: Some models require `guidance_scale > 1.0` with negative prompt
-4. **Wrong timestep format**: Some schedulers expect float, others expect int/long
-5. **Missing post-processing**: Raw VAE output may need denormalization
-
-**Quick test**: Run with diffusers directly using the same seed and compare latents at each step.
-
-### OOM (Out of Memory)
-
-**Solutions** (in order of preference):
-1. `--enforce-eager` to disable torch.compile (saves compile memory)
-2. `--enable-cpu-offload` for model-level offload
-3. `--enable-layerwise-offload` for block-level offload (better for large models)
-4. `--vae-use-slicing --vae-use-tiling` for VAE memory reduction
-5. Reduce resolution: `--height 480 --width 832`
-6. Use TP: `--tensor-parallel-size 2`
-
-### Different Output vs Diffusers Reference
-
-**Common causes**:
-1. **Attention backend difference**: FlashAttention vs SDPA may produce slightly different results. Set `DIFFUSION_ATTENTION_BACKEND=TORCH_SDPA` to match diffusers
-2. **Float precision**: vLLM-Omni may use bfloat16 where diffusers uses float32 for some operations
-3. **Missing normalization**: Check all LayerNorm/RMSNorm are preserved
-4. **Scheduler rounding**: Some schedulers have numerical sensitivity
-
-### Tensor Parallel Errors
-
-**Symptom**: `AssertionError: not divisible` or incorrect output with TP>1
-
-**Fix**:
-1. Verify `num_heads % tp_size == 0` and `num_kv_heads % tp_size == 0`
-2. Ensure `ColumnParallelLinear` / `RowParallelLinear` are used correctly
-3. Check that norms between parallel layers use distributed norm if needed
-4. Verify `load_weights` handles TP sharding for norm weights
-5. Use `self.to_qkv.num_heads` (local heads per GPU) for QKV split sizes, not total heads
-
-**Missing `input_is_parallel=True`**:
-
-`RowParallelLinear` expects sharded input from `ColumnParallelLinear`:
-```python
-self.w1 = ColumnParallelLinear(dim, hidden_dim, return_bias=False)
-self.w2 = RowParallelLinear(hidden_dim, dim, input_is_parallel=True, return_bias=False)
-```
-
-### Sequence Parallel Errors
-
-**Symptom**: Incorrect output or crashes with `--ulysses-degree N` or `--usp N`
-
-**Possible causes**:
-1. **Inline operations between shard/gather points**: `torch.cat()`, `pad_sequence()` etc. not at `nn.Module` boundaries. Fix: extract into submodule.
-2. **Wrong `split_dim`**: Check the tensor shape at the shard point. Sequence dimension is typically `dim=1` for `[B, S, D]` tensors.
-3. **RoPE not sharded**: If RoPE is computed separately, add it to `_sp_plan` with `split_output=True`.
-4. **Sequence not divisible by SP degree**: Use `auto_pad=True` in `SequenceParallelInput` or switch to `ulysses_mode="advanced_uaa"`.
-
-**Debugging**: Add `expected_dims=N` to `SequenceParallelInput`/`Output` for shape validation at runtime.
-
-### CFG Parallel Errors
-
-**Symptom**: CFG parallel not activating, no speedup
-
-**Fix checklist**:
-1. Pipeline inherits `CFGParallelMixin`
-2. `guidance_scale > 1.0`
-3. Negative prompt provided (even if empty string)
-4. `--cfg-parallel-size 2` specified
-5. `diffuse()` method calls `predict_noise_maybe_with_cfg()` and `scheduler_step_maybe_with_cfg()`
-
-**Symptom**: Different output with CFG parallel vs sequential
-
-**Possible cause**: Non-deterministic scheduler. Fix: pass `generator=torch.Generator(device).manual_seed(seed)` to `scheduler_step_maybe_with_cfg()`.
-
-### HSDP Errors
-
-**Symptom**: HSDP not activating or errors during weight loading
-
-**Fix checklist**:
-1. Transformer defines `_hsdp_shard_conditions` class attribute
-2. Shard condition functions return `True` for correct modules (test with `model.named_modules()`)
-3. Not combining with TP (HSDP and TP are incompatible)
-4. For standalone HSDP, `hsdp_shard_size` is specified explicitly
-
-**Verify**: Check logs for "HSDP Inference: replicate_size=..., shard_size=..." and "Sharded N modules + root".
-
-### Cache-DiT Not Applied
-
-**Symptom**: No speedup, no cache-related log messages
-
-**Fix checklist**:
-1. Model not in `_NO_CACHE_ACCELERATION` in `registry.py`
-2. Pipeline class name matches `CUSTOM_DIT_ENABLERS` key (if using custom enabler)
-3. `cache_backend="cache_dit"` specified
-4. Check logs for "Cache-dit enabled successfully on xxx"
-
-**Verify pipeline name**: `print(pipeline.__class__.__name__)` — must match registry key.
-
-### Cache-DiT Quality Degradation
-
-**Symptom**: Artifacts or lower quality with cache-dit
-
-**Fix**: Reduce aggressiveness:
-```python
-cache_config={
-    "residual_diff_threshold": 0.12,      # Lower from 0.24
-    "max_warmup_steps": 6,                # Increase from 4
-    "max_continuous_cached_steps": 2,      # Reduce if higher
-}
-```
-
-If quality is still poor, the model may need a custom enabler with per-block-list `ParamsModifier` tuning.
-
-### Model Not Detected / Wrong Pipeline Class
-
-**Symptom**: `ValueError: Model class ... not found in diffusion model registry`
-
-**Cause**: The model's `model_index.json` has a `_class_name` for the pipeline that doesn't match registry keys.
-
-**Fix**: The registry key must match the diffusers pipeline class name from `model_index.json`. If using a different name, map it in the registry:
-```python
-"DiffusersPipelineClassName": ("your_folder", "your_file", "YourVllmClassName"),
-```
-
-## Debugging Workflow
-
-1. **Add verbose logging**: Use `logger.info()` to print tensor shapes at each stage
-2. **Compare step-by-step**: Run diffusers and vllm-omni side by side, comparing tensors after each major operation
-3. **Use small configs**: Reduce `num_inference_steps=2`, small resolution for fast iteration
-4. **Test transformer isolation**: Feed the same input to both diffusers and vllm-omni transformers, compare outputs
-5. **Binary search for bugs**: Comment out blocks/layers to isolate where divergence starts

.claude/skills/add-tts-model/SKILL.md

@@ -1,284 +0,0 @@
----
-name: add-tts-model
-description: "Integrate a new text-to-speech model into vLLM-Omni from HuggingFace reference implementation through production-ready serving with streaming and CUDA graph acceleration. Use when adding a new TTS model, wiring stage separation for speech synthesis, enabling online voice generation serving, debugging TTS integration behavior, or building audio output pipelines."
----
-
-# TTS Model Integration Workflow
-
-## Overview
-
-```
-HF Reference -> Stage Separation -> Online Serving -> Async Chunk -> CUDA Graph
-   (Phase 1)      (Phase 2)          (Phase 3)        (Phase 4)     (Phase 5)
-```
-
-## Phase 1: HuggingFace Reference
-
-**Goal**: Understand the reference implementation and verify it produces correct audio.
-
-### Steps
-
-1. **Run the reference model** end-to-end using the official HuggingFace / GitHub code
-2. **Document the architecture**:
-   - What are the sub-models? (AR decoder, codec decoder, vocoder, etc.)
-   - What is the token vocabulary? (semantic codes, RVQ codebooks, special tokens)
-   - What is the output format? (sample rate, channels, codec type)
-3. **Capture reference outputs** for comparison during integration
-4. **Identify the config structure**: `config.json` fields, `model_type`, sub-model configs
-
-### Key Questions
-
-- How many codebooks? What are the codebook sizes?
-- What special tokens exist? (`<|voice|>`, `<|audio_start|>`, `<|im_end|>`, etc.)
-- What is the token-to-ID mapping for codec codes?
-- What is the hop length / frame rate of the codec?
-- Does the model support voice cloning? How? (reference audio encoding, speaker embeddings, etc.)
-
-### Deliverables
-
-- Working reference script that produces audio
-- Architecture diagram / notes
-- Token vocabulary mapping
-- Reference audio samples for regression testing
-
-## Phase 2: Stage Separation (Offline Inference)
-
-**Goal**: Split the model into vLLM-Omni stages and get offline inference working.
-
-### Steps
-
-1. **Register the model** in `vllm_omni/model_executor/models/registry.py`
-2. **Create config classes** (`configuration_<model>.py`) with `model_type` registration
-3. **Implement Stage 0** (AR model):
-   - Subclass appropriate base (e.g., wrap Qwen3 decoder layers)
-   - Implement `forward()` for autoregressive token generation
-   - Handle special token logic (start/stop tokens, codec token mapping)
-   - If dual-AR (like Fish Speech), implement Fast AR as a nested module
-4. **Implement Stage 1** (Decoder):
-   - Load codec weights (may need lazy loading from separate checkpoint)
-   - Implement `forward()`: codec codes -> audio waveform
-   - Return `OmniOutput` with `multimodal_outputs`
-5. **Create stage config YAML** defining both stages, memory allocation, and model paths
-6. **Create stage input processor** for prompt building
-7. **Write end2end.py** test script
-
-### Critical Parameters to Get Right
-
-| Parameter | Impact if Wrong |
-|-----------|----------------|
-| Hop length | Audio duration wrong, streaming noise |
-| Token ID mapping | Garbage codes -> noise output |
-| Codebook count/size | Shape mismatch crashes |
-| Stop token | Generation never stops or stops too early |
-| dtype / autocast | Numerical issues, silent quality degradation |
-| Repetition penalty | Must match reference (often 1.0 for TTS) |
-
-### Debugging Priority (from experience)
-
-When audio output is wrong, check in this order:
-
-1. **RoPE / attention**: Are position encodings correct? Is the attention mask right?
-2. **Normalization**: RMSNorm epsilon, layer norm placement (pre vs post)
-3. **Hop length**: Product of all upsample rates in the codec decoder
-4. **Token mapping**: Are codec IDs correctly offset from the vocabulary base?
-5. **Sampling parameters**: Temperature, top_k, top_p, repetition_penalty
-6. **Tensor layout**: Codebook-major vs frame-major ordering
-7. **dtype**: Float32 for codec decoders (autocast can corrupt audio)
-
-### Deliverables
-
-- Model files in `vllm_omni/model_executor/models/<model_name>/`
-- Stage config YAML
-- Working `end2end.py` with correct audio output
-- README.md in the example directory
-
-## Phase 3: Online Serving
-
-**Goal**: Expose the model via `/v1/audio/speech` API endpoint.
-
-### Steps
-
-1. **Register in `serving_speech.py`**:
-   - Add model stage name to `_TTS_MODEL_STAGES` set
-   - Add model detection flag (e.g., `_is_fish_speech`)
-   - Implement prompt builder method (e.g., `_build_fish_speech_prompt()`)
-2. **Handle model-specific parameters**:
-   - Voice cloning: `ref_audio` encoding and prompt injection
-   - `max_new_tokens` override in sampling params
-   - Model-specific default values
-3. **Create client scripts**: `speech_client.py`, `run_server.sh`
-4. **Test all response formats**: wav, mp3, flac, pcm
-5. **Add Gradio demo**: Interactive web UI with streaming support
-
-### Voice Cloning Pattern
-
-```python
-import base64
-from pathlib import Path
-
-def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
-    """Build prompt with reference audio for voice cloning in serving_speech.py."""
-    audio_bytes = Path(ref_audio_path).read_bytes()
-    codes = codec.encode(audio_bytes)  # Encode on CPU using model's codec (e.g., DAC)
-    token_ids = [code + codec.vocab_offset for code in codes.flatten().tolist()]
-    return [
-        {"role": "system", "content": f"<|voice|>{''.join(chr(t) for t in token_ids)}"},
-        {"role": "user", "content": text},
-    ]
-```
-
-### Deliverables
-
-- Updated `serving_speech.py` with model-specific prompt builder
-- Client scripts and server launcher
-- Gradio demo with streaming and voice cloning UI
-- Documentation (offline + online serving docs)
-
-## Phase 4: Async Chunk (Streaming)
-
-**Goal**: Enable inter-stage streaming so audio chunks are produced while AR generation continues.
-
-### Steps
-
-1. **Update stage config YAML**:
-   ```yaml
-   async_chunk: true
-   codec_chunk_frames: 25      # frames per chunk
-   codec_left_context_frames: 25  # overlap for smooth boundaries
-   ```
-2. **Implement chunk handling in Stage 1**:
-   - Accept partial input (chunk of codec codes)
-   - Handle left context for smooth audio boundaries
-   - Return partial audio in `OmniOutput`
-3. **Test streaming**:
-   - Verify audio quality matches non-streaming output
-   - Check for artifacts at chunk boundaries
-   - Measure TTFA (time to first audio)
-4. **Update online serving** to support `stream=true` with PCM output
-
-### Streaming Architecture
-
-```
-Stage 0 (AR)                    Stage 1 (Decoder)
-  |                                |
-  |-- chunk 0 (25 frames) ------> decode -> audio chunk 0 -> client
-  |-- chunk 1 (25 frames) ------> decode -> audio chunk 1 -> client
-  |-- chunk 2 (25 frames) ------> decode -> audio chunk 2 -> client
-  ...
-```
-
-### Key Considerations
-
-- **Left context overlap**: Prevents audible artifacts at chunk boundaries
-- **Hop length matters**: `context_audio_samples = context_frames * hop_length`
-- **First chunk latency**: Can use larger initial chunk for better quality, then smaller chunks
-
-### Deliverables
-
-- Updated stage config with async_chunk enabled
-- Smooth streaming audio without boundary artifacts
-- TTFA metrics
-
-## Phase 5: CUDA Graph Acceleration
-
-**Goal**: Capture the AR loop as a CUDA graph for significant speedup.
-
-### Steps
-
-1. **Identify the hot loop**: The AR decoding loop that runs N steps per token
-2. **Create static buffers**:
-   - KV caches with fixed max sequence length
-   - Pre-built causal masks and position tensors per step
-   - Static input/output tensors
-3. **Implement graph capture**:
-   - Warm up with real data
-   - Capture the forward pass
-   - Replay with updated inputs
-4. **Handle constraints**:
-   - Use `torch.argmax` instead of `torch.multinomial` (graph-safe)
-   - Fixed batch size (fall back to eager for other sizes)
-   - No dynamic control flow inside the graph
-
-### Example: Code Predictor CUDA Graph (Qwen3-TTS)
-
-```python
-import torch
-
-class CodePredictorGraph:
-    """Captures the 16-step code predictor AR loop as a single CUDA graph."""
-
-    def setup_graph(self, device: torch.device, kv_heads: int = 4, head_dim: int = 64):
-        self.num_steps = 16
-        self.kv_cache = torch.zeros(1, kv_heads, self.num_steps, head_dim, device=device)
-        self.positions = torch.arange(self.num_steps, device=device)
-        self.causal_mask = torch.tril(torch.ones(self.num_steps, self.num_steps, device=device))
-        self.input_buf = torch.zeros(1, 1, kv_heads * head_dim, device=device)
-        self.output_buf = torch.zeros(1, self.num_steps, device=device, dtype=torch.long)
-        # Warm up, then: self.graph = torch.cuda.CUDAGraph(); self.graph.capture(...)
-
-    def run_graph(self, initial_input: torch.Tensor) -> torch.Tensor:
-        self.input_buf.copy_(initial_input)
-        self.graph.replay()
-        return self.output_buf.clone()
-```
-
-### Performance Expectations
-
-Based on Qwen3-TTS code predictor experience:
-- **3-5x speedup** for the graphed component
-- Only effective for fixed batch sizes (typically batch_size=1)
-- Falls back to eager mode for unsupported configurations
-
-### Deliverables
-
-- CUDA graph implementation for the AR hot loop
-- Benchmark script comparing eager vs graph performance
-- Documentation of constraints and fallback behavior
-
-## Integration Checklist
-
-Use this checklist when integrating a new TTS model:
-
-### Phase 1: HF Reference
-- [ ] Reference model runs and produces correct audio
-- [ ] Architecture documented (stages, codebooks, tokens, sample rate)
-- [ ] Reference audio samples saved for comparison
-
-### Phase 2: Stage Separation
-- [ ] Model registered in `registry.py`
-- [ ] Config classes created with `model_type` registration
-- [ ] Stage 0 (AR) implemented and generates correct tokens
-- [ ] Stage 1 (Decoder) produces correct audio from tokens
-- [ ] Stage config YAML created
-- [ ] `end2end.py` produces audio matching reference quality
-- [ ] README.md written
-
-### Phase 3: Online Serving
-- [ ] Model added to `serving_speech.py`
-- [ ] Prompt builder handles text input correctly
-- [ ] Voice cloning works (if supported)
-- [ ] All response formats work (wav, mp3, flac, pcm)
-- [ ] Client scripts and server launcher created
-- [ ] Gradio demo working
-- [ ] Documentation added (offline + online docs, nav, supported models)
-
-### Phase 4: Async Chunk
-- [ ] Stage config updated with `async_chunk: true`
-- [ ] Stage 1 handles partial chunks correctly
-- [ ] No audio artifacts at chunk boundaries
-- [ ] Streaming via API (`stream=true`) works
-- [ ] TTFA measured and acceptable
-
-### Phase 5: CUDA Graph
-- [ ] Hot loop identified and profiled
-- [ ] Static buffers allocated
-- [ ] Graph captured and replays correctly
-- [ ] Benchmark shows meaningful speedup
-- [ ] Fallback to eager works for unsupported configs
-
-## References
-
-- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) -- supported models and usage
-- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) -- complete example of Phases 1-3
-- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) -- complete example of all 5 phases
-- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)

.claude/skills/readme.md

@@ -1,34 +0,0 @@
-# Claude Skills for vLLM-Omni
-
-This directory contains Claude Code skills maintained for the `vllm-omni`
-repository. These skills capture repeatable workflows for common contributor
-tasks such as model integration, pull request review, and release note
-generation.
-
-## Directory Structure
-
-Each skill lives in its own directory under `.claude/skills/`. A skill may
-include:
-
-- `SKILL.md`: the main workflow and operating instructions
-- `references/`: focused reference material used by the skill
-- `scripts/`: small helper scripts used by the skill
-
-## Available Skills
-
-- `add-diffusion-model`: guides integration of a new diffusion model into
-  `vllm-omni`
-- `add-omni-model`: covers addition of new omni-modality model support
-- `add-tts-model`: covers integration of new TTS models and related serving
-  workflows
-- `generate-release-note`: helps prepare release notes for repository changes
-- `review-pr`: provides a structured workflow for reviewing pull requests
-
-## Maintenance Guidelines
-
-- Keep skill names short and task-oriented.
-- Prefer repository-local paths, commands, and examples.
-- Avoid hardcoding fast-changing support matrices unless the skill is actively
-  maintained alongside those changes.
-- Treat skills as contributor tooling: optimize for clarity, actionability, and
-  low maintenance overhead.

.github/ISSUE_TEMPLATE/750-RFC.yml

@@ -7,7 +7,7 @@ body:
 - type: markdown
   attributes:
     value: >
-      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=in%3Atitle%20RFC%20sort%3Aupdated-desc) for reference.
+      #### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
 - type: textarea
   attributes:
     label: Motivation.
@@ -21,7 +21,7 @@ body:
     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/1jcgR3cDaUQH3VczD4ZcKaJAoYWHjCmnYzHkCNyz-9fk/edit?usp=sharing).
+      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

.gitignore

@@ -83,9 +83,6 @@ target/
 profile_default/
 ipython_config.py
 
-# uv
-uv.lock
-
 # pyenv
 #   For a library or package, you might want to ignore these files since the code is
 #   intended to run in multiple environments; otherwise, check them in:
@@ -158,19 +155,7 @@ cython_debug/
 
 # Claude
 CLAUDE.md
-/.claude/*
+.claude/
-!.claude/skills/
-!.claude/skills/readme.md
-!.claude/skills/add-diffusion-model/
-!.claude/skills/add-diffusion-model/SKILL.md
-!.claude/skills/add-diffusion-model/references/
-!.claude/skills/add-diffusion-model/references/*.md
-!.claude/skills/add-tts-model/
-!.claude/skills/add-tts-model/SKILL.md
-!.claude/skills/review-pr/
-!.claude/skills/review-pr/SKILL.md
-!.claude/skills/review-pr/references/
-!.claude/skills/review-pr/references/*.md
 
 # Codex
 AGENTS.md

benchmarks/qwen3-tts/README.md

@@ -32,9 +32,6 @@ 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 a Voice Clone model
-MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base 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
 

benchmarks/qwen3-tts/run_benchmark.sh

@@ -23,9 +23,6 @@
 #   # Use 1.7B model:
 #   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
 #
-#   # Use Voice Clone model
-#   MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-Base TASK_TYPE=Base 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
 #
@@ -38,7 +35,6 @@
 #   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)
-#   TASK_TYPE        - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
 
 set -euo pipefail
 
@@ -57,7 +53,6 @@ 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)"
-TASK_TYPE="${TASK_TYPE:-CustomVoice}"
 
 # Parse args
 RUN_ASYNC=true
@@ -82,7 +77,6 @@ echo " Concurrency:  ${CONCURRENCY}"
 echo " Port:         ${PORT}"
 echo " Stage config: ${STAGE_CONFIG}"
 echo " Results:      ${RESULT_DIR}"
-echo " Task type:    ${TASK_TYPE}"
 echo "============================================================"
 
 # Prepare stage config with correct GPU device and memory settings
@@ -201,8 +195,7 @@ run_bench() {
         --max-concurrency ${conc_args} \
         --num-warmups "${NUM_WARMUPS}" \
         --config-name "${config_name}" \
-        --result-dir "${RESULT_DIR}" \
+        --result-dir "${RESULT_DIR}"
-        --task-type "${TASK_TYPE}"
 
     stop_server
 
@@ -229,8 +222,7 @@ if [ "${RUN_HF}" = true ]; then
         --num-warmups "${NUM_WARMUPS}" \
         --gpu-device "${GPU_DEVICE}" \
         --config-name "hf_transformers" \
-        --result-dir "${RESULT_DIR}" \
+        --result-dir "${RESULT_DIR}"
-        --task-type "${TASK_TYPE}"
 
     # Allow GPU memory to settle
     sleep 5

benchmarks/qwen3-tts/transformers/bench_tts_hf.py

@@ -38,10 +38,6 @@ PROMPTS = [
     "It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
 ]
 
-REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
-REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
-INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
-
 
 @dataclass
 class BenchmarkResult:
@@ -79,29 +75,6 @@ class BenchmarkResult:
     per_request: list = field(default_factory=list)
 
 
-def generate_audio(model, prompt: str, args):
-    if args.task_type == "Base":
-        return model.generate_voice_clone(
-            text=prompt,
-            language=args.language,
-            ref_audio=REF_AUDIO,
-            ref_text=REF_TEXT,
-        )
-
-    if args.task_type == "VoiceDesign":
-        return model.generate_voice_design(
-            text=prompt,
-            language=args.language,
-            instruct=INSTRUCT,
-        )
-
-    return model.generate_custom_voice(
-        text=prompt,
-        language=args.language,
-        speaker=args.voice,
-    )
-
-
 def run_benchmark(args):
     from qwen_tts import Qwen3TTSModel
 
@@ -122,7 +95,11 @@ def run_benchmark(args):
         print(f"Warming up with {args.num_warmups} requests...")
         for i in range(args.num_warmups):
             p = PROMPTS[i % len(PROMPTS)]
-            wavs, sr = generate_audio(model, p, args)
+            wavs, sr = model.generate_custom_voice(
+                text=p,
+                language=args.language,
+                speaker=args.voice,
+            )
         # Sync GPU
         torch.cuda.synchronize(device)
         print("Warmup done.")
@@ -147,7 +124,11 @@ def run_benchmark(args):
             torch.cuda.synchronize(device)
             st = time.perf_counter()
 
-            wavs, sr = generate_audio(model, prompt, args)
+            wavs, sr = model.generate_custom_voice(
+                text=prompt,
+                language=args.language,
+                speaker=args.voice,
+            )
 
             torch.cuda.synchronize(device)
             elapsed = time.perf_counter() - st
@@ -287,7 +268,6 @@ def parse_args():
     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("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
     parser.add_argument(
         "--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
     )

benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py

@@ -37,9 +37,6 @@ PROMPTS = [
     "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.",
 ]
-REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
-REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
-INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
 
 
 @dataclass
@@ -96,39 +93,22 @@ def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width
     return num_samples / sample_rate
 
 
-def create_payload(
-    prompt: str, task_type: str = "CustomVoice", voice: str = "vivian", language: str = "English"
-) -> dict:
-    payload = {
-        "input": prompt,
-        "language": language,
-        "stream": True,
-        "response_format": "pcm",
-        "task_type": task_type,
-    }
-
-    if task_type == "Base":
-        payload["ref_audio"] = REF_AUDIO
-        payload["ref_text"] = REF_TEXT
-    elif task_type == "CustomVoice":
-        payload["voice"] = voice
-    elif task_type == "VoiceDesign":
-        payload["instructions"] = INSTRUCT
-
-    return payload
-
-
 async def send_tts_request(
     session: aiohttp.ClientSession,
     api_url: str,
     prompt: str,
-    task_type: str = "CustomVoice",
     voice: str = "vivian",
     language: str = "English",
     pbar: tqdm | None = None,
 ) -> RequestResult:
     """Send a streaming TTS request and measure latency metrics."""
-    payload = create_payload(prompt, task_type, voice, language)
+    payload = {
+        "input": prompt,
+        "voice": voice,
+        "language": language,
+        "stream": True,
+        "response_format": "pcm",
+    }
 
     result = RequestResult(prompt=prompt)
     st = time.perf_counter()
@@ -173,7 +153,6 @@ async def run_benchmark(
     num_prompts: int,
     max_concurrency: int,
     num_warmups: int = 3,
-    task_type: str = "CustomVoice",
     voice: str = "vivian",
     language: str = "English",
 ) -> BenchmarkResult:
@@ -196,7 +175,7 @@ async def run_benchmark(
         warmup_tasks = []
         for i in range(num_warmups):
             prompt = PROMPTS[i % len(PROMPTS)]
-            warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
+            warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
         await asyncio.gather(*warmup_tasks)
         print("  Warmup done.")
 
@@ -210,7 +189,7 @@ async def run_benchmark(
 
     async def limited_request(prompt):
         async with semaphore:
-            return await send_tts_request(session, api_url, prompt, task_type, voice, language, pbar)
+            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]
@@ -327,7 +306,6 @@ async def main(args):
             num_prompts=args.num_prompts,
             max_concurrency=concurrency,
             num_warmups=args.num_warmups,
-            task_type=args.task_type,
             voice=args.voice,
             language=args.language,
         )
@@ -356,7 +334,6 @@ def parse_args():
         "--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("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
     parser.add_argument("--voice", type=str, default="vivian")
     parser.add_argument("--language", type=str, default="English")
     parser.add_argument(

docker/Dockerfile.ci

@@ -1,5 +1,5 @@
 ARG VLLM_BASE_IMAGE=vllm/vllm-openai
-ARG VLLM_BASE_TAG=v0.19.0
+ARG VLLM_BASE_TAG=v0.18.0
 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
 ARG APP_DIR=/workspace/vllm-omni
 WORKDIR ${APP_DIR}
@@ -11,7 +11,9 @@ RUN apt-get update && \
     apt-get clean && \
     rm -rf /var/lib/apt/lists/*
 
-RUN uv pip install --system ".[dev]"
+# 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
 

docker/Dockerfile.cuda

@@ -1,22 +0,0 @@
-ARG BASE_IMAGE=vllm/vllm-openai:v0.19.0
-FROM ${BASE_IMAGE}
-
-ARG COMMON_WORKDIR=/app
-
-WORKDIR ${COMMON_WORKDIR}
-
-# Step 1: Setup - Install system dependencies
-RUN apt-get update && \
-    apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
-    apt-get clean && \
-    rm -rf /var/lib/apt/lists/*
-
-RUN mkdir -p ${COMMON_WORKDIR}/vllm-omni
-
-# Step 2: Copy vllm-omni code and install
-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 "."
-
-RUN ln -sf /usr/bin/python3 /usr/bin/python
-
-ENTRYPOINT []

docker/Dockerfile.npu

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1
+ARG VLLM_ASCEND_TAG=v0.17.0rc1
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-# WORKDIR /vllm-workspace/vllm
+WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
+RUN git fetch origin --tags && git checkout v0.18.0
 
-# WORKDIR /vllm-workspace/vllm-ascend
+WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
+# Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+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/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/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 && \
+    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 install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
+    python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.npu.a3

@@ -1,20 +1,20 @@
 ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
-ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
+ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
 FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
 
-# WORKDIR /vllm-workspace/vllm
+WORKDIR /vllm-workspace/vllm
-# RUN git fetch origin --tags && git checkout v0.18.0
+RUN git fetch origin --tags && git checkout v0.18.0
 
-# WORKDIR /vllm-workspace/vllm-ascend
+WORKDIR /vllm-workspace/vllm-ascend
-# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
-# # Install vllm-ascend
+# Install vllm-ascend
-# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
+# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
-# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
+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/ascend-toolkit/set_env.sh && \
-#     source /usr/local/Ascend/nnal/atb/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 && \
+    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 install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
-#     python3 -m pip cache purge
+    python3 -m pip cache purge
 
 ARG APP_DIR=/vllm-workspace/vllm-omni
 WORKDIR ${APP_DIR}

docker/Dockerfile.rocm

@@ -1,4 +1,4 @@
-ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
+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.

docker/Dockerfile.xpu

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

docs/.nav.yml

@@ -98,7 +98,6 @@ nav:
       - design/feature/ray_based_execution.md
       - design/feature/omni_connectors/
       - design/feature/cfg_parallel.md
-      - design/feature/expert_parallel.md
       - design/feature/sequence_parallel.md
       - design/feature/tensor_parallel.md
       - design/feature/vae_parallel.md

docs/api/README.md

@@ -5,6 +5,7 @@
 Main entry points for vLLM-Omni inference and serving.
 
 - [vllm_omni.entrypoints.async_omni.AsyncOmni][]
+- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
 - [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
 - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
 - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]

docs/contributing/ci/CI_5levels.md

@@ -168,6 +168,7 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜

docs/contributing/ci/tests_style.md

@@ -73,6 +73,7 @@ vllm_omni/                                    tests/
 │   └── arg_utils.py                            │   └── test_arg_utils.py               ⬜
 │
 ├── entrypoints/                        →     ├── entrypoints/
+│   ├── async_omni_diffusion.py                 │   ├── test_async_omni_diffusion_config.py  ✅
 │   ├── stage_utils.py                          │   ├── test_stage_utils.py            ✅
 │   ├── cli/                                     │   ├── cli/                           (benchmarks/test_serve_cli.py covers CLI serve)
 │   │   └── ...                                  │   │   └── test_*.py                  ⬜

docs/contributing/model/adding_diffusion_model.md

@@ -802,7 +802,7 @@ omni = Omni(model="your-model", enable_layerwise_offload=True)
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute name containing transformer blocks
+    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks

docs/design/feature/expert_parallel.md

@@ -1,221 +0,0 @@
-# Expert Parallel
-
-This section describes how to add Expert Parallel (EP) to a diffusion transformer that uses Mixture-of-Experts (MoE) layers.
-We use **HunyuanImage3.0** as the reference implementation.
-
----
-
-## Table of Contents
-
-- [Overview](#overview)
-- [Step-by-Step Implementation](#step-by-step-implementation)
-- [Testing](#testing)
-- [Reference Implementations](#reference-implementations)
-- [Summary](#summary)
-
----
-
-## Overview
-
-### What is Expert Parallel?
-
-**Expert Parallel** is a parallelism strategy in Mixture-of-Experts (MoE) models that distributes different expert networks across distinct computational devices. Each device holds and computes only a subset of experts (local experts), with tokens dispatched to and gathered from remote devices via collective communication operations (e.g., All-to-All, All-Gather).
-
-| Backend | Description |
-|---------|-------------|
-| `allgather_reducescatter` | Default backend based on allgather/reducescatter primitives, suitable for general EP+DP deployments.|
-
-## Configuration
-
-Enable EP by setting the `--enable-expert-parallel` flag. The EP size is automatically calculated as:
-
-```text
-EP_SIZE = TP_SIZE × SP_SIZE × CFG_SIZE × DP_SIZE
-```
-
-
-Where:
-
-- `TP_SIZE`: Tensor parallel size
-- `SP_SIZE`: Sequence parallel size
-- `CFG_SIZE`: Classifier-free guidance parallel size
-- `DP_SIZE`: Data parallel size
-- `EP_SIZE`: Expert parallel size (computed automatically)
-
-Note:
-- Expert parallelism is only applicable to Mixture-of-Experts (MoE) models.
-- The EP group is created **per pipeline stage**, meaning it includes all ranks that participate in model parallelism except pipeline parallelism.
-- The underlying communication pattern for expert parallelism is **All-to-All** among the ranks in the EP group.
-
-For example, consider a configuration with `TP=2`, `SP=1`, `CFG=2`, and `DP=4` (total 2×1×2×4 = 16 GPUs).
-
-- Expert layers are handled by an EP group of size 16.
-
-- Attention layers use tensor parallelism of size 2 within each of the 8 DP groups (because `DP×CFG×SP = 4×2×1 = 8` groups, each containing the 2 TP ranks). Inside each such group, the attention weights are sharded across the 2 GPUs.
-
-
-## Step-by-Step Implementation
-
-### Step 1: Configure Expert Parallelism Settings
-
-Calculate local experts per rank:
-
-```
-ep_size = 8  # Expert Parallel size (typically equals TP size)
-num_experts = 64
-num_local_experts = num_experts // ep_size  # 8 experts per card
-
-# Check divisibility
-assert num_experts % ep_size == 0, "Experts must be divisible by EP size"
-```
-
-### Step 2: Use Sparse MoE Block to enable EP routing.
-
-Example:
-```
-from vllm.model_executor.layers.linear import ReplicatedLinear
-class HunYuanSparseMoeBlock(nn.Module):
-    def __init__(
-        self,
-        config: PretrainedConfig,
-        layer_id: int = -1,
-        prefix: str = "",
-    ):
-        super().__init__()
-        self.tp_size = get_tensor_model_parallel_world_size()
-        self.n_routed_experts = config.num_experts  # 64
-
-        # Calculate local experts per rank (key for EP)
-        if self.tp_size > self.n_routed_experts:
-            raise ValueError(f"TP size {self.tp_size} > experts {self.n_routed_experts}")
-
-        # Routing gate (replicated on all ranks, computes scores for all tokens to all experts)
-        self.gate = ReplicatedLinear(
-            config.hidden_size,
-            config.num_experts,
-            bias=False,
-            quant_config=None,
-            prefix=f"{prefix}.gate",
-        )
-
-        # EP expert layer (factory loads platform-specific implementation)
-        self.experts = HunyuanFusedMoE(...)
-```
-**Key Points:**
-- gate is **ReplicatedLinear** (replicated on all ranks)
-- experts is created via **HunyuanFusedMoE factory**, which automatically handles EP dispatch
-
-### Step 3: Initialize EP Runtime
-
-Initialize the EP communication context before model loading.
-```
-from vllm.utils.import_utils import resolve_obj_by_qualname
-# Call during __init__ or model loading
-op_name = "hunyuan_fused_moe"
-
-# Prepare EP runtime: establish communication groups, assign local expert indices, init _expert_map
-current_omni_platform.prepare_diffusion_op_runtime(op_name)
-
-# Factory automatically resolves platform implementation (GPU: FusedMoE / NPU: AscendFusedMoE)
-impl = resolve_obj_by_qualname(
-    current_omni_platform.get_diffusion_model_impl_qualname(op_name)
-)
-```
-
-### Step 4: Expert Weight Mapping & Loading
-
-Each rank loads only the expert weights assigned to its local allocation.
-```
-# Get expert parameter mapping (different per rank)
-expert_mapping = HunyuanFusedMoE.make_expert_params_mapping(
-    model=self,
-    ckpt_gate_proj_name="gate_proj",
-    ckpt_down_proj_name="down_proj",
-    ckpt_up_proj_name="up_proj",
-    num_experts=64,
-    num_redundant_experts=0,
-)
-# Returns: [(param_name, weight_name, expert_id, shard_id), ...]
-# Note: Each rank only contains mappings for its local expert_ids
-
-# Filter non-local experts during loading
-for name, loaded_weight in weights:
-    if "mlp.experts" in name:
-        # Parse expert_id from weight name (implementation needed)
-        expert_id = parse_expert_id_from_name(name)
-        local_expert_start = (ep_rank) * num_local_experts
-        local_expert_end = (ep_rank + 1) * num_local_experts
-
-        if not (local_expert_start <= expert_id < local_expert_end):
-            continue  # Skip non-local expert weights
-```
-### Step 5: Forward Pass with EP
-
-Example (MoE Forward):
-```
-def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
-    orig_shape = hidden_states.shape
-    hidden_states = hidden_states.view(-1, hidden_states.shape[-1])
-
-    # 1. Global routing computation (all tokens, all expert scores)
-    # hidden_states: [num_tokens, hidden_dim] (full tensor)
-    router_logits, _ = self.gate(hidden_states)  # [num_tokens, num_experts]
-
-    # 2. EP dispatch and compute (HunyuanFusedMoE handles all_to_all internally)
-    # - Dispatch: Send tokens to target ranks based on router_logits
-    # - Local Compute: Each rank processes only its num_local_experts
-    # - Combine: Results returned to original token positions
-    final_hidden_states = self.experts(
-        hidden_states=hidden_states,
-        router_logits=router_logits,
-    )
-
-    # 3. Add shared expert output (not EP, computed on all ranks)
-    if self.shared_mlp is not None:
-        shared_out = self.shared_mlp(hidden_states)
-        final_hidden_states = final_hidden_states + shared_out
-
-    # 4. Tensor Parallel All-Reduce (synchronize across TP group)
-    if self.tp_size > 1:
-        final_hidden_states = self.experts.maybe_all_reduce_tensor_model_parallel(
-            final_hidden_states
-        )
-
-    return final_hidden_states.view(orig_shape)
-```
-
-## Testing
-After adding Expert Parallel support, test via command line:
-```bash
-cd examples/offline_inference/text_to_image
-python text_to_image.py \
-    --model Your-org/your-model \
-    --prompt "a cup of coffee on the table" \
-    --output "ep_enabled.png" \
-    --num-inference-steps 50 \
-    --guidance-scale 5.0 \
-    --tensor-parallel-size 8 \
-    --seed 1234 \
-    --enable-expert-parallel
-```
-
-vLLM‑Omni currently focuses on core diffusion model inference acceleration, so the Expert Parallel implementation includes only the basic multi‑GPU expert sharding functionality (enabled via --enable-expert-parallel). Advanced features such as communication backend selection (--all2all-backend), load balancing (--enable-eplb and its configuration), and multi‑node deployment belong to the extended capabilities of the main vLLM project and have not yet been integrated into Omni.
-
-## Reference Implementations
-
-Complete examples in the codebase:
-
-| Model | Path | Pattern | Notes |
-|-------|------|---------|-------|
-| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py` | Standard EP | Full implementation with validation |
-| **EP Tests** | `vllm-omni/tests/e2e/offline_inference/test_expert_parallel.py` | E2E testing | EP correctness and performance |
-| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic |
-
----
-## Summary
-
-Adding Expert Parallel support to diffusion model:
-
-1. **Identify MoE layers** - Locate the router and expert networks in each transformer block.
-2. **Validate EP constraints** – Ensure num_experts is divisible by expert_parallel_size.
-3. **Test** - Run with enable-expert-parallel, check memory reduction, speedup, and output quality against single‑GPU baseline.

docs/design/feature/omni_connectors/mooncake_transfer_engine_connector.md

@@ -33,8 +33,8 @@ runtime:
         zmq_port: 50051               # ZMQ base port (see "Port Offset Scheme" below)
         protocol: "rdma"              # "rdma" or "tcp"
         device_name: ""               # RDMA device (e.g., "mlx5_0"), empty for auto-detect
-        memory_pool_size: 4294967296  # 4 GB (CPU); use 2147483648 (2 GB) for GPU
+        memory_pool_size: 2147483648  # 2GB memory pool
-        memory_pool_device: "cpu"     # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
+        memory_pool_device: "cpu"     # "cpu" for pinned memory, "cuda" for GPUDirect RDMA
 ```
 
 Wire stages to the connector:
@@ -64,8 +64,8 @@ stage_args:
 
 | Parameter | Default | Description |
 |---|---|---|
-| `memory_pool_size` | 4 GB (CPU) / 2 GB (GPU) | Total size of the RDMA-registered memory pool in bytes. Recommended 4 GB for CPU pinned memory; 2 GB for GPU VRAM to conserve device memory. |
+| `memory_pool_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
-| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended, works on all topologies). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity, PIX topology). |
+| `memory_pool_device` | `"cpu"` | `"cpu"`: pinned host memory (recommended). `"cuda"`: GPU VRAM for GPUDirect RDMA (requires NIC-GPU direct PCIe connectivity). |
 
 ### Networking
 
@@ -107,10 +107,10 @@ receiver_connect  = remote_side_channel_port + tp_rank
 
 ## Memory Pool Modes
 
-| Mode | Config | Recommended Pool Size | Data Flow | Best For |
+| Mode | Config | Data Flow | Best For |
-|---|---|---|---|---|
+|---|---|---|---|
-| CPU Pinned | `memory_pool_device: "cpu"` | 4 GB | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
+| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
-| GPUDirect | `memory_pool_device: "cuda"` | 2 GB | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
+| GPUDirect | `memory_pool_device: "cuda"` | GPU → GPU pool → RDMA (NIC reads GPU BAR1) → GPU pool | NIC-GPU direct PCIe (PIX topology) |
 
 > **Note**: GPUDirect RDMA requires the NIC and GPU to share a direct PCIe
 > switch (PIX topology). On systems where they are connected via PXB or NODE,

docs/design/feature/vae_parallel.md

@@ -1,15 +1,14 @@
 # VAE Patch Parallelism
 
 This document describes how to add **VAE Patch Parallelism** support to a diffusion model.
-We use **Qwen-Image** as the reference implementation for decode parallel, and **Wan2.2** for encode parallel.
+We use **Qwen-Image** as the reference implementation.
 
 ---
 
 ## Table of Contents
 
 - [Overview](#overview)
-- [Step-by-Step Implementation (Decode)](#step-by-step-implementation-decode)
+- [Step-by-Step Implementation](#step-by-step-implementation)
-- [Encode Parallel Implementation](#encode-parallel-implementation)
 - [Testing](#testing)
 - [Reference Implementations](#reference-implementations)
 - [Summary](#summary)
@@ -20,13 +19,13 @@ We use **Qwen-Image** as the reference implementation for decode parallel, and *
 
 ### What is Vae Patch parallel?
 
-**VAE Patch Parallelism** is an acceleration technique for both **encoding** and **decoding**. Instead of processing the entire tensor at once, the tensor is:
+**VAE Patch Parallelism** is a decoding acceleration technique. Instead of decoding the entire latent tensor at once, the latent tensor is:
 
 + Split into multiple spatial tiles
 
 + Distributed across multiple ranks
 
-+ Encoded/Decoded in parallel
++ Decoded in parallel
 
 + Merged to reconstruct the final output
 
@@ -36,17 +35,10 @@ This approach:
 
 + Reduces peak memory usage per device
 
-+ Accelerates encoding/decoding latency
++ Accelerates decoding latency
-
-### When to Use Encode vs Decode Parallel
-
-| Operation | Use Case | Example |
-|-----------|----------|---------|
-| **Decode Parallel** | Text-to-Image, Text-to-Video | Latent → Image/Video |
-| **Encode Parallel** | Image-to-Video (I2V) | Image → Latent (for conditioning) |
 
 ### Architecture
-We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE encoding/decoding.
+We introduce **DistributedVaeExecutor** as the core component responsible for distributed VAE decoding.
 
 The executor is model-agnostic and accepts three function parameters:
 
@@ -92,7 +84,7 @@ Therefore:
 
 + Merge must perform blending to avoid seams
 
-## Step-by-Step Implementation (Decode)
+## Step-by-Step Implementation
 
 ### Step 1: Implement DistributedAutoencoderKLQwenImage
 `QwenImagePipeline` use `AutoencoderKLQwenImage` for vae, so implement a distributed version:
@@ -213,14 +205,14 @@ def tile_merge(self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid
 We need to override tiled_decode, the main logic is:
 + check distributed is enabled
 + select split/exec/merge
-+ Invoke self.distributed_executor.execute to decode
++ Invoke self.distributed_decoder.execute to decode
 ```
 def tiled_decode(self, z: torch.Tensor, return_dict: bool = True):
     if not self.is_distributed_enabled():
         return super().tiled_decode(z, return_dict=return_dict)
 
     logger.info("Decode run with distributed executor")
-    result = self.distributed_executor.execute(
+    result = self.distributed_decoder.execute(
         z,
         DistributedOperator(split=self.tile_split, exec=self.tile_exec, merge=self.tile_merge),
         broadcast_result=True,
@@ -251,166 +243,6 @@ class YourModelPipeline(nn.Module):
 +       ).to(self.device)
 ```
 
-## Encode Parallel Implementation
-
-For models that require VAE encoding (e.g., Image-to-Video), you can also parallelize the encode operation. We use **Wan2.2** as the reference implementation.
-
-### Step 1: Implement encode_tile_split
-
-Similar to decode, split the input tensor into tiles. Key considerations:
-
-+ **Patchify handling**: If the model uses `patch_size`, scale tile parameters accordingly
-+ **Temporal chunking**: Video VAEs may have temporal compression (e.g., 4x)
-
-```python
-def encode_tile_split(self, x: torch.Tensor) -> tuple[list[TileTask], GridSpec]:
-    _, _, num_frames, height, width = x.shape
-    encode_spatial_compression_ratio = self.spatial_compression_ratio
-
-    # Scale tile parameters for patchified coordinate system
-    tile_sample_min_height = self.tile_sample_min_height
-    tile_sample_min_width = self.tile_sample_min_width
-    tile_sample_stride_height = self.tile_sample_stride_height
-    tile_sample_stride_width = self.tile_sample_stride_width
-
-    if self.config.patch_size is not None:
-        # When input is patchified, scale tile parameters accordingly
-        encode_spatial_compression_ratio = self.spatial_compression_ratio // self.config.patch_size
-        tile_sample_min_height = tile_sample_min_height // self.config.patch_size
-        tile_sample_min_width = tile_sample_min_width // self.config.patch_size
-        tile_sample_stride_height = tile_sample_stride_height // self.config.patch_size
-        tile_sample_stride_width = tile_sample_stride_width // self.config.patch_size
-
-    latent_height = height // encode_spatial_compression_ratio
-    latent_width = width // encode_spatial_compression_ratio
-
-    tile_latent_min_height = tile_sample_min_height // encode_spatial_compression_ratio
-    tile_latent_min_width = tile_sample_min_width // encode_spatial_compression_ratio
-    tile_latent_stride_height = tile_sample_stride_height // encode_spatial_compression_ratio
-    tile_latent_stride_width = tile_sample_stride_width // encode_spatial_compression_ratio
-
-    blend_height = tile_latent_min_height - tile_latent_stride_height
-    blend_width = tile_latent_min_width - tile_latent_stride_width
-
-    tiletask_list = []
-    # Use temporal compression ratio from config instead of hardcoding
-    temporal_compression = self.config.scale_factor_temporal
-
-    for i in range(0, height, tile_sample_stride_height):
-        for j in range(0, width, tile_sample_stride_width):
-            time_list = []
-            frame_range = 1 + (num_frames - 1) // temporal_compression
-            for k in range(frame_range):
-                if k == 0:
-                    tile = x[:, :, :1, i : i + tile_sample_min_height, j : j + tile_sample_min_width]
-                else:
-                    tile = x[
-                        :, :,
-                        1 + temporal_compression * (k - 1) : 1 + temporal_compression * k,
-                        i : i + tile_sample_min_height,
-                        j : j + tile_sample_min_width,
-                    ]
-                time_list.append(tile)
-            tiletask_list.append(
-                TileTask(len(tiletask_list), (i // tile_sample_stride_height, j // tile_sample_stride_width),
-                         time_list, workload=time_list[0].shape[3] * time_list[0].shape[4])
-            )
-
-    grid_spec = GridSpec(
-        split_dims=(3, 4),
-        grid_shape=(tiletask_list[-1].grid_coord[0] + 1, tiletask_list[-1].grid_coord[1] + 1),
-        tile_spec={
-            "latent_height": latent_height, "latent_width": latent_width,
-            "blend_height": blend_height, "blend_width": blend_width,
-            "tile_latent_stride_height": tile_latent_stride_height,
-            "tile_latent_stride_width": tile_latent_stride_width,
-        },
-        output_dtype=self.dtype,
-    )
-    return tiletask_list, grid_spec
-```
-
-### Step 2: Implement encode_tile_exec
-
-```python
-def encode_tile_exec(self, task: TileTask) -> torch.Tensor:
-    """Encode a single sample tile into latent space."""
-    self.clear_cache()
-    time = []
-    for k, tile in enumerate(task.tensor):
-        self._enc_conv_idx = [0]
-        encoded = self.encoder(tile, feat_cache=self._enc_feat_map, feat_idx=self._enc_conv_idx)
-        encoded = self.quant_conv(encoded)
-        time.append(encoded)
-    result = torch.cat(time, dim=2)
-    self.clear_cache()
-    return result
-```
-
-### Step 3: Implement encode_tile_merge
-
-```python
-def encode_tile_merge(
-    self, coord_tensor_map: dict[tuple[int, ...], torch.Tensor], grid_spec: GridSpec
-) -> torch.Tensor:
-    """Merge encoded tiles into a full latent tensor."""
-    grid_h, grid_w = grid_spec.grid_shape
-    result_rows = []
-    for i in range(grid_h):
-        result_row = []
-        for j in range(grid_w):
-            tile = coord_tensor_map[(i, j)]
-            if i > 0:
-                tile = self.blend_v(coord_tensor_map[(i - 1, j)], tile, grid_spec.tile_spec["blend_height"])
-            if j > 0:
-                tile = self.blend_h(coord_tensor_map[(i, j - 1)], tile, grid_spec.tile_spec["blend_width"])
-            result_row.append(tile[:, :, :,
-                : grid_spec.tile_spec["tile_latent_stride_height"],
-                : grid_spec.tile_spec["tile_latent_stride_width"]])
-        result_rows.append(torch.cat(result_row, dim=-1))
-
-    enc = torch.cat(result_rows, dim=3)[
-        :, :, :, : grid_spec.tile_spec["latent_height"], : grid_spec.tile_spec["latent_width"]
-    ]
-    return enc
-```
-
-### Step 4: Override tiled_encode method
-
-Override `tiled_encode` instead of `encode`. The parent's `_encode()` handles patchify before calling `tiled_encode()`, so input `x` is already patchified.
-
-```python
-def tiled_encode(self, x: torch.Tensor) -> torch.Tensor:
-    """
-    Encode using distributed VAE executor.
-
-    Note: x is already patchified by parent's _encode() before calling this method.
-    """
-    if not self.is_distributed_enabled():
-        return super().tiled_encode(x)
-
-    self.clear_cache()
-    result = self.distributed_executor.execute(
-        x,
-        DistributedOperator(
-            split=self.encode_tile_split,
-            exec=self.encode_tile_exec,
-            merge=self.encode_tile_merge,
-        ),
-        broadcast_result=True,  # Latents needed by all ranks for diffusion
-    )
-    self.clear_cache()
-    return result
-```
-
-**Key differences from decode parallel:**
-
-| Aspect | Decode Parallel | Encode Parallel |
-|--------|-----------------|-----------------|
-| `broadcast_result` | Often `False` (only rank 0 needs output) | `True` (all ranks need latents for diffusion) |
-| Patchify | Applied in merge (unpatchify) | Handled by parent `_encode()` before `tiled_encode()` |
-| Temporal chunking | Frame-by-frame | Chunk-based (e.g., 1 + 4n frames) |
-
 ## Testing
 Verify numerical consistency between:
 + vae_patch_parallel_size = 1
@@ -440,20 +272,18 @@ When vae_patch_parallel_size is larger than the DiT world size, it will automati
 
 Complete examples in the codebase:
 
-| Model | Path | Decode Parallel | Encode Parallel |
+| Model | Path | Notes |
-|-------|------|-----------------|-----------------|
+|-------|------|-------|
-| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
+| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
-| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
+| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
-| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
+| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | Distributed AutoencoderKLQwenImage |
 
 ---
 
 ## Summary
 
-Adding VAE Patch Parallel support to diffusion model:
+Adding Vae Patch Parallel support to diffusion model:
 
-1. **Implement Distributed VAE** - Inherit from base VAE class and `DistributedVaeMixin`
+1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
-2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
+2. **Change vae model in pipeline to Distributed Vae**
-3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
+3. **Test** - Verify with `tensor_parallel_size=N` quality
-4. **Change VAE model in pipeline** - Use the distributed version
-5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`

docs/design/module/async_omni_architecture.md

@@ -69,7 +69,7 @@
 [5] Orchestrator._orchestration_loop (loop)
     -> poll stage output
        - llm stage: await get_output_async()
-       - diffusion stage: get_diffusion_output_nowait()
+       - diffusion stage: get_diffusion_output_async()
     -> (llm stage) output_processors[i].process_outputs(...)
     -> _route_output(...)
     -> if finished and not final_stage and non-async-chunk:
@@ -112,7 +112,7 @@ sequenceDiagram
     ORCH->>S0: add_request_async
 
     loop poll route forward
-        ORCH->>S0: get_output_async / get_diffusion_output_nowait
+        ORCH->>S0: get_output_async / get_diffusion_output_async
         ORCH->>ORCH: _route_output
         alt need forward to next stage
             ORCH->>SN: add_request_async

docs/getting_started/installation/README.md

@@ -6,5 +6,4 @@ vLLM-Omni supports the following hardware platforms:
     - [NVIDIA CUDA](gpu.md)
     - [AMD ROCm](gpu.md)
     - [Intel XPU](gpu.md)
-    - [MThreads MUSA](gpu.md)
 - [NPU](npu.md)

docs/getting_started/installation/gpu.md

@@ -22,10 +22,6 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:requirements"
 
-=== "MThreads MUSA"
-
-    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:requirements"
-
 ## Set up using Python
 
 ### Create a new Python environment
@@ -48,10 +44,6 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-wheels"
 
-=== "MThreads MUSA"
-
-    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-wheels"
-
 [](){ #build-from-source }
 
 ### Build wheel from source
@@ -68,10 +60,6 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-wheel-from-source"
 
-=== "MThreads MUSA"
-
-    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-wheel-from-source"
-
 ## Set up using Docker
 
 ### Pre-built images
@@ -88,10 +76,6 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:pre-built-images"
 
-=== "MThreads MUSA"
-
-    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:pre-built-images"
-
 ### Build your own docker image
 
 === "AMD ROCm"
@@ -101,7 +85,3 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
 === "Intel XPU"
 
     --8<-- "docs/getting_started/installation/gpu/xpu.inc.md:build-docker"
-
-=== "MThreads MUSA"
-
-    --8<-- "docs/getting_started/installation/gpu/musa.inc.md:build-docker"

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

@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
 # --8<-- [start:build-wheel-from-source]
 
 #### Installation of vLLM
-If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
+If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
 
 ```bash
-uv pip install vllm==0.19.0 --torch-backend=auto
+uv pip install vllm==0.18.0 --torch-backend=auto
 ```
 
-The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
+The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
 
 #### Installation of vLLM-Omni
 Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,12 +66,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr
 ```bash
 git clone https://github.com/vllm-project/vllm.git
 cd vllm
-git checkout v0.19.0
+git checkout v0.18.0
 ```
 Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
 ```bash
-#For CUDA 13.0
+export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
-export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
 ```
 Install vllm with command below (If you have no existing PyTorch).
 ```bash

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

@@ -1,65 +0,0 @@
-# --8<-- [start:requirements]
-
-- GPU: Moore Threads GPU with MUSA SDK installed (validated on MTT S5000)
-
-# --8<-- [end:requirements]
-# --8<-- [start:set-up-using-python]
-
-vLLM-Omni for MUSA requires building from source. Pre-built wheels are not currently available.
-
-!!! note
-    MUSA platform requires vLLM-MUSA to be installed first.
-
-# --8<-- [start:pre-built-wheels]
-
-# --8<-- [end:pre-built-wheels]
-
-# --8<-- [start:build-wheel-from-source]
-
-#### Prerequisites
-
-- **MUSA SDK**: Download from [MUSA SDK Download](https://developer.mthreads.com/sdk/download/musa)
-- **torchada**: CUDA→MUSA compatibility layer for PyTorch (`pip install torchada`)
-- **mthreads-ml-py**: MTML Python bindings (`pip install mthreads-ml-py`)
-- **MATE**: MUSA AI Tensor Engine ([GitHub](https://github.com/MooreThreads/mate))
-
-#### Installation of vLLM-MUSA
-
-```bash
-git clone https://github.com/MooreThreads/vllm-musa.git
-cd vllm-musa
-git checkout v0.18.0-dev
-pip install . --no-build-isolation -v
-```
-
-#### Installation of vLLM-Omni
-
-```bash
-git clone https://github.com/vllm-project/vllm-omni.git
-cd vllm-omni
-VLLM_OMNI_TARGET_DEVICE=musa pip install -e . --no-build-isolation
-```
-
-For Gradio demos:
-
-```bash
-pip install -e '.[demo]' --no-build-isolation
-```
-
-#### Environment Variables
-
-```bash
-export MUSA_VISIBLE_DEVICES=0,1
-export VLLM_WORKER_MULTIPROC_METHOD=spawn
-export VLLM_MUSA_CUSTOM_OP_USE_NATIVE=false
-```
-
-# --8<-- [end:build-wheel-from-source]
-
-# --8<-- [start:build-docker]
-
-# --8<-- [end:build-docker]
-
-# --8<-- [start:pre-built-images]
-
-# --8<-- [end:pre-built-images]

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

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

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

@@ -10,10 +10,10 @@ The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built
 ```bash
 # Update the vllm-ascend image
 # Atlas A2:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
 # Atlas A3:
-# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3
+# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1-a3
-export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
+export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
 docker run --rm \
     --name vllm-omni-npu \
     --shm-size=1g \
@@ -33,6 +33,17 @@ docker run --rm \
     -p 8000:8000 \
     -it $IMAGE bash
 
+cd /vllm-workspace/vllm
+git fetch origin --tags
+git checkout v0.18.0
+
+# Because vllm-ascend will release v0.18.0rc1 after vllm-omni 0.16.0,
+# we have to pin vllm-ascend at the current commit.
+cd /vllm-workspace/vllm-ascend
+git pull origin main
+git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
+pip install -v -e .
+
 # Inside the container, install vLLM-Omni from source
 cd /vllm-workspace
 git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
@@ -57,10 +68,15 @@ You can also build vLLM-Omni from the latest main branch if you want to use the
 
 ```bash
 # Pin vLLM version to 0.18.0
-git clone -b v0.18.0 https://github.com/vllm-project/vllm.git
+cd /vllm-workspace/vllm
-VLLM_TARGET_DEVICE=empty pip install -v -e .
+git fetch origin --tags
+git checkout v0.18.0
 
-git clone -b v0.18.0rc1 https://github.com/vllm-project/vllm-ascend.git
+# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
+cd /vllm-workspace/vllm-ascend
+git pull origin main
+git fetch origin --tags
+git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
 pip install -v -e .
 
 # Install vLLM-Omni from the latest main branch

docs/getting_started/quickstart.md

@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
 source .venv/bin/activate
 
 # On CUDA
-uv pip install vllm==0.19.0 --torch-backend=auto
+uv pip install vllm==0.18.0 --torch-backend=auto
 
 # On ROCm
-uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
+uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
 
 git clone https://github.com/vllm-project/vllm-omni.git
 cd vllm-omni

docs/mkdocs/hooks/generate_argparse.py

@@ -120,7 +120,6 @@ def extract_omni_serve_subparser_init():
                         "_FlexibleArgumentParser": _FlexibleArgumentParser,
                         "FlexibleArgumentParser": _FlexibleArgumentParser,
                         "make_arg_parser": lambda parser: parser,  # no-op for doc
-                        "_ensure_vllm_platform": lambda: None,  # no-op for doc
                         "VLLM_SUBCMD_PARSER_EPILOG": "",
                         "logger": logger,
                         "DummySubparsers": DummySubparsers,

docs/models/supported_models.md

@@ -31,13 +31,9 @@ th {
 | `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
-| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
 | `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
-| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
-| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
 | `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
-| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
 | `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
 | `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
 | `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
@@ -60,6 +56,5 @@ th {
 | `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
 | `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
 | `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
-|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
 
 ✅︎ indicates the model is supported on that backend. Empty cells mean not listed as supported on that backend.

docs/serving/speech_api.md

@@ -118,7 +118,6 @@ Content-Type: application/json
 | `instructions` | string | "" | Voice style/emotion instructions |
 | `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
 | `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
-| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
 
 **Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
 
@@ -144,23 +143,9 @@ Lists available voices for the loaded model.
 
 ```json
 {
-    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian", "custom_voice_1"],
+    "voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
-    "uploaded_voices": [
-        {
-            "name": "custom_voice_1",
-            "consent": "user_consent_id",
-            "created_at": 1738660000,
-            "file_size": 1024000,
-            "mime_type": "audio/wav",
-            "ref_text": "The exact transcript of the audio sample.",
-            "speaker_description": "warm narrator"
-        }
-    ]
 }
 ```
-
-`uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
-
 ```
 POST /v1/audio/voices
 Content-Type: multipart/form-data
@@ -176,7 +161,6 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 | `consent` | string | Yes | Consent recording ID |
 | `name` | string | Yes | Name for the new voice |
 | `ref_text` | string | No | Transcript of the audio. When provided, enables in-context voice cloning (higher quality). Without it, only the speaker embedding is extracted. |
-| `speaker_description` | string | No | Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata and returned in `GET /v1/audio/voices`. |
 
 **Response Example:**
 
@@ -188,15 +172,11 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
     "consent": "user_consent_id",
     "created_at": 1738660000,
     "mime_type": "audio/wav",
-    "file_size": 1024000,
+    "file_size": 1024000
-    "ref_text": "The exact transcript of the audio sample.",
-    "speaker_description": "warm narrator"
   }
 }
 ```
 
-Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
-
 **Usage Example:**
 
 ```bash
@@ -204,8 +184,7 @@ curl -X POST http://localhost:8091/v1/audio/voices \
   -F "audio_sample=@/path/to/voice_sample.wav" \
   -F "consent=user_consent_id" \
   -F "name=custom_voice_1" \
-  -F "ref_text=The exact transcript of the audio sample." \
+  -F "ref_text=The exact transcript of the audio sample."
-  -F "speaker_description=warm narrator"
 ```
 
 ## Streaming Text Input (WebSocket)

docs/user_guide/diffusion/cache_acceleration/cache_dit.md

@@ -164,7 +164,7 @@ cache_config={
 
 **Performance Tips**:
 
-- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., FLUX.2-klein) use a larger value for `Fn_compute_blocks` for a balanced performance.
+- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., [FLUX.2-klein](https://github.com/wtomin/vllm-omni/blob/main/vllm_omni/diffusion/cache/cache_dit_backend.py#L363)) use a larger value for `Fn_compute_blocks` for a balanced performance.
 - Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality.
 - Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed.
 

docs/user_guide/diffusion/cpu_offload_diffusion.md

@@ -91,19 +91,12 @@ Models must define the blocks attribute name for layerwise offloading:
 
 ```python
 class WanTransformer3DModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["blocks"]  # Attribute names containing transformer blocks
+    _layerwise_offload_blocks_attr = "blocks"  # Attribute name containing transformer blocks
 
     def __init__(self):
         self.blocks = nn.ModuleList([...])  # Transformer blocks
 ```
 
-For models with multiple block types:
-
-```python
-class Flux2Transformer2DModel(nn.Module):
-    _layerwise_offload_blocks_attrs = ["transformer_blocks", "single_transformer_blocks"]
-```
-
 ### Limitations
 - Cold start latency increases because of
     1) components are loaded to CPU first at the very first during initialization,
@@ -139,16 +132,12 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
 
 ## Supported Models
 
-| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attrs (Layerwise specific) |
+| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) |
-|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
+|--------------|----------------|-----------|---------------------|-------------------|-------------|
-| LongCatImagePipeline | `meituan-longcat/LongCat-Image` | `LongCatImageTransformer2DModel` | - | ✓ | `"transformer_blocks"`, `"single_transformer_blocks"` |
-| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
-| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
-| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
-| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
-| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
 | Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
+| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
 
 **Notes:**
 - Model-Level Offloading is expected to be supported by all common diffusion models (DiT and encoders) naturally
-- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attrs` pointing to transformer blocks
+- Layerwise Offloading requires DiT class to define `_layerwise_offload_blocks_attr` pointing to transformer blocks

docs/user_guide/diffusion/quantization/autoround.md

@@ -1,91 +0,0 @@
-# AutoRound Quantization
-
-## Overview
-
-[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
-
-The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
-
-### Supported Schemes
-
-| Scheme | Bits | Status |
-|--------|------|--------|
-| W4A16 | 4 | ✅ Supported |
-| W8A16 | 8 | Planned |
-
-W4A16 is the first supported scheme. Additional schemes will be added in future releases.
-
-## Configuration
-
-1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
-
-```python
-from vllm_omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
-
-outputs = omni.generate(
-    "A cat sitting on a windowsill",
-    OmniDiffusionSamplingParams(num_inference_steps=28),
-)
-outputs[0].save_images("output.png")
-```
-
-2. **CLI**: pass the quantized model path directly.
-
-```bash
-python examples/offline_inference/text_to_image/text_to_image.py \
-  --model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
-  --prompt "A cat sitting on a windowsill" \
-  --num-inference-steps 28 \
-  --output outputs/flux_w4a16.png
-```
-
-No `--quantization` flag is needed — the quantization method is read from the checkpoint.
-
-## How It Works
-
-The checkpoint's `config.json` contains:
-
-```json
-{
-  "quantization_config": {
-    "quant_method": "auto-round",
-    "bits": 4,
-    "group_size": 128,
-    "sym": true,
-    "packing_format": "auto_round:auto_gptq",
-    "block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
-  }
-}
-```
-
-At load time:
-
-1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
-2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
-3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
-
-## Supported Models
-
-| Model | HF Checkpoint | Scheme | Group Size | Backend |
-|-------|--------------|--------|------------|---------|
-| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
-
-## Creating a Quantized Checkpoint
-
-Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
-
-```bash
-# W4A16 (4-bit weight, 16-bit activation)
-auto-round \
-    --model black-forest-labs/FLUX.1-dev \
-    --scheme W4A16 \
-    --batch_size 1 \
-    --disable_opt_rtn \
-    --dataset coco2014 \
-    --iters 0
-```
-
-The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.

docs/user_guide/diffusion/quantization/overview.md

@@ -11,7 +11,6 @@ vLLM-Omni provides a unified quantization framework that supports both diffusion
 | FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
 | Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
 | GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
-| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |
 
 ### Multi-stage Omni Models (Pre-quantized Checkpoints)
 
@@ -103,9 +102,6 @@ config = build_quant_config("fp8")
 # Dict with parameters
 config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
 
-# AutoRound / INC (auto-detected from checkpoint, or explicit)
-config = build_quant_config("auto-round", bits=4, group_size=128)
-
 # Per-component dict
 config = build_quant_config({
     "transformer": {"method": "fp8"},

docs/user_guide/diffusion_features.md

@@ -12,7 +12,7 @@
 
 vLLM-Omni supports various advanced features for diffusion models:
 
-- Acceleration: **cache methods**, **parallelism methods**, **startup optimizations**
+- Acceleration: **cache methods**, **parallelism methods**
 - Memory optimization: **cpu offloading**, **quantization**
 - Extensions: **LoRA inference**
 - Execution modes: **step execution**
@@ -44,12 +44,6 @@ Parallelism methods distribute computation across GPUs without quality loss (mat
 | **[HSDP](diffusion/parallelism/hsdp.md)** | Weight sharding via FSDP2, redistributed on-demand at runtime | Very large models (14B+) on limited VRAM, combinable with SP |
 | **[Expert Parallelism](diffusion/parallelism/expert_parallel.md)** | Shards MoE expert MLP blocks across devices | MoE diffusion models (e.g., HunyuanImage3.0) |
 
-#### Startup Optimization
-
-| Method | Description | Best For |
-|--------|-------------|----------|
-| **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** | Loads safetensors shards in parallel using a thread pool | All diffusion models; reduces startup from minutes to seconds |
-
 **Note:** Some acceleration methods can be combined together for optimized performance. See [Feature Compatibility Table](#feature-compatibility) and [Feature Compatibility Tutorial](feature_compatibility.md) for detailed configuration examples.
 
 ### Memory Optimization
@@ -111,36 +105,34 @@ The following tables show which models support each feature:
 | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
 | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 | **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
-| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
+| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
-| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
 | **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
+| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
+| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ |
-| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
-| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
+| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
-| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ |
-| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
+| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
-| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
+| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ | ❌ |
 
 > Notes:
 > 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
-> 2. `Tongyi-MAI/Z-Image-Turbo` and `SII-GAIR/daVinci-MagiHuman-Base-1080p` are distilled models with minimal NFEs; CFG-Parallel is not necessary.
+> 2. `Tongyi-MAI/Z-Image-Turbo` is a distilled model with minimal NFEs; CFG-Parallel is not necessary.
 
 ### VideoGen
 
 | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
 |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
-| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
+| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
-| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
 | **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 | **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
-| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
+| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 
 ### AudioGen
@@ -185,59 +177,6 @@ The following tables show which models support each feature:
     6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
 
 
-## Multi-Thread Weight Loading
-
-Large diffusion models can take several minutes to load weights at startup (e.g., ~3 min for Qwen-Image, ~5 min for Wan2.2 I2V 14B). Multi-thread weight loading speeds up this process by loading safetensors shards in parallel using a thread pool instead of sequentially.
-
-This optimization is **enabled by default** with 4 threads. No configuration is needed for the default behavior.
-
-### Configuration
-
-| Parameter | CLI Flag | Default | Description |
-|-----------|----------|---------|-------------|
-| `enable_multithread_weight_load` | `--disable-multithread-weight-load` | `True` (enabled) | Pass the flag to disable multi-thread loading |
-| `num_weight_load_threads` | `--num-weight-load-threads` | `4` | Number of threads for parallel weight loading |
-
-!!! tip
-    The default of 4 threads balances speed and disk I/O contention. On fast NVMe storage you may benefit from more threads (e.g., 8). On HDD or network storage, the default of 4 avoids saturating I/O bandwidth.
-
-### Online Serving
-
-```bash
-# Default (multi-thread enabled, 4 threads)
-vllm serve Qwen/Qwen-Image --omni --port 8091
-
-# Custom thread count
-vllm serve Wan-AI/Wan2.2-I2V-A14B-Diffusers --omni --num-weight-load-threads 8
-
-# Disable multi-thread loading
-vllm serve Qwen/Qwen-Image --omni --disable-multithread-weight-load
-```
-
-### Offline Inference
-
-```python
-from vllm_omni import Omni
-
-# Default (multi-thread enabled, 4 threads)
-omni = Omni(model="Qwen/Qwen-Image")
-
-# Custom thread count
-omni = Omni(
-    model="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
-    num_weight_load_threads=8,
-)
-```
-
-### Benchmarks
-
-Measured on NVIDIA H800:
-
-| Model | Before | After | Speedup |
-|-------|--------|-------|---------|
-| **Qwen/Qwen-Image** (53.7 GiB) | 168s | 27s | **6.2x** |
-| **Wan-AI/Wan2.2-I2V-A14B-Diffusers** (64.5 GiB) | 283s | 56s | **5.1x** |
-
 ## Learn More
 
 **Cache Acceleration:**
@@ -263,10 +202,6 @@ Measured on NVIDIA H800:
 
 - **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
 
-**Startup Optimization:**
-
-- **[Multi-Thread Weight Loading](#multi-thread-weight-loading)** - Speed up model startup by loading safetensors shards in parallel
-
 **Advanced Topics:**
 
 - **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance

docs/user_guide/examples/online_serving/qwen3_tts.md

@@ -159,7 +159,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "Hello, how are you?",
-        "voice": "vivian",
+        "speaker": "vivian",
         "language": "English"
     }' --output output.wav
 
@@ -168,7 +168,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "I am so excited!",
-        "voice": "vivian",
+        "speaker": "vivian",
         "instructions": "Speak with great enthusiasm"
     }' --output excited.wav
 
@@ -185,7 +185,7 @@ client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
 
 response = client.audio.speech.create(
     model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
-    voice="vivian",
+    speaker="vivian",
     input="Hello, how are you?",
 )
 
@@ -201,7 +201,7 @@ response = httpx.post(
     "http://localhost:8091/v1/audio/speech",
     json={
         "input": "Hello, how are you?",
-        "voice": "vivian",
+        "speaker": "vivian",
         "language": "English",
     },
     timeout=300.0,
@@ -237,16 +237,12 @@ List all available voices/speakers from the loaded model, including both built-i
       "consent": "user_consent_id",
       "created_at": 1738660000,
       "file_size": 1024000,
-      "mime_type": "audio/wav",
+      "mime_type": "audio/wav"
-      "ref_text": "The exact transcript of the audio sample.",
-      "speaker_description": "warm narrator"
     }
   ]
 }
 ```
 
-Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
-
 #### POST /v1/audio/voices
 
 Upload a new voice sample for voice cloning in Base task TTS requests.
@@ -256,7 +252,6 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 - `consent` (required): Consent recording ID
 - `name` (required): Name for the new voice
 - `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
-- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata.
 
 **Response Example:**
 ```json
@@ -267,23 +262,18 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
     "consent": "user_consent_id",
     "created_at": 1738660000,
     "mime_type": "audio/wav",
-    "file_size": 1024000,
+    "file_size": 1024000
-    "ref_text": "The exact transcript of the audio sample.",
-    "speaker_description": "warm narrator"
   }
 }
 ```
 
-Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
-
 **Usage Example:**
 ```bash
-curl -X POST http://localhost:8091/v1/audio/voices \
+curl -X POST http://localhost:8000/v1/audio/voices \
   -F "audio_sample=@/path/to/voice_sample.wav" \
   -F "consent=user_consent_id" \
   -F "name=custom_voice_1" \
-  -F "ref_text=The exact transcript of the audio sample." \
+  -F "ref_text=The exact transcript of the audio sample."
-  -F "speaker_description=warm narrator"
 ```
 
 ### Endpoint
@@ -300,7 +290,7 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
 ```json
 {
     "input": "Text to synthesize",
-    "voice": "vivian",
+    "speaker": "vivian",
     "response_format": "wav",
     "task_type": "CustomVoice",
     "language": "Auto",
@@ -320,7 +310,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 
 ### Voice and language (summary)
 
-- **Speaker**: Use the `voice` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
+- **Speaker**: Use the `speaker` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
 - **Language**: Use the `language` field for the codec language tag (`Auto`, `Chinese`, `English`, etc.). Default is `Auto` for automatic detection.
 - **CustomVoice**: Requires a valid `voice` from the model’s speaker set. **VoiceDesign**: Use `instructions` to describe the voice. **Base**: Use `ref_audio` and `ref_text` for voice cloning.
 
@@ -332,7 +322,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
 | ----------------- | ------ | -------------- | ----------------------------------------------------------- |
 | `input`           | string | **required**   | Text to synthesize                                          |
 | `model`           | string | server's model | Model to use (optional, should match server if specified)   |
-| `voice`           | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
+| `speaker`         | string | "vivian"       | Speaker name (e.g., vivian, ryan, aiden)                    |
 | `response_format` | string | "wav"          | Audio format: wav, mp3, flac, pcm, aac, opus                |
 | `speed`           | float  | 1.0            | Playback speed (0.25-4.0, not supported with `stream=true`) |
 
@@ -367,7 +357,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
     -H "Content-Type: application/json" \
     -d '{
         "input": "Hello, how are you?",
-        "voice": "vivian",
+        "speaker": "vivian",
         "language": "English",
         "stream": true,
         "response_format": "pcm"

examples/offline_inference/bagel/end2end.py

@@ -2,10 +2,6 @@ import argparse
 import os
 
 from vllm_omni.inputs.data import OmniPromptType
-from vllm_omni.model_executor.stage_input_processors.bagel import (
-    GEN_THINK_SYSTEM_PROMPT,
-    VLM_THINK_SYSTEM_PROMPT,
-)
 
 
 def parse_args():
@@ -69,17 +65,6 @@ def parse_args():
         help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
     )
     parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
-    parser.add_argument(
-        "--cfg-interval",
-        type=float,
-        nargs=2,
-        default=None,
-        help="CFG interval [start, end] (default: pipeline default)",
-    )
-    parser.add_argument(
-        "--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel"
-    )
-    parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min")
     parser.add_argument(
         "--enable-diffusion-pipeline-profiler",
         action="store_true",
@@ -91,12 +76,6 @@ def parse_args():
         default=None,
         help="Quantization method (e.g. 'fp8').",
     )
-    parser.add_argument(
-        "--think",
-        action="store_true",
-        default=False,
-        help="Enable thinking mode: AR stage decodes <think>...</think> planning tokens before image generation.",
-    )
 
     args = parser.parse_args()
     return args
@@ -131,12 +110,8 @@ def main():
     from vllm_omni.entrypoints.omni import Omni
 
     omni_kwargs = {}
-    stage_configs_path = args.stage_configs_path
+    if args.stage_configs_path:
-    if args.think and stage_configs_path is None:
+        omni_kwargs["stage_configs_path"] = args.stage_configs_path
-        stage_configs_path = "vllm_omni/model_executor/stage_configs/bagel_think.yaml"
-        print(f"[Info] Think mode enabled, using stage config: {stage_configs_path}")
-    if stage_configs_path:
-        omni_kwargs["stage_configs_path"] = stage_configs_path
 
     omni_kwargs.update(
         {
@@ -161,8 +136,7 @@ def main():
             if not args.image_path or not os.path.exists(args.image_path):
                 raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
             loaded_image = Image.open(args.image_path).convert("RGB")
-            think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
-            final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>"
             prompt_dict = {
                 "prompt": final_prompt_text,
                 "multi_modal_data": {"img2img": loaded_image},
@@ -174,10 +148,7 @@ def main():
         elif args.modality == "img2text":
             if args.image_path:
                 loaded_image = Image.open(args.image_path).convert("RGB")
-                think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else ""
+                final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                final_prompt_text = (
-                    f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
-                )
                 prompt_dict = {
                     "prompt": final_prompt_text,
                     "multi_modal_data": {"image": loaded_image},
@@ -185,13 +156,11 @@ def main():
                 }
                 formatted_prompts.append(prompt_dict)
         elif args.modality == "text2text":
-            think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
-            final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>"
             prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
             formatted_prompts.append(prompt_dict)
         else:
-            think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
+            final_prompt_text = f"<|im_start|>{p}<|im_end|>"
-            final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>"
             prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
             if args.negative_prompt is not None:
                 prompt_dict["negative_prompt"] = args.negative_prompt
@@ -209,12 +178,6 @@ def main():
                 "cfg_text_scale": args.cfg_text_scale,
                 "cfg_img_scale": args.cfg_img_scale,
             }
-            if args.cfg_interval is not None:
-                extra["cfg_interval"] = tuple(args.cfg_interval)
-            if args.cfg_renorm_type is not None:
-                extra["cfg_renorm_type"] = args.cfg_renorm_type
-            if args.cfg_renorm_min is not None:
-                extra["cfg_renorm_min"] = args.cfg_renorm_min
             if args.negative_prompt is not None:
                 extra["negative_prompt"] = args.negative_prompt
             diffusion_params.extra_args = extra  # type: ignore
@@ -223,13 +186,6 @@ def main():
 
     img_idx = 0
     for req_output in omni_outputs:
-        if args.think:
-            ro = getattr(req_output, "request_output", None)
-            if ro and getattr(ro, "outputs", None):
-                txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
-                if txt:
-                    print(txt)
-
         images = getattr(req_output, "images", None)
 
         if not images:
@@ -238,7 +194,6 @@ def main():
         for j, img in enumerate(images):
             save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
             img.save(save_path)
-            print(f"[Output] Saved image to {save_path}")
         img_idx += 1
 
     print(omni_outputs)

examples/offline_inference/dynin_omni/README.md

@@ -1,110 +0,0 @@
-# Dynin-Omni Offline End2End Example
-
-This folder contains a unified offline inference entrypoint:
-
-- `end2end.py`
-
-## 1. Environment Setup
-
-Run from repository root:
-
-```bash
-cd <REPO_ROOT>
-```
-
-If needed, install this repo in editable mode:
-
-```bash
-pip install -e .
-```
-
-## 2. Extra Dependencies (EMOVA)
-
-Install the following packages for EMOVA-related components:
-
-```bash
-pip install \
-  "phonemizer==3.3.0" \
-  "Unidecode==1.4.0" \
-  "hydra-core==1.3.2" \
-  "pytorch-lightning==1.1.0" \
-  "wget==3.2" \
-  "wrapt==2.1.1" \
-  "onnx==1.20.1" \
-  "frozendict==2.4.7" \
-  "inflect==7.5.0" \
-  "braceexpand==0.1.7" \
-  "webdataset==1.0.2" \
-  "torch-stft==0.1.4" \
-  "editdistance==0.8.1"
-```
-
-## 3. Hardware and VRAM Requirements
-
-This example uses a 3-stage pipeline on one GPU by default
-([`dynin_omni.yaml`](../../../vllm_omni/model_executor/stage_configs/dynin_omni.yaml)):
-
-- Stage-0 (`token2text`): `gpu_memory_utilization: 0.5`
-- Stage-1 (`token2image`): `gpu_memory_utilization: 0.1`
-- Stage-2 (`token2audio`): `gpu_memory_utilization: 0.1`
-
-### Requested GPU Memory Budget from `gpu_memory_utilization`
-
-| Stage | Utilization | A100 80GB | H200 141GB |
-| :-- | :-- | :-- | :-- |
-| Stage-0 (token2text) | 0.5 | ~40.0 GB | ~70.5 GB |
-| Stage-1 (token2image) | 0.1 | ~8.0 GB | ~14.1 GB |
-| Stage-2 (token2audio) | 0.1 | ~8.0 GB | ~14.1 GB |
-| Total requested budget | 0.7 | ~56.0 GB | ~98.7 GB |
-
-### Observed Runtime Signal (from your log)
-
-- Stage-0 reported: `Model loading took 15.12 GiB memory` (weights footprint signal).
-- Stages 1/2 can still add runtime memory depending on task path and backend allocations.
-- Keep extra headroom for CUDA/PyTorch overhead and temporary allocations.
-
-### GPU Compatibility
-
-- Confirmed target GPUs for this setup: **NVIDIA H200**, **NVIDIA A100**.
-- CI/e2e coverage in this repo also includes CUDA **L4** markers for Dynin tests.
-
-## 4. End2End Run Examples
-
-```bash
-# t2t
-python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
-  --task t2t --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
-
-# i2t
-python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
-  --task i2t --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text "Please describe this image in detail."
-
-# s2t
-python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
-  --task s2t --model snu-aidas/Dynin-Omni --audio <AUDIO_PATH> --text "Transcribe the given audio."
-
-# t2i
-python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
-  --task t2i --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
-
-# v2t
-python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
-  --task v2t --model snu-aidas/Dynin-Omni --video <VIDEO_PATH> --text "Describe this video in detail."
-
-# i2i
-python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
-  --task i2i --model snu-aidas/Dynin-Omni --image <IMAGE_PATH> --text <INSTRUCTION_TEXT>
-
-# t2s
-python <REPO_ROOT>/examples/offline_inference/dynin_omni/end2end.py \
-  --task t2s --model snu-aidas/Dynin-Omni --text <INSTRUCTION_TEXT>
-```
-
-## 5. Notes
-
-- Outputs are saved under task-specific directories in `/tmp` by default.
-- You can override output path with `--output-dir`.
-- If you want to force local config resolution, pass `--dynin-config-path <PATH_TO_DYNIN_OMNI_YAML>`.
-- If you see the warning
-  `max_num_batched_tokens (32768) exceeds max_num_seqs * max_model_len (4096)`,
-  reduce `max_num_batched_tokens` in stage config (for example, `4096` in CI config).

examples/offline_inference/image_to_image/image_to_image.md

@@ -51,6 +51,5 @@ Key arguments:
 - `--vae-use-tiling`: enable VAE tiling for memory optimization.
 - `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
 - `--enable-cpu-offload`: enable CPU offloading for diffusion models.
-- `--strength`: **Z-Image only** - controls the denoising start timestep for I2I (default: 0.6). Range: [0.0, 1.0]. Lower values preserve more of the original image; higher values allow more creative changes.
 
 > ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.

examples/offline_inference/image_to_video/image_to_video.py

@@ -146,7 +146,7 @@ def parse_args() -> argparse.Namespace:
         "--audio-sample-rate",
         type=int,
         default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
+        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
     )
     parser.add_argument(
         "--cache-backend",
@@ -471,9 +471,15 @@ def main():
 
     video_array = _ensure_frame_list(video_array)
 
-    if audio is not None:
+    use_ltx2_export = is_ltx2
-        from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+    encode_video = None
+    if use_ltx2_export:
+        try:
+            from diffusers.pipelines.ltx2.export_utils import encode_video
+        except ImportError:
+            encode_video = None
 
+    if use_ltx2_export and encode_video is not None:
         if isinstance(video_array, list):
             frames_np = np.stack(video_array, axis=0)
         elif isinstance(video_array, np.ndarray):
@@ -484,24 +490,29 @@ def main():
         if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
             frames_np = frames_np[..., :3]
 
-        frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
+        frames_np = np.clip(frames_np, 0.0, 1.0)
+        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
+        video_tensor = torch.from_numpy(frames_u8)
 
-        audio_np = audio
+        audio_out = None
-        if isinstance(audio_np, list):
+        if audio is not None:
-            audio_np = audio_np[0] if audio_np else None
+            if isinstance(audio, list):
-        if isinstance(audio_np, torch.Tensor):
+                audio = audio[0] if audio else None
-            audio_np = audio_np.detach().cpu().float().numpy()
+            if isinstance(audio, np.ndarray):
-        if isinstance(audio_np, np.ndarray):
+                audio = torch.from_numpy(audio)
-            audio_np = np.squeeze(audio_np).astype(np.float32)
+            if isinstance(audio, torch.Tensor):
+                audio_out = audio
+                if audio_out.dim() > 1:
+                    audio_out = audio_out[0]
+                audio_out = audio_out.float().cpu()
 
-        video_bytes = mux_video_audio_bytes(
+        encode_video(
-            frames_u8,
+            video_tensor,
-            audio_np,
+            fps=fps,
-            fps=float(fps),
+            audio=audio_out,
-            audio_sample_rate=args.audio_sample_rate,
+            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
+            output_path=str(output_path),
         )
-        with open(str(output_path), "wb") as f:
-            f.write(video_bytes)
     else:
         export_to_video(video_array, str(output_path), fps=fps)
     print(f"Saved generated video to {output_path}")

examples/offline_inference/magi_human/README.md

@@ -1,72 +0,0 @@
-# MagiHuman Generation
-
-MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
-
-Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
-
-## Setup
-
-### Install MagiCompiler (recommended)
-
-MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
-
-```bash
-# Clone the repo
-git clone https://github.com/SandAI-org/MagiCompiler.git
-cd MagiCompiler
-
-# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
-sudo apt update && sudo apt install -y graphviz
-
-# Python dependencies
-pip install -r requirements.txt
-
-# Install MagiCompiler
-pip install .          # end users (recommended)
-# pip install -e .     # developers (editable install)
-```
-
-### Hardware requirements
-
-Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
-
-Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
-
-## Run Examples
-
-Get into the example folder:
-```bash
-cd examples/offline_inference/magi_human
-```
-
-### End-to-End Generation (Text to Video+Audio)
-
-Generate a video with synchronized speech natively generated by the model.
-
-```bash
-python end2end.py \
-    --model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
-    --prompt "A young woman with long, wavy golden blonde hair..." \
-    --tensor-parallel-size 4 \
-    --output output_magihuman.mp4
-```
-
-## Common Parameters
-
-| Parameter | Default | Description |
-|-----------|---------|-------------|
-| `--model` | *(Required)* | Local model path or HuggingFace ID |
-| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
-| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
-| `--height` | `256` | Initial resolution height |
-| `--width` | `448` | Initial resolution width |
-| `--num-inference-steps` | `8` | Denoising steps |
-| `--seed` | `52` | Random seed |
-| `--output` | `output_magihuman.mp4` | Output video with audio path |
-
-## Example materials
-
-??? abstract "end2end.py"
-    ``````py
-    --8<-- "examples/offline_inference/magi_human/end2end.py"
-    ``````

examples/offline_inference/magi_human/end2end.py

@@ -1,122 +0,0 @@
-import argparse
-
-from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-
-def parse_args():
-    parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
-    parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
-    parser.add_argument(
-        "--prompt",
-        type=str,
-        default="",
-        help="Text prompt containing visual description, dialogue, and background sound.",
-    )
-    parser.add_argument(
-        "--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
-    )
-    parser.add_argument(
-        "--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
-    )
-    parser.add_argument("--height", type=int, default=256, help="Video height.")
-    parser.add_argument("--width", type=int, default=448, help="Video width.")
-    parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
-    parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
-    return parser.parse_args()
-
-
-def main():
-    args = parse_args()
-
-    print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
-    omni = Omni(
-        model=args.model,
-        init_timeout=1200,
-        tensor_parallel_size=args.tensor_parallel_size,
-        devices=list(range(args.tensor_parallel_size)),
-    )
-
-    prompt = args.prompt
-    if not prompt:
-        prompt = (
-            "A young woman with long, wavy golden blonde hair and bright blue eyes, "
-            "wearing a fitted ivory silk blouse with a delicate lace collar, sits "
-            "stationary in front of a softly lit, blurred warm-toned interior. Her "
-            "overall disposition is warm, composed, and gently confident. The camera "
-            "holds a static medium close-up, framing her from the shoulders up, "
-            "with shallow depth of field keeping her face in sharp focus. Soft "
-            "directional key light falls from the upper left, casting a gentle "
-            "highlight along her cheekbone and nose bridge. She draws a quiet breath, "
-            "the levator labii superiors relaxing as her lips part. She speaks in "
-            "clear, warm, unhurried American English: "
-            "\"The most beautiful things in life aren't things at all — "
-            "they're moments, feelings, and the people who make you feel truly alive.\" "
-            "Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
-            "shapes each vowel with precision. A faint, genuine smile engages the "
-            "zygomaticus major, lifting her lip corners fractionally. Her brows rest "
-            "in a soft, neutral arch throughout. She maintains steady, forward-facing "
-            "eye contact. Head position remains level; no torso displacement occurs.\n\n"
-            "Dialogue:\n"
-            "<Young blonde woman, American English>: "
-            "\"The most beautiful things in life aren't things at all — "
-            "they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
-            "Background Sound:\n"
-            "<Soft, warm indoor ambience with a faint distant piano melody>"
-        )
-
-    sampling_params = OmniDiffusionSamplingParams(
-        height=args.height,
-        width=args.width,
-        num_inference_steps=args.num_inference_steps,
-        seed=args.seed,
-        extra_args={
-            "seconds": 5,
-            "sr_height": 1080,
-            "sr_width": 1920,
-            "sr_num_inference_steps": 5,
-        },
-    )
-
-    print(f"Generating with prompt: {prompt[:80]}...")
-    outputs = omni.generate(
-        prompts=[prompt],
-        sampling_params_list=[sampling_params],
-    )
-
-    print(f"Generation complete. Output type: {type(outputs)}")
-    if outputs:
-        first = outputs[0]
-
-        if hasattr(first, "images") and first.images:
-            video_frames = first.images[0]
-            print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
-
-            audio_waveform = None
-            mm = first.multimodal_output or {}
-            if mm:
-                audio_waveform = mm.get("audio")
-                if audio_waveform is not None:
-                    print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
-
-            output_fps = float(mm.get("fps", 25))
-            output_sr = int(mm.get("audio_sample_rate", 24000))
-            print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
-
-            video_bytes = mux_video_audio_bytes(
-                video_frames,
-                audio_waveform,
-                fps=output_fps,
-                audio_sample_rate=output_sr,
-            )
-            with open(args.output, "wb") as f:
-                f.write(video_bytes)
-            print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
-        print("SUCCESS: MagiHuman pipeline generation completed.")
-    else:
-        print("WARNING: No outputs returned.")
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/omnivoice/README.md

@@ -1,73 +0,0 @@
-# OmniVoice
-
-This directory contains an offline demo for running OmniVoice TTS models with vLLM Omni. It generates speech from text and saves WAV files locally.
-
-## Model Overview
-
-[OmniVoice](https://huggingface.co/k2-fsa/OmniVoice) is a zero-shot multilingual TTS model supporting 600+ languages. It uses a diffusion language model (Qwen3-0.6B backbone) with iterative masked unmasking to generate speech.
-
-Three inference modes are supported:
-
-- **Auto Voice**: Generate speech without any reference — the model picks a voice automatically.
-- **Voice Clone**: Clone a voice from a reference audio + transcription.
-- **Voice Design**: Control voice style via natural language instruction (e.g., "female, low pitch, british accent").
-
-## Setup
-
-Ensure the model is downloaded:
-
-```bash
-huggingface-cli download k2-fsa/OmniVoice
-```
-
-> **Note:** Voice cloning requires `transformers>=5.3.0` for `HiggsAudioV2TokenizerModel`. Auto voice and voice design modes work with `transformers>=4.57.0`.
-
-## Quick Start
-
-Auto voice (text only):
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice --text "Hello, this is a test."
-```
-
-Voice design (with style instruction):
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice \
-    --text "Hello, this is a test." \
-    --instruct "female, low pitch, british accent"
-```
-
-Voice clone (with reference audio):
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice \
-    --text "Hello, this is a test." \
-    --ref-audio ref.wav \
-    --ref-text "This is the reference transcription."
-```
-
-## Language Support
-
-Specify a language for improved quality:
-
-```bash
-python end2end.py --model k2-fsa/OmniVoice \
-    --text "你好，这是一个测试。" \
-    --lang zh
-```
-
-## Architecture
-
-OmniVoice uses a two-stage pipeline:
-
-- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text.
-- **Stage 1 (Decoder)**: HiggsAudioV2 RVQ quantizer + DAC acoustic decoder. Converts tokens to 24kHz waveform.
-
-Both stages use `GPUGenerationWorker` with `OmniGenerationScheduler`.
-
-## Notes
-
-- Output audio is saved to `output.wav` by default. Use `--output` to change the path.
-- The model estimates duration from text automatically via `RuleDurationEstimator`.
-- Use `--stage-init-timeout` to increase the stage initialization timeout for first-time model downloads.

examples/offline_inference/omnivoice/end2end.py

@@ -1,164 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""End-to-end OmniVoice TTS inference via vLLM-Omni.
-
-Supports:
-- Auto voice mode: text only → generated speech
-- Voice cloning mode: text + reference audio → cloned voice speech
-
-Usage:
-    # Auto voice
-    python end2end.py --model k2-fsa/OmniVoice --text "Hello world"
-
-    # Voice cloning
-    python end2end.py --model k2-fsa/OmniVoice --text "Hello" \
-        --ref-audio ref.wav --ref-text "reference transcription"
-"""
-
-import argparse
-import os
-
-import numpy as np
-import soundfile as sf
-
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-
-
-def run_e2e():
-    parser = argparse.ArgumentParser(description="OmniVoice E2E TTS inference")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="k2-fsa/OmniVoice",
-        help="Model name or path (HuggingFace or local)",
-    )
-    parser.add_argument(
-        "--stage-config",
-        type=str,
-        default="vllm_omni/model_executor/stage_configs/omnivoice.yaml",
-    )
-    parser.add_argument(
-        "--text",
-        type=str,
-        default="Hello, this is a test of the OmniVoice text to speech system.",
-    )
-    parser.add_argument(
-        "--ref-audio",
-        type=str,
-        default=None,
-        help="Reference audio for voice cloning (WAV file)",
-    )
-    parser.add_argument(
-        "--ref-text",
-        type=str,
-        default=None,
-        help="Transcription of reference audio",
-    )
-    parser.add_argument(
-        "--lang",
-        type=str,
-        default=None,
-        help="Language code (e.g., 'en', 'zh')",
-    )
-    parser.add_argument(
-        "--instruct",
-        type=str,
-        default=None,
-        help="Voice design instruction (e.g., 'female, low pitch, british accent')",
-    )
-    parser.add_argument(
-        "--output",
-        type=str,
-        default="output.wav",
-        help="Output audio file path",
-    )
-    parser.add_argument(
-        "--stage-init-timeout",
-        type=int,
-        default=600,
-        help="Stage initialization timeout in seconds",
-    )
-    args = parser.parse_args()
-
-    if not os.path.exists(args.stage_config):
-        raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
-
-    print(f"Initializing OmniVoice with model={args.model}")
-
-    omni = Omni(
-        model=args.model,
-        stage_configs_path=args.stage_config,
-        trust_remote_code=True,
-        log_stats=True,
-    )
-
-    print("Model initialized. Preparing inputs...")
-
-    # Build prompt
-    mm_processor_kwargs = {}
-    multi_modal_data = {}
-
-    if args.ref_audio:
-        if not os.path.exists(args.ref_audio):
-            raise FileNotFoundError(f"Reference audio not found: {args.ref_audio}")
-
-        import librosa
-
-        audio_signal, sr = librosa.load(args.ref_audio, sr=None)
-        multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr)
-        mm_processor_kwargs["ref_text"] = args.ref_text or ""
-        mm_processor_kwargs["sample_rate"] = sr
-
-    if args.lang:
-        mm_processor_kwargs["lang"] = args.lang
-    if args.instruct:
-        mm_processor_kwargs["instruct"] = args.instruct
-
-    prompts = {"prompt": args.text}
-    if multi_modal_data:
-        prompts["multi_modal_data"] = multi_modal_data
-    if mm_processor_kwargs:
-        prompts["mm_processor_kwargs"] = mm_processor_kwargs
-
-    sampling_params_list = [OmniDiffusionSamplingParams()]
-
-    print(f"Generating speech for: {args.text}")
-
-    outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
-
-    print(f"Received {len(outputs)} outputs.")
-    for i, output in enumerate(outputs):
-        try:
-            ro = output.request_output
-            if ro is None:
-                print("No request_output found.")
-                continue
-
-            mm = getattr(ro, "multimodal_output", None)
-            if not mm and ro.outputs:
-                mm = getattr(ro.outputs[0], "multimodal_output", None)
-
-            if mm:
-                print(f"Multimodal output keys: {mm.keys()}")
-                if "audio" in mm:
-                    audio_out = mm["audio"]
-                    sr = mm.get("sr", 24000)
-                    if isinstance(audio_out, np.ndarray):
-                        audio_np = audio_out
-                    else:
-                        audio_np = audio_out.cpu().numpy().squeeze()
-                    out_path = args.output if i == 0 else f"output_{i}.wav"
-                    sf.write(out_path, audio_np, sr)
-                    print(f"Saved audio to {out_path} ({sr}Hz, {len(audio_np) / sr:.2f}s)")
-            else:
-                print("No multimodal output found.")
-        except Exception as e:
-            print(f"Error inspecting output: {e}")
-
-    omni.close()
-    print("Done.")
-
-
-if __name__ == "__main__":
-    run_e2e()

examples/offline_inference/text_to_image/README.md

@@ -33,7 +33,6 @@ This folder provides several entrypoints for experimenting with text-to-image di
 | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 |
 | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 |
 | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) |
-| `HunyuanImage-3.0` | 1024 x 1024 | 80.0 (TP≥3)  | 160 |
 
 !!! info
 *Peak VRAM:  based on basic single-card usage, batch size =1, without any acceleration/optimization features. FLUX.2-dev requires `--enable-cpu-offload` on a single 80 GiB GPU.
@@ -91,8 +90,6 @@ python text_to_image.py \
 | `--enable-cpu-offload` | flag | off | Enable CPU offloading for diffusion models |
 | `--lora-path` | str | — | Path to PEFT LoRA adapter folder |
 | `--lora-scale` | float | `1.0` | Scale factor for LoRA weights |
-| `--use-system-prompt` | str | `None` | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text. Recommended: `en_unified`. Only for HunyuanImage-3.0.|
-| `--system-prompt` | str | `None` | Custom system prompt text. Only used when `--use-system-prompt` is set to `custom`. Only for HunyuanImage-3.0.|
 
 **NextStep-1.1 specific arguments:**
 

examples/offline_inference/text_to_image/text_to_image.py

@@ -242,19 +242,6 @@ def parse_args() -> argparse.Namespace:
         action="store_true",
         help="Enable logging of diffusion pipeline stats.",
     )
-    parser.add_argument(
-        "--use-system-prompt",
-        type=str,
-        default=None,
-        choices=["None", "dynamic", "en_vanilla", "en_recaption", "en_think_recaption", "en_unified", "custom"],
-        help="System prompt preset for generation. Recommended: en_unified.",
-    )
-    parser.add_argument(
-        "--system-prompt",
-        type=str,
-        default=None,
-        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
-    )
     return parser.parse_args()
 
 
@@ -376,7 +363,7 @@ def main():
         f"vae_patch_parallel_size={args.vae_patch_parallel_size}, "
         f"enable_expert_parallel={args.enable_expert_parallel}."
     )
-    print(f"  CPU offload: {args.enable_cpu_offload}; CPU Layerwise Offload: {args.enable_layerwise_offload}")
+    print(f"  CPU offload: {args.enable_cpu_offload}")
     print(f"  Image size: {args.width}x{args.height}")
     if args.lora_path:
         print(f"  LoRA: scale={args.lora_scale}")
@@ -395,13 +382,13 @@ def main():
         )
 
     generation_start = time.perf_counter()
+
     extra_args = {
         "timesteps_shift": args.timesteps_shift,
         "cfg_schedule": args.cfg_schedule,
         "use_norm": args.use_norm,
-        "use_system_prompt": args.use_system_prompt,
-        "system_prompt": args.system_prompt,
     }
+
     if lora_request:
         extra_args["lora_request"] = lora_request
         extra_args["lora_scale"] = args.lora_scale

examples/offline_inference/text_to_video/text_to_video.py

@@ -56,13 +56,8 @@ def parse_args() -> argparse.Namespace:
         "Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
         "hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
     )
-    parser.add_argument(
-        "--model-class-name",
-        default=None,
-        help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
-    )
     parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
+    parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).")
     parser.add_argument("--seed", type=int, default=42, help="Random seed.")
     parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
     parser.add_argument(
@@ -160,7 +155,7 @@ def parse_args() -> argparse.Namespace:
         "--audio-sample-rate",
         type=int,
         default=24000,
-        help="Sample rate for audio output when saved (default: 24000).",
+        help="Sample rate for audio output when saved (default: 24000 for LTX2).",
     )
     parser.add_argument(
         "--vae-patch-parallel-size",
@@ -190,7 +185,6 @@ def parse_args() -> argparse.Namespace:
 
 def main():
     args = parse_args()
-    model_class_name = args.model_class_name
 
     preset = _detect_preset(args.model)
     for key, default_val in preset.items():
@@ -235,7 +229,6 @@ def main():
         enable_cpu_offload=args.enable_cpu_offload,
         parallel_config=parallel_config,
         enforce_eager=args.enforce_eager,
-        model_class_name=model_class_name,
         cache_backend=args.cache_backend,
         cache_config=cache_config,
         enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
@@ -430,8 +423,17 @@ def main():
 
     video_array = _ensure_frame_list(video_array)
 
+    use_ltx2_export = False
+    if args.model and "ltx" in str(args.model).lower():
+        use_ltx2_export = True
     if audio is not None:
-        from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
+        use_ltx2_export = True
+
+    if use_ltx2_export:
+        try:
+            from diffusers.pipelines.ltx2.export_utils import encode_video
+        except ImportError:
+            raise ImportError("diffusers is required for LTX2 encode_video.")
 
         if isinstance(video_array, list):
             frames_np = np.stack(video_array, axis=0)
@@ -440,24 +442,28 @@ def main():
         else:
             frames_np = np.asarray(video_array)
 
-        frames_u8 = (np.clip(frames_np, 0.0, 1.0) * 255).round().clip(0, 255).astype("uint8")
+        frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
+        video_tensor = torch.from_numpy(frames_u8)
 
-        audio_np = audio
+        audio_out = None
-        if isinstance(audio_np, list):
+        if audio is not None:
-            audio_np = audio_np[0] if audio_np else None
+            if isinstance(audio, list):
-        if isinstance(audio_np, torch.Tensor):
+                audio = audio[0] if audio else None
-            audio_np = audio_np.detach().cpu().float().numpy()
+            if isinstance(audio, np.ndarray):
-        if isinstance(audio_np, np.ndarray):
+                audio = torch.from_numpy(audio)
-            audio_np = np.squeeze(audio_np).astype(np.float32)
+            if isinstance(audio, torch.Tensor):
+                audio_out = audio
+                if audio_out.dim() > 1:
+                    audio_out = audio_out[0]
+                audio_out = audio_out.float().cpu()
 
-        video_bytes = mux_video_audio_bytes(
+        encode_video(
-            frames_u8,
+            video_tensor,
-            audio_np,
+            fps=args.fps,
-            fps=float(args.fps),
+            audio=audio_out,
-            audio_sample_rate=args.audio_sample_rate,
+            audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
+            output_path=str(output_path),
         )
-        with open(str(output_path), "wb") as f:
-            f.write(video_bytes)
     else:
         export_to_video(video_array, str(output_path), fps=args.fps)
     print(f"Saved generated video to {output_path}")

examples/offline_inference/vace/vace_video_generation.md

@@ -1,88 +0,0 @@
-# VACE Video Generation
-
-[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
-
-| Model | Architecture | Model Weights (bf16) | HuggingFace |
-|-------|-------------|----------------------|-------------|
-| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
-| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
-
-## Text-to-Video (T2V)
-
-```bash
-python vace_video_generation.py \
-  --mode t2v \
-  --prompt "A sleek robot stands in a vast warehouse filled with boxes" \
-  --height 480 --width 832 --num-frames 81 \
-  --num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
-  --output t2v_output.mp4
-```
-
-## Image-to-Video (I2V)
-
-First frame is kept, remaining frames are generated:
-
-```bash
-python vace_video_generation.py \
-  --mode i2v \
-  --image astronaut.jpg \
-  --prompt "An astronaut emerging from a cracked egg on the moon" \
-  --height 480 --width 832 --num-frames 81 \
-  --output i2v_output.mp4
-```
-
-## First-Last-Frame Interpolation (FLF2V)
-
-```bash
-python vace_video_generation.py \
-  --mode flf2v \
-  --image first_frame.jpg --last-image last_frame.jpg \
-  --prompt "A bird takes off from a branch and lands on another" \
-  --height 512 --width 512 --num-frames 81 \
-  --output flf2v_output.mp4
-```
-
-## Inpainting
-
-Center vertical stripe is masked and regenerated:
-
-```bash
-python vace_video_generation.py \
-  --mode inpaint \
-  --image scene.jpg \
-  --prompt "Shrek walks out of a building" \
-  --height 480 --width 832 --num-frames 81 \
-  --output inpaint_output.mp4
-```
-
-## Reference Image-guided (R2V)
-
-```bash
-python vace_video_generation.py \
-  --mode r2v \
-  --image reference.jpg \
-  --prompt "Camera slowly zooms out from the character" \
-  --height 480 --width 832 --num-frames 81 \
-  --output r2v_output.mp4
-```
-
-## Key Arguments
-
-- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
-- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
-- `--image`: Input image for I2V, inpainting, and R2V modes.
-- `--last-image`: Last frame image for FLF2V mode.
-- `--prompt`: Text description of desired video.
-- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
-- `--num-frames`: Number of frames (default 81).
-- `--guidance-scale`: CFG scale (default 5.0).
-- `--flow-shift`: Scheduler flow shift (default 5.0).
-- `--num-inference-steps`: Number of denoising steps (default 30).
-- `--fps`: Frames per second for the saved MP4 (default 16).
-- `--output`: Path to save the generated video.
-- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
-- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
-- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
-- `--tensor-parallel-size`: Tensor parallel size.
-
-> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.

examples/offline_inference/vace/vace_video_generation.py

@@ -1,209 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""VACE video generation example.
-
-VACE (Video All-in-one Creation Engine) supports multiple video tasks:
-  - T2V:        Text-to-Video
-  - I2V:        Image-to-Video (first frame conditioning)
-  - V2LF:       Video-to-Last-Frame
-  - FLF2V:      First-Last-Frame interpolation
-  - Inpainting:  Masked region generation
-  - R2V:        Reference image-guided generation
-
-Usage examples:
-  # T2V (text-to-video)
-  python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
-
-  # I2V (image-to-video, first frame kept)
-  python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
-
-  # FLF2V (first-last frame interpolation)
-  python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
-
-  # R2V (reference image guided)
-  python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
-"""
-
-import argparse
-import time
-from pathlib import Path
-
-import numpy as np
-import PIL.Image
-import torch
-
-from vllm_omni.diffusion.data import DiffusionParallelConfig
-from vllm_omni.entrypoints.omni import Omni
-from vllm_omni.inputs.data import OmniDiffusionSamplingParams
-from vllm_omni.platforms import current_omni_platform
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="VACE video generation.")
-    parser.add_argument(
-        "--model",
-        default="Wan-AI/Wan2.1-VACE-14B-diffusers",
-        help="VACE model ID or local path.",
-    )
-    parser.add_argument(
-        "--mode",
-        default="t2v",
-        choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
-        help="Generation mode.",
-    )
-    parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
-    parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
-    parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
-    parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
-    parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
-    parser.add_argument("--seed", type=int, default=42, help="Random seed.")
-    parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
-    parser.add_argument("--height", type=int, default=480, help="Video height.")
-    parser.add_argument("--width", type=int, default=832, help="Video width.")
-    parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
-    parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
-    parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
-    parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
-    parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
-    parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
-    parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
-    parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
-    parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
-    parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
-    return parser.parse_args()
-
-
-def build_prompts(args):
-    """Build prompt dict with multi_modal_data based on mode."""
-    h, w, nf = args.height, args.width, args.num_frames
-
-    gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
-    mask_black = PIL.Image.new("L", (w, h), 0)
-    mask_white = PIL.Image.new("L", (w, h), 255)
-
-    prompt_data = {
-        "prompt": args.prompt,
-        "negative_prompt": args.negative_prompt,
-    }
-
-    if args.mode == "t2v":
-        return prompt_data
-
-    if args.mode == "r2v":
-        assert args.image, "--image required for R2V mode"
-        ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
-        return prompt_data
-
-    if args.mode == "i2v":
-        assert args.image, "--image required for I2V mode"
-        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {
-            "video": [img] + [gray] * (nf - 1),
-            "mask": [mask_black] + [mask_white] * (nf - 1),
-        }
-        return prompt_data
-
-    if args.mode == "v2lf":
-        assert args.image, "--image required for V2LF mode"
-        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {
-            "video": [gray] * (nf - 1) + [img],
-            "mask": [mask_white] * (nf - 1) + [mask_black],
-        }
-        return prompt_data
-
-    if args.mode == "flf2v":
-        assert args.image and args.last_image, "--image and --last-image required for FLF2V"
-        first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
-        prompt_data["multi_modal_data"] = {
-            "video": [first] + [gray] * (nf - 2) + [last],
-            "mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
-        }
-        return prompt_data
-
-    if args.mode == "inpaint":
-        assert args.image, "--image required for inpaint mode"
-        img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
-        d = 80
-        frames, masks = [], []
-        for _ in range(nf):
-            base = np.array(img).copy()
-            mask = PIL.Image.new("L", (w, h), 0)
-            stripe = PIL.Image.new("L", (2 * d, h), 255)
-            mask.paste(stripe, (w // 2 - d, 0))
-            base[np.array(mask) > 128] = 128
-            frames.append(PIL.Image.fromarray(base))
-            masks.append(mask)
-        prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
-        return prompt_data
-
-    raise ValueError(f"Unknown mode: {args.mode}")
-
-
-def main():
-    args = parse_args()
-    generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
-
-    parallel_config = DiffusionParallelConfig(
-        ulysses_degree=args.ulysses_degree,
-        ring_degree=args.ring_degree,
-        cfg_parallel_size=args.cfg_parallel_size,
-    )
-
-    omni = Omni(
-        model=args.model,
-        vae_use_tiling=args.vae_use_tiling,
-        flow_shift=args.flow_shift,
-        enforce_eager=args.enforce_eager,
-        parallel_config=parallel_config,
-    )
-
-    prompt_data = build_prompts(args)
-
-    print(f"\n{'=' * 60}")
-    print(f"VACE {args.mode.upper()} Generation")
-    print(f"  Model: {args.model}")
-    print(f"  Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
-    print(f"{'=' * 60}\n")
-
-    start = time.perf_counter()
-    outputs = omni.generate(
-        prompt_data,
-        OmniDiffusionSamplingParams(
-            height=args.height,
-            width=args.width,
-            num_frames=args.num_frames,
-            num_inference_steps=args.num_inference_steps,
-            guidance_scale=args.guidance_scale,
-            generator=generator,
-        ),
-    )
-    elapsed = time.perf_counter() - start
-
-    video = outputs[0].images
-    if isinstance(video, list):
-        video = video[0]
-    if isinstance(video, torch.Tensor):
-        video = video.cpu().numpy()
-    if video.ndim == 5:
-        video = video[0]
-    print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
-
-    output_path = Path(args.output)
-    output_path.parent.mkdir(parents=True, exist_ok=True)
-
-    from diffusers.utils import export_to_video
-
-    if np.issubdtype(video.dtype, np.integer):
-        video = video.astype(np.float32) / 255.0
-    export_to_video(list(video), str(output_path), fps=args.fps)
-    print(f"Saved to {output_path}")
-
-    omni.close()
-
-
-if __name__ == "__main__":
-    main()

examples/offline_inference/x_to_video_audio/download_dreamid_omni.py

@@ -82,6 +82,7 @@ def main(output_dir: str):
 
     data = {
         "_class_name": "DreamIDOmniPipeline",
+        "fusion": "DreamID-Omni/dreamid_omni.safetensors",
     }
 
     with open(os.path.join(output_dir, "model_index.json"), "w", encoding="utf-8") as f:
@@ -89,12 +90,6 @@ def main(output_dir: str):
 
     print(f"model_index.json created at {os.path.join(output_dir, 'model_index.json')}")
 
-    transformer_dir = os.path.join(output_dir, "transformer")
-    os.makedirs(transformer_dir, exist_ok=True)
-    with open(os.path.join(transformer_dir, "config.json"), "w", encoding="utf-8") as f:
-        json.dump({"fusion": "DreamID-Omni/dreamid_omni.safetensors"}, f)
-    print(f"transformer/config.json created at {os.path.join(transformer_dir, 'config.json')}")
-
     # now we download the dependency code
     download_dependency()
 

examples/offline_inference/x_to_video_audio/x_to_video_audio.md

@@ -24,9 +24,7 @@ dreamid_omni/
 │   ├── models_t5_umt5-xxl-enc-bf16.pth
 │   ├── Wan2.2_VAE.pth
 │
-├── model_index.json
+├── model_index.json # create by download_dreamid_omni.py
-└── transformer/
-    └── config.json   # create by download_dreamid_omni.py
 ```
 
 ### Run the Inference

examples/offline_inference/x_to_video_audio/x_to_video_audio.py

@@ -132,8 +132,8 @@ def main() -> None:
     if not outputs:
         raise RuntimeError("No output returned from DreamID-Omni.")
     output = outputs[0].request_output
-    generated_video = output.images[0][0]
+    generated_video = output[0].images[0][0]
-    generated_audio = output.images[0][1]
+    generated_audio = output[0].images[0][1]
     try:
         from dreamid_omni.utils.io_utils import save_video
     except Exception as e:

examples/online_serving/dynin_omni/README.md

@@ -1,97 +0,0 @@
-# Dynin-Omni Online Serving Example
-
-## Installation
-
-Please refer to [README.md](../../../README.md).
-
-## Launch the Server
-
-First, find the `transformers_modules` path:
-
-```bash
-python - <<'PY'
-from transformers.utils.hub import HF_MODULES_CACHE
-print(HF_MODULES_CACHE)
-PY
-```
-
-Then export it for both `PYTHONPATH` and `HF_MODULES_CACHE`:
-
-```bash
-export PYTHONPATH=<transformers_modules_path>:$PYTHONPATH
-export HF_MODULES_CACHE=<transformers_modules_path>
-```
-
-Run from repository root:
-
-```bash
-vllm-omni serve snu-aidas/Dynin-Omni \
-  --omni \
-  --port 8091 \
-  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
-```
-
-If `vllm-omni` is not in PATH, run:
-
-```bash
-PYTHONPATH="$(pwd)" python -m vllm_omni.entrypoints.cli.main serve snu-aidas/Dynin-Omni \
-  --omni \
-  --port 8091 \
-  --stage-configs-path "$(pwd)/vllm_omni/model_executor/stage_configs/dynin_omni.yaml"
-```
-
-Wait until the server logs show both `All stages initialized successfully` and
-`Application startup complete.` before sending requests.
-
-## Send Requests via Python Client
-
-Move to the example directory:
-
-```bash
-cd examples/online_serving/dynin_omni
-```
-
-### Text -> Image
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-  --query-type t2i \
-  --prompt "A realistic indoor living room with natural daylight."
-```
-
-### Image -> Image
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-  --query-type i2i \
-  --image-path ../../offline_inference/dynin_omni/data/image/sofa_under_water.jpg \
-  --prompt "Transform this surreal underwater setting into a realistic indoor living room while preserving the sofa layout."
-```
-
-### Text -> Speech
-
-```bash
-python openai_chat_completion_client_for_multimodal_generation.py \
-  --query-type t2s \
-  --prompt "Hello. This is Dynin-omni."
-```
-
-## CLI Arguments
-
-- `--query-type` (`t2i|t2s|i2i`)
-- `--model` (default: `snu-aidas/Dynin-Omni`)
-- `--host` / `--port` (OpenAI-compatible vLLM endpoint)
-- `--prompt` (custom text)
-- `--image-path` (required for `i2i`)
-- `--modalities` (optional output modalities override)
-- `--output-dir` (default: `/tmp/dynin_online_outputs`)
-
-## Notes
-
-- This client currently supports only `t2i`, `t2s`, and `i2i`.
-- `t2t` is intentionally not exposed in this online example.
-- This example intentionally uses the OpenAI-compatible chat completion endpoint.
-- Task routing for non-text outputs relies on Dynin task trigger tokens (`<|t2i|>`, `<|i2i|>`, `<|t2s|>`) injected by the client.
-- Outputs are saved under `/tmp/dynin_online_outputs` by default.
-- Dynin stage-0 warmup can take a while on first startup; do not send requests before startup completes.
-- Dynin itself can execute text-returning tasks such as `t2t`, `s2t`, `i2t`, and `v2t`, but this online serving example currently runs stage-0 in `generation` mode. In that path, the generation worker does not surface the final text as `output.text`, so OpenAI chat responses for those text-output tasks may complete internally but still return empty text.

examples/online_serving/dynin_omni/openai_chat_completion_client_for_multimodal_generation.py

@@ -1,342 +0,0 @@
-#!/usr/bin/env python3
-# SPDX-License-Identifier: Apache-2.0
-
-from __future__ import annotations
-
-import argparse
-import base64
-import json
-import mimetypes
-import os
-import time
-from pathlib import Path
-from typing import Any
-
-DEFAULT_MODEL = "snu-aidas/Dynin-Omni"
-DEFAULT_OUTPUT_DIR = "/tmp/dynin_online_outputs"
-
-QUERY_CHOICES = ("t2i", "t2s", "i2i")
-DEFAULT_PROMPT_BY_QUERY = {
-    "t2i": "A high quality detailed living room interior photo.",
-    "t2s": "Please read this sentence naturally: Hello from Dynin-Omni online serving.",
-    "i2i": "Transform this image into a realistic indoor living room while preserving layout.",
-}
-DEFAULT_MODALITIES_BY_QUERY = {
-    "t2i": ["image"],
-    "t2s": ["audio"],
-    "i2i": ["image"],
-}
-OFFLINE_PARITY_STAGE_COUNT = 3
-OFFLINE_PARITY_STAGE_SAMPLING = {
-    "max_tokens": 1,
-    "temperature": 0.0,
-    "top_p": 1.0,
-    "detokenize": False,
-}
-
-
-def _infer_mime_type(path: Path) -> str:
-    mime_type, _ = mimetypes.guess_type(str(path))
-    return mime_type or "application/octet-stream"
-
-
-def _encode_file_as_data_url(path: Path) -> str:
-    mime_type = _infer_mime_type(path)
-    raw = path.read_bytes()
-    encoded = base64.b64encode(raw).decode("utf-8")
-    return f"data:{mime_type};base64,{encoded}"
-
-
-def _to_image_url(path_or_url: str) -> str:
-    value = str(path_or_url)
-    if value.startswith(("http://", "https://", "data:image/")):
-        return value
-    path = Path(value).expanduser().resolve()
-    if not path.exists():
-        raise FileNotFoundError(f"Image file not found: {path}")
-    return _encode_file_as_data_url(path)
-
-
-def _build_user_content(query_type: str, prompt: str, image_path: str | None) -> list[dict[str, Any]]:
-    if query_type == "t2i":
-        return [{"type": "text", "text": f"<|t2i|> {prompt}"}]
-
-    if query_type == "t2s":
-        return [{"type": "text", "text": f"<|t2s|> {prompt}"}]
-
-    if query_type == "i2i":
-        if not image_path:
-            raise ValueError("--image-path is required for query type i2i")
-        return [
-            {"type": "text", "text": f"<|i2i|> {prompt}"},
-            {"type": "image_url", "image_url": {"url": _to_image_url(image_path)}},
-        ]
-
-    raise ValueError(f"Unsupported query_type: {query_type}")
-
-
-def _collect_text_from_content(content: Any) -> list[str]:
-    texts: list[str] = []
-    if isinstance(content, str):
-        stripped = content.strip()
-        if stripped:
-            texts.append(stripped)
-        return texts
-
-    if isinstance(content, dict):
-        for key in ("text", "content", "value", "output_text"):
-            text_value = content.get(key)
-            if isinstance(text_value, str) and text_value.strip():
-                texts.append(text_value.strip())
-        return texts
-
-    if isinstance(content, list):
-        for item in content:
-            texts.extend(_collect_text_from_content(item))
-        return texts
-
-    content_text = getattr(content, "text", None)
-    if isinstance(content_text, str) and content_text.strip():
-        texts.append(content_text.strip())
-    content_value = getattr(content, "content", None)
-    if isinstance(content_value, str) and content_value.strip():
-        texts.append(content_value.strip())
-    output_text = getattr(content, "output_text", None)
-    if isinstance(output_text, str) and output_text.strip():
-        texts.append(output_text.strip())
-    return texts
-
-
-def _extract_text_outputs(chat_completion: Any) -> list[str]:
-    texts: list[str] = []
-    for choice in getattr(chat_completion, "choices", []) or []:
-        message = getattr(choice, "message", None)
-        if message is None:
-            continue
-        content = getattr(message, "content", None)
-        texts.extend(_collect_text_from_content(content))
-        reasoning_content = getattr(message, "reasoning_content", None)
-        if isinstance(reasoning_content, str) and reasoning_content.strip():
-            texts.append(reasoning_content.strip())
-        choice_text = getattr(choice, "text", None)
-        if isinstance(choice_text, str) and choice_text.strip():
-            texts.append(choice_text.strip())
-    top_level_output_text = getattr(chat_completion, "output_text", None)
-    if isinstance(top_level_output_text, str) and top_level_output_text.strip():
-        texts.append(top_level_output_text.strip())
-    return texts
-
-
-def _extract_image_data_urls(chat_completion: Any) -> list[str]:
-    urls: list[str] = []
-    for choice in getattr(chat_completion, "choices", []) or []:
-        message = getattr(choice, "message", None)
-        if message is None:
-            continue
-        content = getattr(message, "content", None)
-        if not isinstance(content, list):
-            continue
-        for item in content:
-            if not isinstance(item, dict):
-                continue
-            if item.get("type") != "image_url":
-                continue
-            image_url = (item.get("image_url") or {}).get("url")
-            if isinstance(image_url, str) and image_url.startswith("data:image"):
-                urls.append(image_url)
-    return urls
-
-
-def _extract_audio_payloads(chat_completion: Any) -> list[bytes]:
-    payloads: list[bytes] = []
-    for choice in getattr(chat_completion, "choices", []) or []:
-        message = getattr(choice, "message", None)
-        if message is None:
-            continue
-        message_audio = getattr(message, "audio", None)
-        if message_audio is None:
-            continue
-        data_b64 = getattr(message_audio, "data", None)
-        if isinstance(data_b64, str) and data_b64:
-            try:
-                payloads.append(base64.b64decode(data_b64))
-            except Exception:
-                continue
-    return payloads
-
-
-def _decode_data_url(data_url: str) -> tuple[bytes, str]:
-    header, data = data_url.split(",", 1)
-    mime_type = "image/png"
-    if ";" in header and ":" in header:
-        mime_type = header.split(":", 1)[1].split(";", 1)[0]
-    return base64.b64decode(data), mime_type
-
-
-def _image_extension_from_mime(mime_type: str) -> str:
-    if mime_type == "image/jpeg":
-        return ".jpg"
-    if mime_type == "image/webp":
-        return ".webp"
-    if mime_type == "image/gif":
-        return ".gif"
-    return ".png"
-
-
-def _save_outputs(
-    *,
-    query_type: str,
-    chat_completion: Any,
-    output_dir: Path,
-) -> None:
-    output_dir.mkdir(parents=True, exist_ok=True)
-    stamp = time.strftime("%Y%m%d_%H%M%S")
-
-    text_outputs = _extract_text_outputs(chat_completion)
-    image_data_urls = _extract_image_data_urls(chat_completion)
-    audio_payloads = _extract_audio_payloads(chat_completion)
-
-    if text_outputs:
-        text_path = output_dir / f"{query_type}_{stamp}.txt"
-        text_path.write_text("\n\n".join(text_outputs) + "\n", encoding="utf-8")
-        print(f"[dynin-online] text saved: {text_path}")
-        print(text_outputs[0])
-
-    for idx, image_url in enumerate(image_data_urls):
-        image_bytes, mime_type = _decode_data_url(image_url)
-        ext = _image_extension_from_mime(mime_type)
-        image_path = output_dir / f"{query_type}_{stamp}_{idx}{ext}"
-        image_path.write_bytes(image_bytes)
-        print(f"[dynin-online] image saved: {image_path}")
-
-    for idx, audio_bytes in enumerate(audio_payloads):
-        audio_path = output_dir / f"{query_type}_{stamp}_{idx}.wav"
-        audio_path.write_bytes(audio_bytes)
-        print(f"[dynin-online] audio saved: {audio_path}")
-
-    if not text_outputs and not image_data_urls and not audio_payloads:
-        print("[dynin-online] no output extracted from response")
-        raw_path = output_dir / f"{query_type}_{stamp}_raw_response.json"
-        try:
-            if hasattr(chat_completion, "model_dump_json"):
-                serialized = chat_completion.model_dump_json(indent=2)
-            else:
-                if hasattr(chat_completion, "model_dump"):
-                    raw_payload: Any = chat_completion.model_dump(mode="json")
-                else:
-                    raw_payload = chat_completion
-                try:
-                    serialized = json.dumps(raw_payload, ensure_ascii=False, indent=2)
-                except Exception:
-                    serialized = json.dumps({"repr": repr(raw_payload)}, ensure_ascii=False, indent=2)
-            raw_path.write_text(serialized + "\n", encoding="utf-8")
-            print(f"[dynin-online] raw response saved: {raw_path}")
-        except Exception:
-            pass
-
-
-def _build_offline_parity_sampling_params_list() -> list[dict[str, Any]]:
-    return [dict(OFFLINE_PARITY_STAGE_SAMPLING) for _ in range(OFFLINE_PARITY_STAGE_COUNT)]
-
-
-def run_request(args: argparse.Namespace) -> None:
-    from openai import OpenAI
-
-    client = OpenAI(
-        api_key="EMPTY",
-        base_url=f"http://{args.host}:{args.port}/v1",
-    )
-    prompt = args.prompt.strip() if args.prompt else DEFAULT_PROMPT_BY_QUERY[args.query_type]
-    user_content = _build_user_content(
-        query_type=args.query_type,
-        prompt=prompt,
-        image_path=args.image_path,
-    )
-    if args.modalities:
-        modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
-    else:
-        modalities = DEFAULT_MODALITIES_BY_QUERY[args.query_type]
-
-    extra_body = {
-        "sampling_params_list": _build_offline_parity_sampling_params_list(),
-    }
-    chat_completion = client.chat.completions.create(
-        model=args.model,
-        messages=[{"role": "user", "content": user_content}],
-        modalities=modalities,
-        extra_body=extra_body,
-    )
-    _save_outputs(
-        query_type=args.query_type,
-        chat_completion=chat_completion,
-        output_dir=Path(args.output_dir).expanduser(),
-    )
-
-
-def parse_args() -> argparse.Namespace:
-    parser = argparse.ArgumentParser(description="Dynin-Omni online chat completion client")
-    parser.add_argument(
-        "--query-type",
-        "-q",
-        type=str,
-        default="t2i",
-        choices=QUERY_CHOICES,
-        help="Dynin query type",
-    )
-    parser.add_argument(
-        "--model",
-        "-m",
-        type=str,
-        default=DEFAULT_MODEL,
-        help="Model name/path",
-    )
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="localhost",
-        help="Host/IP of the vLLM Omni API server",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=8091,
-        help="Port of the vLLM Omni API server",
-    )
-    parser.add_argument(
-        "--prompt",
-        "-p",
-        type=str,
-        default="",
-        help="Custom prompt text",
-    )
-    parser.add_argument(
-        "--image-path",
-        "-i",
-        type=str,
-        default=None,
-        help="Image path/URL for i2i",
-    )
-    parser.add_argument(
-        "--modalities",
-        type=str,
-        default="",
-        help="Comma-separated output modalities override (e.g., text,image,audio)",
-    )
-    parser.add_argument(
-        "--output-dir",
-        "-o",
-        type=str,
-        default=DEFAULT_OUTPUT_DIR,
-        help="Directory to save outputs",
-    )
-    return parser.parse_args()
-
-
-def main() -> None:
-    args = parse_args()
-    os.environ.setdefault("HF_HUB_DISABLE_XET", "1")
-    run_request(args)
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/image_to_image/README.md

@@ -314,7 +314,6 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `seed`                   | int   | None    | Random seed (reproducible)            |
 | `negative_prompt`        | str   | None    | Negative prompt                       |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate          |
-| `strength`               | float | 0.6     | **Z-Image only** - Denoising start timestep for I2I. Range: [0.0, 1.0]. Lower preserves more of original image. |
 | `layers`                 | int   | 4       | Number of layers (Qwen-Image-Layered) |
 | `resolution`             | int   | 640     | Resolution, 640 or 1024 (Qwen-Image-Layered) |
 

examples/online_serving/omnivoice/README.md

@@ -1,131 +0,0 @@
-# OmniVoice
-
-## Model Overview
-
-| Model | Description |
-|-------|-------------|
-| `k2-fsa/OmniVoice` | Zero-shot multilingual TTS (600+ languages) with diffusion language model (Qwen3-0.6B backbone) |
-
-> **Note:** Requires `transformers>=5.3.0` for voice cloning (HiggsAudioV2 tokenizer). Auto voice and voice design work with `transformers>=4.57.0`.
-
-## Launch the Server
-
-```bash
-vllm serve k2-fsa/OmniVoice \
-    --omni \
-    --port 8091 \
-    --trust-remote-code
-```
-
-Or use the convenience script:
-
-```bash
-./run_server.sh
-```
-
-## Send TTS Request
-
-### Using curl
-
-```bash
-# Basic TTS (auto voice)
-curl -X POST http://localhost:8091/v1/audio/speech \
-    -H "Content-Type: application/json" \
-    -d '{
-        "input": "Hello, how are you?",
-        "voice": "default",
-        "response_format": "wav"
-    }' --output output.wav
-```
-
-### Using Python
-
-```python
-import httpx
-
-response = httpx.post(
-    "http://localhost:8091/v1/audio/speech",
-    json={
-        "input": "Hello, how are you?",
-        "voice": "default",
-        "response_format": "wav",
-    },
-    timeout=300.0,
-)
-
-with open("output.wav", "wb") as f:
-    f.write(response.content)
-```
-
-### Using OpenAI SDK
-
-```python
-from openai import OpenAI
-
-client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
-
-response = client.audio.speech.create(
-    model="k2-fsa/OmniVoice",
-    voice="default",
-    input="Hello, how are you?",
-)
-
-response.stream_to_file("output.wav")
-```
-
-### Using the CLI Client
-
-```bash
-cd examples/online_serving/omnivoice
-
-# Basic TTS
-python speech_client.py --text "Hello, how are you?"
-
-# Specify language for improved quality
-python speech_client.py --text "Bonjour, comment allez-vous?" --language French
-```
-
-The CLI client supports:
-
-- `--api-base`: API base URL (default: `http://localhost:8091`)
-- `--model` (or `-m`): Model name (default: `k2-fsa/OmniVoice`)
-- `--text`: Text to synthesize (required)
-- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
-- `--language`: Language hint (default: Auto)
-- `--output` (or `-o`): Output file path (default: `omnivoice_output.wav`)
-
-## Inference Modes
-
-OmniVoice supports three inference modes. Currently, **auto voice** is supported through the online Speech API. Voice cloning and voice design are available via offline inference (see `examples/offline_inference/omnivoice/`).
-
-| Mode | Description | Online API | Offline |
-|------|-------------|:----------:|:-------:|
-| Auto Voice | Generate speech without reference | Yes | Yes |
-| Voice Clone | Clone from reference audio + transcript | - | Yes |
-| Voice Design | Control style via natural language instruction | - | Yes |
-
-## Architecture
-
-OmniVoice uses a single-stage diffusion pipeline:
-
-- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative masked unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text, then decodes to 24kHz waveform via HiggsAudioV2 RVQ quantizer + DAC acoustic decoder.
-
-## API Parameters
-
-OmniVoice uses the standard `/v1/audio/speech` endpoint. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full documentation.
-
-Key parameters:
-
-| Parameter | Type | Default | Description |
-|-----------|------|---------|-------------|
-| `input` | string | **required** | Text to synthesize |
-| `voice` | string | "default" | Voice name |
-| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
-| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
-
-## Troubleshooting
-
-1. **TTS model did not produce audio output**: Ensure the model is fully downloaded (`huggingface-cli download k2-fsa/OmniVoice`)
-2. **Connection refused**: Make sure the server is running on the correct port
-3. **Out of memory**: Reduce `--gpu-memory-utilization` (default stage config uses 0.5)
-4. **Slow first request**: The model performs warmup on first inference; subsequent requests are faster

examples/online_serving/omnivoice/run_server.sh

@@ -1,19 +0,0 @@
-#!/bin/bash
-# Launch vLLM-Omni server for OmniVoice TTS
-#
-# Usage:
-#   ./run_server.sh
-#   CUDA_VISIBLE_DEVICES=0 ./run_server.sh
-
-set -e
-
-MODEL="${MODEL:-k2-fsa/OmniVoice}"
-PORT="${PORT:-8091}"
-
-echo "Starting OmniVoice server with model: $MODEL"
-
-vllm serve "$MODEL" \
-    --host 0.0.0.0 \
-    --port "$PORT" \
-    --trust-remote-code \
-    --omni

examples/online_serving/omnivoice/speech_client.py

@@ -1,84 +0,0 @@
-"""Client for OmniVoice TTS via /v1/audio/speech endpoint.
-
-Examples:
-    # Basic TTS (auto voice)
-    python speech_client.py --text "Hello, how are you?"
-
-    # Specify language
-    python speech_client.py --text "Bonjour, comment allez-vous?" --language French
-"""
-
-import argparse
-
-import httpx
-
-DEFAULT_API_BASE = "http://localhost:8091"
-DEFAULT_API_KEY = "EMPTY"
-
-
-def run_tts(args) -> None:
-    """Generate speech via /v1/audio/speech API."""
-    payload = {
-        "model": args.model,
-        "input": args.text,
-        "voice": "default",
-        "response_format": args.response_format,
-    }
-
-    if args.language:
-        payload["language"] = args.language
-
-    print(f"Model: {args.model}")
-    print(f"Text: {args.text}")
-    if args.language:
-        print(f"Language: {args.language}")
-    print("Generating audio...")
-
-    api_url = f"{args.api_base}/v1/audio/speech"
-    headers = {
-        "Content-Type": "application/json",
-        "Authorization": f"Bearer {args.api_key}",
-    }
-
-    with httpx.Client(timeout=300.0) as client:
-        response = client.post(api_url, json=payload, headers=headers)
-
-    if response.status_code != 200:
-        print(f"Error: {response.status_code}")
-        print(response.text)
-        return
-
-    try:
-        text = response.content.decode("utf-8")
-        if text.startswith('{"error"'):
-            print(f"Error: {text}")
-            return
-    except UnicodeDecodeError:
-        pass
-
-    output_path = args.output or "omnivoice_output.wav"
-    with open(output_path, "wb") as f:
-        f.write(response.content)
-    print(f"Audio saved to: {output_path}")
-
-
-def main():
-    parser = argparse.ArgumentParser(description="OmniVoice TTS client")
-    parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
-    parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
-    parser.add_argument("--model", "-m", default="k2-fsa/OmniVoice", help="Model name")
-    parser.add_argument("--text", required=True, help="Text to synthesize")
-    parser.add_argument("--language", default=None, help="Language hint (e.g., English, Chinese, French)")
-    parser.add_argument(
-        "--response-format",
-        default="wav",
-        choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
-        help="Audio format (default: wav)",
-    )
-    parser.add_argument("--output", "-o", default=None, help="Output file path")
-    args = parser.parse_args()
-    run_tts(args)
-
-
-if __name__ == "__main__":
-    main()

examples/online_serving/qwen3_omni/README.md

@@ -36,45 +36,6 @@ cd examples/online_serving/qwen3_omni
 python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type use_image --port 8091 --host "localhost"
 ```
 
-#### Realtime WebSocket client (`openai_realtime_client.py`)
-
-[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, uploads a local audio file as **PCM16 mono @ 16 kHz** chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and prints **streaming transcription** (`transcription.delta` / `transcription.done`).
-
-**Dependencies:**
-
-```bash
-pip install websockets librosa numpy
-```
-
-(ffmpeg may be required by `librosa` for some formats; see the FAQ below.)
-
-**From this directory** (`examples/online_serving/qwen3_omni`):
-
-```bash
-python openai_realtime_client.py \
-  --host localhost \
-  --port 8091 \
-  --model Qwen/Qwen3-Omni-30B-A3B-Instruct \
-  --audio_path /path/to/your.wav
-```
-
-If `--audio_path` is omitted, the script uses a bundled default clip (`mary_had_lamb` via vLLM assets).
-
-**Arguments:**
-
-| Flag | Default | Description |
-|------|---------|-------------|
-| `--host` | `localhost` | API server host |
-| `--port` | `8000` | API server port (match your `vllm serve` port, e.g. `8091`) |
-| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (also sent in `session.update`) |
-| `--audio_path` | *(optional)* | Path to input audio; resampled to 16 kHz mono inside the client |
-
-Ensure the vLLM-Omni server is running with realtime support for this endpoint, for example:
-
-```bash
-vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
-```
-
 The Python client supports the following command-line arguments:
 
 - `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video`

examples/online_serving/qwen3_omni/openai_realtime_client.py

@@ -1,146 +0,0 @@
-"""
-This script demonstrates how to use the vLLM-Omni Realtime WebSocket API to perform
-audio transcription by uploading an audio file.
-
-Before running this script, you must start the vLLM-Omni server with a realtime-capable
-model, for example:
-
-    vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni
-
-Requirements:
-- vllm with audio support
-- websockets
-- librosa
-- numpy
-
-The script:
-1. Connects to the Realtime WebSocket endpoint
-2. Converts an audio file to PCM16 @ 16kHz
-3. Sends audio chunks to the server
-4. Receives and prints transcription as it streams
-"""
-
-import argparse
-import asyncio
-import base64
-import json
-
-import librosa
-import numpy as np
-import websockets
-from vllm.assets.audio import AudioAsset
-
-
-def audio_to_pcm16_base64(audio_path: str) -> str:
-    """
-    Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
-    """
-    # Load audio and resample to 16kHz mono
-    audio, _ = librosa.load(audio_path, sr=16000, mono=True)
-    # Convert to PCM16
-    pcm16 = (audio * 32767).astype(np.int16)
-    # Encode as base64
-    return base64.b64encode(pcm16.tobytes()).decode("utf-8")
-
-
-async def realtime_transcribe(audio_path: str, host: str, port: int, model: str):
-    """
-    Connect to the Realtime API and transcribe an audio file.
-    """
-    uri = f"ws://{host}:{port}/v1/realtime"
-
-    async with websockets.connect(uri) as ws:
-        # Wait for session.created
-        response = json.loads(await ws.recv())
-        if response["type"] == "session.created":
-            print(f"Session created: {response['id']}")
-        else:
-            print(f"Unexpected response: {response}")
-            return
-
-        # Validate model
-        await ws.send(json.dumps({"type": "session.update", "model": model}))
-
-        # Signal ready to start
-        await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
-
-        # Convert audio file to base64 PCM16
-        print(f"Loading audio from: {audio_path}")
-        audio_base64 = audio_to_pcm16_base64(audio_path)
-
-        # Send audio in chunks (4KB of raw audio = ~8KB base64)
-        chunk_size = 4096
-        audio_bytes = base64.b64decode(audio_base64)
-        total_chunks = (len(audio_bytes) + chunk_size - 1) // chunk_size
-
-        print(f"Sending {total_chunks} audio chunks...")
-        for i in range(0, len(audio_bytes), chunk_size):
-            chunk = audio_bytes[i : i + chunk_size]
-            await ws.send(
-                json.dumps(
-                    {
-                        "type": "input_audio_buffer.append",
-                        "audio": base64.b64encode(chunk).decode("utf-8"),
-                    }
-                )
-            )
-
-        # Signal all audio is sent
-        await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
-        print("Audio sent. Waiting for transcription...\n")
-
-        # Receive transcription
-        print("Transcription: ", end="", flush=True)
-        while True:
-            response = json.loads(await ws.recv())
-            if response["type"] == "transcription.delta":
-                print(response["delta"], end="", flush=True)
-            elif response["type"] == "transcription.done":
-                print(f"\n\nFinal transcription: {response['text']}")
-                if response.get("usage"):
-                    print(f"Usage: {response['usage']}")
-                break
-            elif response["type"] == "error":
-                print(f"\nError: {response['error']}")
-                break
-
-
-def main(args):
-    if args.audio_path:
-        audio_path = args.audio_path
-    else:
-        # Use default audio asset
-        audio_path = str(AudioAsset("mary_had_lamb").get_local_path())
-        print(f"No audio path provided, using default: {audio_path}")
-
-    asyncio.run(realtime_transcribe(audio_path, args.host, args.port, args.model))
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(description="Realtime WebSocket Transcription Client")
-    parser.add_argument(
-        "--model",
-        type=str,
-        default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
-        help="Model that is served and should be pinged.",
-    )
-    parser.add_argument(
-        "--audio_path",
-        type=str,
-        default=None,
-        help="Path to the audio file to transcribe.",
-    )
-    parser.add_argument(
-        "--host",
-        type=str,
-        default="localhost",
-        help="vLLM-Omni server host (default: localhost)",
-    )
-    parser.add_argument(
-        "--port",
-        type=int,
-        default=8000,
-        help="vLLM-Omni server port (default: 8000)",
-    )
-    args = parser.parse_args()
-    main(args)

examples/online_serving/qwen3_tts/README.md

@@ -233,7 +233,6 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
 - `consent` (required): Consent recording ID
 - `name` (required): Name for the new voice
 - `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
-- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter").
 
 **Response Example:**
 ```json
@@ -244,23 +243,18 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
     "consent": "user_consent_id",
     "created_at": 1738660000,
     "mime_type": "audio/wav",
-    "file_size": 1024000,
+    "file_size": 1024000
-    "ref_text": "The exact transcript of the audio sample.",
-    "speaker_description": "warm narrator"
   }
 }
 ```
 
-Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
-
 **Usage Example:**
 ```bash
 curl -X POST http://localhost:8000/v1/audio/voices \
   -F "audio_sample=@/path/to/voice_sample.wav" \
   -F "consent=user_consent_id" \
   -F "name=custom_voice_1" \
-  -F "ref_text=The exact transcript of the audio sample." \
+  -F "ref_text=The exact transcript of the audio sample."
-  -F "speaker_description=warm narrator"
 ```
 
 ### Endpoint

examples/online_serving/qwen3_tts/openai_speech_client.py

@@ -71,7 +71,7 @@ def run_tts_generation(args) -> None:
     payload = {
         "model": args.model,
         "input": args.text,
-        "voice": args.speaker,
+        "speaker": args.speaker,
         "response_format": args.response_format,
     }
 

examples/online_serving/text_to_image/README.md

@@ -231,8 +231,6 @@ count, use `size` and `n` rather than `height`, `width`, or
 | `seed`                   | int   | None    | Random seed (reproducible)     |
 | `negative_prompt`        | str   | None    | Negative prompt                |
 | `num_outputs_per_prompt` | int   | 1       | Number of images to generate   |
-| `use_system_prompt` | str | None | System prompt preset: `en_unified`, `en_vanilla`, `en_recaption`, `en_think_recaption`, `dynamic`, `None`, or custom text string. Only for HunyuanImage-3.0. |
-| `system_prompt` | str | None | Custom system prompt text. Only used when `use_system_prompt` is set to `custom`. Only for HunyuanImage-3.0. |
 
 ## Response Format
 

examples/online_serving/text_to_image/openai_chat_client.py

@@ -28,8 +28,6 @@ def generate_image(
     lora_name: str | None = None,
     lora_scale: float | None = None,
     lora_int_id: int | None = None,
-    use_system_prompt: str | None = None,
-    system_prompt: str | None = None,
 ) -> bytes | None:
     """Generate an image using the images generation API.
 
@@ -47,8 +45,6 @@ def generate_image(
         lora_name: LoRA name (optional, defaults to path stem)
         lora_scale: LoRA scale factor (default: 1.0)
         lora_int_id: LoRA integer ID (optional, derived from path if not provided)
-        use_system_prompt: System prompt for generation.
-        system_prompt: Custom system prompt.
 
     Returns:
         Image bytes or None if failed
@@ -74,10 +70,7 @@ def generate_image(
         payload["negative_prompt"] = negative_prompt
     if seed is not None:
         payload["seed"] = seed
-    if use_system_prompt is not None:
+
-        payload["use_system_prompt"] = use_system_prompt
-    if system_prompt is not None:
-        payload["system_prompt"] = system_prompt
     # Add LoRA if provided
     if lora_path:
         lora_body: dict = {
@@ -135,21 +128,9 @@ def main():
         default=None,
         help="LoRA integer id (cache key). If omitted, the server derives a stable id from lora_path.",
     )
-    parser.add_argument(
+
-        "--use-system-prompt",
-        type=str,
-        default=None,
-        help=(
-            "System prompt for generation. Use predefined types: 'en_unified', 'en_vanilla', 'en_recaption', 'en_think_recaption', 'dynamic', or 'None'; Or provide custom text string directly. Recommended en_unified. "
-        ),
-    )
-    parser.add_argument(
-        "--system-prompt",
-        type=str,
-        default=None,
-        help=("Custom system prompt. Used when --use-system-prompt is custom. "),
-    )
     args = parser.parse_args()
+
     print(f"Generating image for: {args.prompt}")
 
     image_bytes = generate_image(
@@ -165,8 +146,6 @@ def main():
         lora_name=args.lora_name,
         lora_scale=args.lora_scale if args.lora_path else None,
         lora_int_id=args.lora_int_id if args.lora_path else None,
-        use_system_prompt=args.use_system_prompt,
-        system_prompt=args.system_prompt,
     )
 
     if image_bytes:

pyproject.toml

@@ -223,4 +223,3 @@ extend-ignore-identifiers-re = [
 ue = "ue"
 semantics = "semantics"
 fullset = "fullset"
-Vai = "Vai"

requirements/common.txt

@@ -1,5 +1,4 @@
 # Common dependencies for all platforms
-av>=14.0.0
 omegaconf>=2.3.0
 librosa>=0.11.0
 resampy>=0.4.3

tests/benchmarks/metrics/test_metrics.py

@@ -1,67 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""
-Unit tests for metrics.py
-"""
-
-import pytest
-from vllm.benchmarks.serve import TaskType
-
-from vllm_omni.benchmarks.metrics.metrics import calculate_metrics
-from vllm_omni.benchmarks.patch.patch import MixRequestFuncOutput
-
-pytestmark = [pytest.mark.core_model, pytest.mark.benchmark, pytest.mark.cpu]
-
-
-def _make_output(prompt_len: int, output_tokens: int = 10) -> MixRequestFuncOutput:
-    """Build a minimal successful MixRequestFuncOutput for metrics aggregation."""
-    output = MixRequestFuncOutput()
-    output.success = True
-    output.prompt_len = prompt_len
-    output.output_tokens = output_tokens
-    output.generated_text = "x" * output_tokens
-    output.ttft = 0.1
-    output.text_latency = 1.0
-    output.latency = 1.0
-    output.start_time = 0.0
-    output.itl = [0.1] * max(output_tokens - 1, 0)
-    output.audio_ttfp = 0.0
-    output.audio_rtf = 0.0
-    output.audio_duration = 0.0
-    output.audio_frames = 0
-    output.input_audio_duration = 0.0
-    output.error = ""
-    return output
-
-
-# ============================================================================
-# total_input Tests
-# ============================================================================
-
-
-def test_total_input_aggregated_from_output_prompt_len():
-    """Test that total_input sums outputs[i].prompt_len, not input_requests[i].prompt_len."""
-    outputs = [_make_output(4992), _make_output(3000)]
-
-    metrics, _ = calculate_metrics(
-        input_requests=[],
-        outputs=outputs,
-        dur_s=10.0,
-        tokenizer=None,
-        selected_percentiles=[99.0],
-        goodput_config_dict={},
-        task_type=TaskType.GENERATION,
-        selected_percentile_metrics=[],
-        max_concurrency=None,
-        request_rate=float("inf"),
-        benchmark_duration=10.0,
-    )
-
-    assert metrics.total_input == 7992, (
-        "total_input should aggregate from outputs[i].prompt_len to reflect the true multimodal input token count"
-    )
-
-
-if __name__ == "__main__":
-    pytest.main([__file__, "-v", "-s"])

tests/benchmarks/patch/test_patch.py

@@ -574,59 +574,5 @@ class TestTextLatencyAttribute:
         )
 
 
-# ============================================================================
-# prompt_len Tests
-# ============================================================================
-
-
-@pytest.mark.asyncio
-async def test_prompt_len_assigned_from_usage(mocker: MockerFixture):
-    # Arrange: request claims prompt_len=100, but server reports 4992 (multimodal).
-    request_input = RequestFuncInput(
-        model="test-model",
-        model_name="test-model",
-        prompt="test prompt",
-        api_url="http://test.com/v1/chat/completions",
-        prompt_len=100,
-        output_len=20,
-    )
-
-    chunks = [
-        create_sse_chunk(
-            {
-                "choices": [{"delta": {"content": "Hello"}}],
-                "modality": "text",
-            }
-        ),
-        create_sse_chunk(
-            {
-                "choices": [{"delta": {"content": " world"}}],
-                "modality": "text",
-            }
-        ),
-        # Final usage chunk emitted because stream_options.include_usage=True.
-        create_sse_chunk(
-            {
-                "choices": [],
-                "usage": {"prompt_tokens": 4992, "completion_tokens": 2, "total_tokens": 4994},
-            }
-        ),
-        b"data: [DONE]\n\n",
-    ]
-
-    mock_response = MockResponse(200, chunks)
-    mock_session = mocker.AsyncMock()
-    mock_session.post = mocker.MagicMock(return_value=mock_response)
-
-    # Act
-    output = await async_request_openai_chat_omni_completions(request_input, mock_session)
-
-    # Assert
-    assert output.success is True
-    assert output.prompt_len == 4992, (
-        "prompt_len should be overridden by usage.prompt_tokens to reflect the true multimodal input token count"
-    )
-
-
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

tests/conftest.py

@@ -16,7 +16,6 @@ if "VLLM_TARGET_DEVICE" not in os.environ:
     os.environ["VLLM_TARGET_DEVICE"] = "cpu"
 
 import concurrent.futures
-import contextlib
 import gc
 import multiprocessing
 import socket
@@ -53,7 +52,6 @@ from vllm_omni.platforms import current_omni_platform
 
 logger = init_logger(__name__)
 
-
 PromptAudioInput = list[tuple[Any, int]] | tuple[Any, int] | None
 PromptImageInput = list[Any] | Any | None
 PromptVideoInput = list[Any] | Any | None
@@ -339,10 +337,10 @@ def log_test_name_before_test(request):
 
 def _run_pre_test_cleanup(enable_force=False):
     if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1" and not enable_force:
-        print("\nPre-test GPU cleanup skipped(Default off is typical when one worker/instance runs many tests.)\n")
+        print("GPU cleanup disabled")
         return
 
-    print("\nPre-test GPU status:")
+    print("Pre-test GPU status:")
 
     num_gpus = torch.cuda.device_count()
     if num_gpus > 0:
@@ -1089,22 +1087,6 @@ def _merge_base64_audio_to_segment(base64_list: list[str]):
     return merged
 
 
-@contextlib.contextmanager
-def _serialize_whisper_small_model_download():
-    """Serialize Whisper ``small`` cache writes across processes (Linux; ``fcntl``)."""
-    import fcntl
-
-    lock_path = Path.home() / ".cache" / "whisper" / ".small_model_download.lock"
-    lock_path.parent.mkdir(parents=True, exist_ok=True)
-    f = open(lock_path, "a+b")
-    try:
-        fcntl.flock(f.fileno(), fcntl.LOCK_EX)
-        yield
-    finally:
-        fcntl.flock(f.fileno(), fcntl.LOCK_UN)
-        f.close()
-
-
 def _whisper_transcribe_in_current_process(output_path: str) -> str:
     import whisper
 
@@ -1125,7 +1107,6 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
     else:
         use_accelerator = False
         device = "cpu"
-    with _serialize_whisper_small_model_download():
     model = whisper.load_model("small", device=device)
     try:
         text = model.transcribe(
@@ -1145,7 +1126,8 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
 
 
 def convert_audio_file_to_text(output_path: str) -> str:
-    """Convert an audio file to text in an isolated subprocess (spawn)."""
+    """Convert an audio file to text in an isolated subprocess."""
+    # Import locally to avoid impacting test module import time.
     ctx = multiprocessing.get_context("spawn")
     with concurrent.futures.ProcessPoolExecutor(max_workers=1, mp_context=ctx) as executor:
         future = executor.submit(_whisper_transcribe_in_current_process, output_path)
@@ -1164,6 +1146,18 @@ def convert_audio_bytes_to_text(raw_bytes: bytes) -> str:
     return text
 
 
+def merge_base64_and_convert_to_text(base64_list):
+    """
+    Merge a list of base64 encoded audio data and convert to text.
+    """
+    merged_audio = _merge_base64_audio_to_segment(base64_list)
+    output_path = f"./test_{uuid.uuid4().hex}.wav"
+    merged_audio.export(output_path, format="wav")
+    print(f"audio data is saved: {output_path}")
+    text = convert_audio_file_to_text(output_path)
+    return text
+
+
 def modify_stage_config(
     yaml_path: str,
     updates: dict[str, Any] = None,
@@ -1748,7 +1742,7 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
         label = str(top.get("label", "")).lower()
         conf = float(top.get("score", 0.0))
 
-        if conf < 0.5:
+        if conf < 0.6:
             gender = "unknown"
         # Some models use non-English labels (e.g., Russian). Normalize to 'male'/'female'.
         elif ("female" in label) or ("жен" in label):
@@ -1777,34 +1771,6 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
         return "unknown"
 
 
-_PRESET_VOICE_GENDER_MAP: dict[str, str] = {
-    "serena": "female",
-    "uncle_fu": "male",
-    "chelsie": "female",
-    "clone": "female",
-    "ethan": "male",
-}
-
-
-def _assert_preset_voice_gender_from_audio(
-    audio_bytes: bytes | None,
-    voice_name: str | None,
-) -> None:
-    """If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown)."""
-    if not voice_name or not audio_bytes:
-        return
-    key = str(voice_name).lower()
-    expected_gender = _PRESET_VOICE_GENDER_MAP.get(key)
-    if expected_gender is None:
-        return
-    estimated_gender = _estimate_voice_gender_from_audio(audio_bytes)
-    print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}")
-    if estimated_gender != "unknown":
-        assert estimated_gender == expected_gender, (
-            f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}"
-        )
-
-
 # Threshold aligned with _compute_pcm_hnr_db docstring (clean clone vs distorted).
 _MIN_PCM_SPEECH_HNR_DB = 1.0
 
@@ -1871,12 +1837,6 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
         if "audio" in modalities:
             assert response.audio_content is not None, "No audio output is generated"
             print(f"audio content is: {response.audio_content}")
-            speaker = request_config.get("speaker")
-            if speaker:
-                _assert_preset_voice_gender_from_audio(
-                    response.audio_bytes,
-                    speaker,
-                )
 
         if "text" in modalities:
             assert response.text_content is not None, "No text output is generated"
@@ -1889,14 +1849,12 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
             keywords = keywords_dict.get(word_type)
             if "text" in modalities:
                 if keywords:
-                    text_lower = response.text_content.lower()
+                    assert any(keyword in response.text_content.lower() for keyword in keywords), (
-                    assert any(str(kw).lower() in text_lower for kw in keywords), (
                         "The output does not contain any of the keywords."
                     )
             else:
                 if keywords:
-                    audio_lower = response.audio_content.lower()
+                    assert any(keyword in response.audio_content.lower() for keyword in keywords), (
-                    assert any(str(kw).lower() in audio_lower for kw in keywords), (
                         "The output does not contain any of the keywords."
                     )
 
@@ -1950,11 +1908,23 @@ def assert_audio_speech_response(
                 f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
             )
 
-        # Voice gender consistency check (preset names in ``_PRESET_VOICE_GENDER_MAP``).
+        # Voice gender consistency check:
         # When the estimator returns 'unknown', we treat it as inconclusive and do NOT fail the test.
-        _assert_preset_voice_gender_from_audio(
+        voice = (request_config.get("voice") or "").lower()
-            response.audio_bytes,
+        if voice and response.audio_bytes:
-            request_config.get("voice"),
+            estimated_gender = _estimate_voice_gender_from_audio(response.audio_bytes)
+            voice_gender_map = {
+                # adjust this mapping to your actual voice names
+                "serena": "female",
+                "uncle_fu": "male",
+                "clone": "female",
+            }
+            expected_gender = voice_gender_map.get(voice)
+            if expected_gender is not None:
+                print(f"Estimated voice gender from audio: {estimated_gender} (voice='{voice}')")
+                if estimated_gender != "unknown":
+                    assert estimated_gender == expected_gender, (
+                        f"Voice '{voice}' is expected {expected_gender}, but estimated gender is '{estimated_gender}'"
                     )
 
 
@@ -2071,11 +2041,7 @@ class OpenAIClientHandler:
 
             if audio_data or text_content:
                 if audio_data:
-                    merged_seg = _merge_base64_audio_to_segment(audio_data)
+                    audio_content = merge_base64_and_convert_to_text(audio_data)
-                    wav_buf = BytesIO()
-                    merged_seg.export(wav_buf, format="wav")
-                    result.audio_bytes = wav_buf.getvalue()
-                    audio_content = convert_audio_bytes_to_text(result.audio_bytes)
                 if audio_content and text_content:
                     similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
 
@@ -2130,8 +2096,7 @@ class OpenAIClientHandler:
 
             if audio_data or text_content:
                 if audio_data:
-                    result.audio_bytes = base64.b64decode(audio_data)
+                    audio_content = convert_audio_to_text(audio_data)
-                    audio_content = convert_audio_bytes_to_text(result.audio_bytes)
                 if audio_content and text_content:
                     similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
 
@@ -2300,9 +2265,8 @@ class OpenAIClientHandler:
             request_config: Request configuration dictionary containing parameters like model, messages, stream.
                 Optional ``use_audio_in_video`` (bool): when true, sets
                 ``extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}`` for Qwen-Omni video+audio
-                extraction.
+                extraction (merged with any existing ``extra_body`` / ``mm_processor_kwargs``).
-                Optional top-level ``speaker`` (str): Qwen3-Omni preset TTS speaker name; sent as
+                Optional ``extra_body`` (dict): passed through to ``chat.completions.create`` after merge.
-                ``extra_body["speaker"]`` to ``chat.completions.create``.
             request_num: Number of requests, defaults to 1 (single request)
 
         Returns:
@@ -2314,8 +2278,9 @@ class OpenAIClientHandler:
         modalities = request_config.get("modalities", ["text", "audio"])
 
         extra_body: dict[str, Any] = {}
-        if "speaker" in request_config:
+        raw_extra = request_config.get("extra_body")
-            extra_body["speaker"] = request_config["speaker"]
+        if raw_extra:
+            extra_body.update(raw_extra)
         if request_config.get("use_audio_in_video"):
             mm = dict(extra_body.get("mm_processor_kwargs") or {})
             mm["use_audio_in_video"] = True
@@ -2347,15 +2312,12 @@ class OpenAIClientHandler:
             # Send concurrent requests: run create + process in worker so e2e_latency includes full round-trip.
             def _one_omni_request():
                 start = time.perf_counter()
-                worker_kwargs: dict[str, Any] = {
+                chat_completion = self.client.chat.completions.create(
-                    "model": request_config.get("model"),
+                    model=request_config.get("model"),
-                    "messages": request_config.get("messages"),
+                    messages=request_config.get("messages"),
-                    "modalities": modalities,
+                    modalities=modalities,
-                    "stream": stream,
+                    stream=stream,
-                }
+                )
-                if extra_body_arg is not None:
-                    worker_kwargs["extra_body"] = extra_body_arg
-                chat_completion = self.client.chat.completions.create(**worker_kwargs)
                 if stream:
                     response = self._process_stream_omni_response(chat_completion)
                 else:

tests/core/sched/test_generation_scheduler_restore.py

@@ -9,10 +9,6 @@ those requests are permanently orphaned.
 import unittest
 from collections import deque
 
-import pytest
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
 
 class FakeAdapter:
     """Minimal mock of OmniChunkTransferAdapter tracking restore calls."""

tests/dfx/perf/stage_configs/qwen3_tts.yaml

@@ -88,7 +88,7 @@ runtime:
         connector_get_max_wait_first_chunk: 3000
         connector_get_max_wait: 300
         codec_chunk_frames: 25
-        codec_left_context_frames: 72
+        codec_left_context_frames: 25
 
   edges:
     - from: 0

tests/dfx/perf/tests/test_ltx2_vllm_omni.json

@@ -1,217 +0,0 @@
-[
-    {
-        "test_name": "test_ltx2_baseline_eager",
-        "description": "Single-device baseline with enforce-eager (no torch.compile)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
-    },
-
-    {
-        "test_name": "test_ltx2_torch_compile",
-        "description": "Single-device with torch.compile (default, no enforce-eager)",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
-    },
-
-    {
-        "test_name": "test_ltx2_cfg2_eager",
-        "description": "CFG-parallel=2 with enforce-eager",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cfg-parallel-size": 2,
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
-    },
-
-    {
-        "test_name": "test_ltx2_cfg2_compile",
-        "description": "CFG-parallel=2 with torch.compile",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cfg-parallel-size": 2,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
-    },
-
-    {
-        "test_name": "test_ltx2_cache_dit_eager",
-        "description": "CacheDiT with enforce-eager",
-        "server_type": "vllm-omni",
-        "server_params": {
-            "model": "Lightricks/LTX-2",
-            "serve_args": {
-                "cache-backend": "cache_dit",
-                "enforce-eager": true,
-                "enable-diffusion-pipeline-profiler": true
-            }
-        },
-        "benchmark_params": [
-            {
-                "name": "256x256_145f_steps6",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 256,
-                "height": 256,
-                "num-frames": 145,
-                "fps": 24,
-                "num-inference-steps": 6,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            },
-            {
-                "name": "480x768_41f_steps20",
-                "dataset": "random",
-                "task": "t2v",
-                "backend": "v1/videos",
-                "width": 768,
-                "height": 480,
-                "num-frames": 41,
-                "fps": 24,
-                "num-inference-steps": 20,
-                "num-prompts": 3,
-                "max-concurrency": 1,
-                "enable-negative-prompt": true
-            }
-        ]
-    }
-]

tests/diffusion/cache/test_teacache_extractors.py

@@ -21,13 +21,12 @@ from unittest.mock import MagicMock, Mock, patch
 import pytest
 import torch
 
-from tests.utils import hardware_test
 from vllm_omni.diffusion.cache.teacache.extractors import extract_flux2_klein_context
 from vllm_omni.diffusion.models.flux2_klein.flux2_klein_transformer import (
     Flux2Transformer2DModel,
 )
 
-pytestmark = [pytest.mark.core_model]
+pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
 
 
 @pytest.fixture(scope="function", autouse=True)
@@ -114,7 +113,6 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
     def get_sample_inputs(self, sample_inputs):
         return sample_inputs
 
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_modulated_input_shape(self, flux2_klein_module, sample_inputs):
         """Test that modulated_input has correct shape matching the model's inner_dim.
 
@@ -128,19 +126,16 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
         inner_dim = flux2_klein_module.inner_dim
         assert context.modulated_input.shape == (batch_size, img_seq_len, inner_dim)
 
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_run_transformer_blocks_callable(self, flux2_klein_module, sample_inputs):
         """Test that run_transformer_blocks is callable."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
         assert callable(context.run_transformer_blocks)
 
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_postprocess_callable(self, flux2_klein_module, sample_inputs):
         """Test that postprocess is callable."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
         assert callable(context.postprocess)
 
-    @hardware_test(res={"cuda": "L4"}, num_cards=1)
     def test_extra_states_contains_full_transformer(self, flux2_klein_module, sample_inputs):
         """Test that extra_states contains run_flux2_full_transformer_with_single."""
         context = extract_flux2_klein_context(flux2_klein_module, **sample_inputs)
@@ -159,7 +154,6 @@ class TestFlux2KleinExtractor(BaseExtractorTest):
         assert context is not None
         assert context.temb is not None
 
-    @pytest.mark.cpu
     def test_invalid_module_raises_error(self):
         """Test that invalid module without transformer_blocks raises ValueError."""
         invalid_module = Mock()

tests/diffusion/distributed/test_autoencoder_kl_wan.py

@@ -1,43 +0,0 @@
-import pytest
-import torch
-
-from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_wan as wan_vae_module
-from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import OmniAutoencoderKLWan
-
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
-
-class _DummyOmniAutoencoderKLWan(OmniAutoencoderKLWan):
-    def __init__(self, *, dtype: torch.dtype):
-        torch.nn.Module.__init__(self)
-        self.register_parameter("dummy_weight", torch.nn.Parameter(torch.ones(1, dtype=dtype)))
-
-
-def test_wan_vae_execution_context_handles_fp32():
-    model = _DummyOmniAutoencoderKLWan(dtype=torch.float32)
-    with model._execution_context():
-        output = model.dummy_weight + 1
-    assert output.dtype == torch.float32
-
-
-def test_wan_vae_execution_context_handles_bf16():
-    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
-    with model._execution_context():
-        output = model.dummy_weight + 1
-    assert output.dtype == torch.bfloat16
-
-
-def test_wan_vae_execution_context_uses_platform_autocast(mocker):
-    sentinel = object()
-    platform = mocker.Mock()
-    platform.create_autocast_context.return_value = sentinel
-    mocker.patch.object(wan_vae_module, "current_omni_platform", platform)
-
-    model = _DummyOmniAutoencoderKLWan(dtype=torch.bfloat16)
-
-    assert model._execution_context() is sentinel
-    platform.create_autocast_context.assert_called_once_with(
-        device_type=model.dummy_weight.device.type,
-        dtype=torch.bfloat16,
-        enabled=True,
-    )

tests/diffusion/distributed/test_autoencoder_kl_wan_encode.py

@@ -1,273 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-
-"""Unit tests for DistributedAutoencoderKLWan encode parallel (CPU-only)."""
-
-import pytest
-import torch
-
-pytestmark = [pytest.mark.cpu, pytest.mark.core_model]
-
-
-class _DummyConfig:
-    def __init__(self, patch_size=None, scale_factor_temporal=4):
-        self.patch_size = patch_size
-        self.scale_factor_temporal = scale_factor_temporal
-
-
-class _DummyWanVae:
-    """Minimal mock of DistributedAutoencoderKLWan for testing encode_tile_split."""
-
-    def __init__(
-        self,
-        config=None,
-        spatial_compression_ratio=8,
-        tile_sample_min_height=256,
-        tile_sample_min_width=256,
-        tile_sample_stride_height=192,
-        tile_sample_stride_width=192,
-    ):
-        self.config = config or _DummyConfig()
-        self.spatial_compression_ratio = spatial_compression_ratio
-        self.tile_sample_min_height = tile_sample_min_height
-        self.tile_sample_min_width = tile_sample_min_width
-        self.tile_sample_stride_height = tile_sample_stride_height
-        self.tile_sample_stride_width = tile_sample_stride_width
-        self.dtype = torch.float32
-
-        # Mock caches
-        self._enc_feat_map = None
-        self._enc_conv_idx = [0]
-
-    def clear_cache(self):
-        self._enc_feat_map = None
-        self._enc_conv_idx = [0]
-
-    def encoder(self, x, feat_cache=None, feat_idx=None):  # noqa: ARG002
-        # Simple mock: just return the input
-        return x
-
-    def quant_conv(self, x):
-        return x
-
-    def blend_v(self, _a, b, _blend_extent):
-        return b
-
-    def blend_h(self, _a, b, _blend_extent):
-        return b
-
-
-def _import_encode_tile_split():
-    """Import the encode_tile_split method from the module."""
-    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
-        DistributedAutoencoderKLWan,
-    )
-
-    return DistributedAutoencoderKLWan.encode_tile_split
-
-
-def _import_encode_tile_exec():
-    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
-        DistributedAutoencoderKLWan,
-    )
-
-    return DistributedAutoencoderKLWan.encode_tile_exec
-
-
-def _import_encode_tile_merge():
-    from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import (
-        DistributedAutoencoderKLWan,
-    )
-
-    return DistributedAutoencoderKLWan.encode_tile_merge
-
-
-class TestEncodeTileSplit:
-    """Tests for encode_tile_split method."""
-
-    def test_basic_split_without_patch_size(self):
-        """Test basic tile splitting without patch_size."""
-        encode_tile_split = _import_encode_tile_split()
-
-        vae = _DummyWanVae(
-            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
-            spatial_compression_ratio=8,
-            tile_sample_min_height=256,
-            tile_sample_min_width=256,
-            tile_sample_stride_height=192,
-            tile_sample_stride_width=192,
-        )
-
-        # Input: (B, C, T, H, W) = (1, 3, 5, 256, 256)
-        x = torch.randn(1, 3, 5, 256, 256)
-
-        tiletask_list, grid_spec = encode_tile_split(vae, x)
-
-        # With stride 192 and input size 256, we should get:
-        # Height: ceil(256/192) = 2 positions (0, 192) but 192+256 > 256, so only 1
-        # Actually for i in range(0, 256, 192): i = 0, 192 but 192 is out of bounds
-        # So we get 1x1 grid
-        assert len(tiletask_list) >= 1
-        assert grid_spec.grid_shape[0] >= 1
-        assert grid_spec.grid_shape[1] >= 1
-
-        # Check temporal chunking: 5 frames -> 1 + (5-1)//4 = 2 chunks
-        first_task = tiletask_list[0]
-        assert len(first_task.tensor) == 2  # 2 temporal chunks
-
-    def test_split_with_patch_size_scales_coordinates(self):
-        """Test that patch_size properly scales tile coordinates."""
-        encode_tile_split = _import_encode_tile_split()
-
-        # Without patch_size
-        vae_no_patch = _DummyWanVae(
-            config=_DummyConfig(patch_size=None, scale_factor_temporal=4),
-            spatial_compression_ratio=8,
-            tile_sample_min_height=256,
-            tile_sample_min_width=256,
-            tile_sample_stride_height=128,
-            tile_sample_stride_width=128,
-        )
-
-        # With patch_size=2 (simulating patchified input)
-        vae_with_patch = _DummyWanVae(
-            config=_DummyConfig(patch_size=2, scale_factor_temporal=4),
-            spatial_compression_ratio=8,
-            tile_sample_min_height=256,
-            tile_sample_min_width=256,
-            tile_sample_stride_height=128,
-            tile_sample_stride_width=128,
-        )
-
-        # Same patchified input size
-        x = torch.randn(1, 3, 5, 256, 256)
-
-        tasks_no_patch, _ = encode_tile_split(vae_no_patch, x)
-        tasks_with_patch, _ = encode_tile_split(vae_with_patch, x)
-
-        # With patch_size=2, stride becomes 128//2=64, so more tiles
-        assert len(tasks_with_patch) >= len(tasks_no_patch)
-
-    def test_temporal_compression_from_config(self):
-        """Test that temporal compression ratio is read from config."""
-        encode_tile_split = _import_encode_tile_split()
-
-        # temporal_compression=4 (default)
-        vae_4x = _DummyWanVae(
-            config=_DummyConfig(scale_factor_temporal=4),
-            tile_sample_min_height=512,
-            tile_sample_min_width=512,
-            tile_sample_stride_height=512,
-            tile_sample_stride_width=512,
-        )
-
-        # temporal_compression=2
-        vae_2x = _DummyWanVae(
-            config=_DummyConfig(scale_factor_temporal=2),
-            tile_sample_min_height=512,
-            tile_sample_min_width=512,
-            tile_sample_stride_height=512,
-            tile_sample_stride_width=512,
-        )
-
-        # 9 frames input
-        x = torch.randn(1, 3, 9, 512, 512)
-
-        tasks_4x, _ = encode_tile_split(vae_4x, x)
-        tasks_2x, _ = encode_tile_split(vae_2x, x)
-
-        # With 4x compression: 1 + (9-1)//4 = 3 chunks
-        assert len(tasks_4x[0].tensor) == 3
-
-        # With 2x compression: 1 + (9-1)//2 = 5 chunks
-        assert len(tasks_2x[0].tensor) == 5
-
-    def test_grid_spec_latent_dimensions(self):
-        """Test that grid_spec contains correct latent dimensions."""
-        encode_tile_split = _import_encode_tile_split()
-
-        vae = _DummyWanVae(
-            config=_DummyConfig(patch_size=None),
-            spatial_compression_ratio=8,
-            tile_sample_min_height=512,
-            tile_sample_min_width=512,
-            tile_sample_stride_height=512,
-            tile_sample_stride_width=512,
-        )
-
-        # Input: 512x512 with compression 8 -> 64x64 latent
-        x = torch.randn(1, 3, 5, 512, 512)
-
-        _, grid_spec = encode_tile_split(vae, x)
-
-        assert grid_spec.tile_spec["latent_height"] == 64
-        assert grid_spec.tile_spec["latent_width"] == 64
-
-
-class TestEncodeTileExec:
-    """Tests for encode_tile_exec method."""
-
-    def test_basic_exec(self):
-        """Test basic tile execution."""
-        encode_tile_exec = _import_encode_tile_exec()
-
-        vae = _DummyWanVae()
-
-        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
-            TileTask,
-        )
-
-        # Create a simple task with 2 temporal chunks
-        tile1 = torch.randn(1, 3, 1, 32, 32)
-        tile2 = torch.randn(1, 3, 4, 32, 32)
-        task = TileTask(tile_id=0, grid_coord=(0, 0), tensor=[tile1, tile2])
-
-        result = encode_tile_exec(vae, task)
-
-        # Result should concatenate temporal dimension
-        assert result.shape[2] == 5  # 1 + 4 frames
-
-
-class TestEncodeTileMerge:
-    """Tests for encode_tile_merge method."""
-
-    def test_basic_merge(self):
-        """Test basic tile merging."""
-        encode_tile_merge = _import_encode_tile_merge()
-
-        vae = _DummyWanVae()
-
-        from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
-            GridSpec,
-        )
-
-        # Create 2x2 grid of tiles
-        tile_00 = torch.ones(1, 16, 2, 32, 32) * 0
-        tile_01 = torch.ones(1, 16, 2, 32, 32) * 1
-        tile_10 = torch.ones(1, 16, 2, 32, 32) * 2
-        tile_11 = torch.ones(1, 16, 2, 32, 32) * 3
-
-        coord_tensor_map = {
-            (0, 0): tile_00,
-            (0, 1): tile_01,
-            (1, 0): tile_10,
-            (1, 1): tile_11,
-        }
-
-        grid_spec = GridSpec(
-            split_dims=(3, 4),
-            grid_shape=(2, 2),
-            tile_spec={
-                "latent_height": 48,
-                "latent_width": 48,
-                "blend_height": 8,
-                "blend_width": 8,
-                "tile_latent_stride_height": 24,
-                "tile_latent_stride_width": 24,
-            },
-        )
-
-        result = encode_tile_merge(vae, coord_tensor_map, grid_spec)
-
-        # Output should be (1, 16, 2, 48, 48)
-        assert result.shape == (1, 16, 2, 48, 48)

tests/diffusion/distributed/test_distributed_vae_executor.py

@@ -11,8 +11,6 @@ from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor impor
     TileTask,
 )
 
-pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
-
 
 class E2EOperator:
     """tiles with (2, 3) -- (H,W)"""
@@ -61,9 +59,9 @@ class E2EOperator:
 class DummyMixin(DistributedVaeMixin):
     def __init__(self):
         self.use_tiling = True
-        self.distributed_executor = MagicMock()
+        self.distributed_decoder = MagicMock()
-        self.distributed_executor.parallel_size = 2
+        self.distributed_decoder.parallel_size = 2
-        self.distributed_executor.group = None
+        self.distributed_decoder.group = None
 
 
 @pytest.fixture(autouse=True)

tests/diffusion/distributed/test_ulysses_uaa_perf.py

@@ -17,7 +17,6 @@ import pytest
 import torch
 import torch.distributed as dist
 
-from tests.utils import hardware_test
 from vllm_omni.diffusion.attention.parallel.ulysses import (
     _all_gather_int,
     _ulysses_all_to_all_any_o,
@@ -70,8 +69,6 @@ PERF_CASES: list[_PerfCase] = [
 
 
 @pytest.mark.parametrize("case", PERF_CASES)
-@pytest.mark.core_model
-@hardware_test(res={"cuda": "L4"}, num_cards=4)
 def test_ulysses_advanced_uaa_comm_overhead(case: _PerfCase) -> None:
     available_gpus = current_omni_platform.get_device_count()
     if available_gpus < case.world_size:

tests/diffusion/layers/test_adalayernorm.py

@@ -1,237 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Unit tests for shared AdaLayerNorm layers used by FLUX and other models."""
-
-import os
-
-import pytest
-import torch
-
-pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
-
-
-@pytest.fixture(autouse=True)
-def _init_distributed():
-    """Initialize the minimal distributed environment required by
-    ReplicatedLinear (tensor-parallel group must exist)."""
-    from vllm.distributed.parallel_state import (
-        cleanup_dist_env_and_memory,
-        init_distributed_environment,
-        initialize_model_parallel,
-    )
-
-    os.environ.setdefault("MASTER_ADDR", "localhost")
-    os.environ.setdefault("MASTER_PORT", "29501")
-    init_distributed_environment(
-        world_size=1,
-        rank=0,
-        local_rank=0,
-        distributed_init_method="env://",
-    )
-    initialize_model_parallel()
-    yield
-    cleanup_dist_env_and_memory()
-
-
-@pytest.fixture(autouse=True)
-def _force_default_gemm(monkeypatch):
-    """Force CPU-compatible GEMM dispatch for tests using CPU tensors.
-
-    vLLM's dispatch_unquantized_gemm() selects the backend by platform
-    (e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
-    CPU test tensors crash with NotImplementedError on ROCm.  Monkeypatch
-    the dispatcher to always return the default (torch.nn.functional.linear)
-    implementation which works on any device."""
-    from vllm.model_executor.layers.utils import default_unquantized_gemm
-
-    monkeypatch.setattr(
-        "vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
-        lambda: default_unquantized_gemm,
-    )
-
-
-def test_adalayernorm_import_from_shared_module():
-    """Verify imports work from the shared adalayernorm module."""
-    from vllm_omni.diffusion.layers.adalayernorm import (  # noqa: F401
-        AdaLayerNormContinuous,
-        AdaLayerNormZero,
-        AdaLayerNormZeroSingle,
-    )
-
-
-def test_adalayernorm_zero_forward_shape():
-    """AdaLayerNormZero produces correct output shapes (x, gate, shift, scale, gate)."""
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
-
-    dim = 64
-    batch = 2
-    seq_len = 4
-    norm = AdaLayerNormZero(dim)
-
-    x = torch.randn(batch, seq_len, dim)
-    emb = torch.randn(batch, dim)
-
-    out_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = norm(x, emb)
-
-    assert out_x.shape == (batch, seq_len, dim)
-    assert gate_msa.shape == (batch, dim)
-    assert shift_mlp.shape == (batch, dim)
-    assert scale_mlp.shape == (batch, dim)
-    assert gate_mlp.shape == (batch, dim)
-
-
-def test_adalayernorm_zero_single_forward_shape():
-    """AdaLayerNormZeroSingle produces (x, gate) with correct shapes."""
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
-
-    dim = 64
-    batch = 2
-    seq_len = 4
-    norm = AdaLayerNormZeroSingle(dim)
-
-    x = torch.randn(batch, seq_len, dim)
-    emb = torch.randn(batch, dim)
-
-    out_x, gate = norm(x, emb)
-
-    assert out_x.shape == (batch, seq_len, dim)
-    assert gate.shape == (batch, dim)
-
-
-def test_adalayernorm_continuous_forward_shape():
-    """AdaLayerNormContinuous produces correct output shape."""
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
-
-    dim = 64
-    cond_dim = 64
-    batch = 2
-    seq_len = 4
-    norm = AdaLayerNormContinuous(dim, cond_dim)
-
-    x = torch.randn(batch, seq_len, dim)
-    conditioning = torch.randn(batch, cond_dim)
-
-    out = norm(x, conditioning)
-
-    assert out.shape == (batch, seq_len, dim)
-
-
-def test_adalayernorm_zero_accepts_quant_config():
-    """Constructor accepts quant_config=None and prefix='test' without error."""
-    from vllm_omni.diffusion.layers.adalayernorm import (
-        AdaLayerNormContinuous,
-        AdaLayerNormZero,
-        AdaLayerNormZeroSingle,
-    )
-
-    # Should not raise with quant_config=None and prefix
-    AdaLayerNormZero(64, quant_config=None, prefix="test.norm1")
-    AdaLayerNormZeroSingle(64, quant_config=None, prefix="test.norm")
-    AdaLayerNormContinuous(64, 64, quant_config=None, prefix="test.norm_out")
-
-
-def test_adalayernorm_uses_replicated_linear():
-    """Verify .linear is a ReplicatedLinear instance (not nn.Linear)."""
-    from vllm.model_executor.layers.linear import ReplicatedLinear
-
-    from vllm_omni.diffusion.layers.adalayernorm import (
-        AdaLayerNormContinuous,
-        AdaLayerNormZero,
-        AdaLayerNormZeroSingle,
-    )
-
-    norm_zero = AdaLayerNormZero(64)
-    assert isinstance(norm_zero.linear, ReplicatedLinear)
-
-    norm_zero_single = AdaLayerNormZeroSingle(64)
-    assert isinstance(norm_zero_single.linear, ReplicatedLinear)
-
-    norm_continuous = AdaLayerNormContinuous(64, 64)
-    assert isinstance(norm_continuous.linear, ReplicatedLinear)
-
-
-# ── Numerical equivalence tests against diffusers originals ──
-
-
-def _copy_weights(src_linear, dst_replicated_linear):
-    """Copy weights from nn.Linear to ReplicatedLinear for comparison."""
-    dst_replicated_linear.weight.data.copy_(src_linear.weight.data)
-    if src_linear.bias is not None and dst_replicated_linear.bias is not None:
-        dst_replicated_linear.bias.data.copy_(src_linear.bias.data)
-
-
-def test_adalayernorm_zero_matches_diffusers():
-    """Verify AdaLayerNormZero produces identical output to diffusers original."""
-    from diffusers.models.normalization import (
-        AdaLayerNormZero as DiffusersAdaLayerNormZero,
-    )
-
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
-
-    dim = 64
-    torch.manual_seed(42)
-    ours = AdaLayerNormZero(dim)
-    ref = DiffusersAdaLayerNormZero(dim)
-
-    # Copy weights: nn.Linear -> ReplicatedLinear
-    _copy_weights(ref.linear, ours.linear)
-
-    x = torch.randn(2, 4, dim)
-    emb = torch.randn(2, dim)
-
-    out_ours = ours(x, emb)
-    out_ref = ref(x, emb=emb)
-
-    for o, r in zip(out_ours, out_ref):
-        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
-
-
-def test_adalayernorm_zero_single_matches_diffusers():
-    """Verify AdaLayerNormZeroSingle produces identical output to diffusers original."""
-    from diffusers.models.normalization import (
-        AdaLayerNormZeroSingle as DiffusersAdaLayerNormZeroSingle,
-    )
-
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
-
-    dim = 64
-    torch.manual_seed(42)
-    ours = AdaLayerNormZeroSingle(dim)
-    ref = DiffusersAdaLayerNormZeroSingle(dim)
-
-    _copy_weights(ref.linear, ours.linear)
-
-    x = torch.randn(2, 4, dim)
-    emb = torch.randn(2, dim)
-
-    out_ours = ours(x, emb)
-    out_ref = ref(x, emb=emb)
-
-    for o, r in zip(out_ours, out_ref):
-        torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
-
-
-def test_adalayernorm_continuous_matches_diffusers():
-    """Verify AdaLayerNormContinuous produces identical output to diffusers original."""
-    from diffusers.models.normalization import (
-        AdaLayerNormContinuous as DiffusersAdaLayerNormContinuous,
-    )
-
-    from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
-
-    dim = 64
-    cond_dim = 64
-    torch.manual_seed(42)
-    # Match constructor args: diffusers defaults elementwise_affine=True, eps=1e-5
-    ours = AdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
-    ref = DiffusersAdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
-
-    _copy_weights(ref.linear, ours.linear)
-
-    x = torch.randn(2, 4, dim)
-    cond = torch.randn(2, cond_dim)
-
-    out_ours = ours(x, cond)
-    out_ref = ref(x, cond)
-
-    torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5)

tests/diffusion/lora/conftest.py

@@ -1,56 +0,0 @@
-# SPDX-License-Identifier: Apache-2.0
-# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
-"""Shared test helpers for diffusion LoRA tests."""
-
-from __future__ import annotations
-
-import torch
-from vllm.model_executor.layers.linear import LinearBase
-
-
-class FakeLinearBase(LinearBase):
-    """Minimal LinearBase stub for LoRA layer discovery."""
-
-    def __init__(self):
-        torch.nn.Module.__init__(self)
-
-
-class DummyBaseLayerWithLoRA(torch.nn.Module):
-    """Fake LoRA wrapper that records set/reset/create calls."""
-
-    def __init__(self, base_layer: torch.nn.Module):
-        super().__init__()
-        self.base_layer = base_layer
-
-        self.set_calls: list[
-            tuple[list[torch.Tensor | None] | torch.Tensor, list[torch.Tensor | None] | torch.Tensor]
-        ] = []
-        self.reset_calls: int = 0
-        self.create_calls: int = 0
-
-    def set_lora(self, index: int, lora_a, lora_b):
-        assert index == 0
-        self.set_calls.append((lora_a, lora_b))
-
-    def reset_lora(self, index: int):
-        assert index == 0
-        self.reset_calls += 1
-
-    def create_lora_weights(self, max_loras, lora_config, model_config):
-        self.create_calls += 1
-
-
-def fake_replace_submodule(
-    root: torch.nn.Module,
-    module_name: str,
-    submodule: torch.nn.Module,
-    replace_calls: list[str] | None = None,
-) -> None:
-    """Replace a submodule by traversing dotted paths correctly."""
-    if replace_calls is not None:
-        replace_calls.append(module_name)
-    parts = module_name.split(".")
-    parent = root
-    for attr in parts[:-1]:
-        parent = getattr(parent, attr)
-    setattr(parent, parts[-1], submodule)