Compare commits

..
Author SHA1 Message Date
hsliuustc0106 8affabd527 update release
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-03-31 20:39:22 +08:00
392 changed files with 4034 additions and 38212 deletions

No files matched your search

-20
View File
@@ -1,20 +0,0 @@
steps:
- label: "Build and upload wheel"
key: "build-wheel"
agents:
queue: cpu_queue_release
commands:
- "curl -LsSf https://astral.sh/uv/install.sh | sh"
- 'export PATH="$HOME/.local/bin:$PATH"'
- "uv venv --python=3.12 && source .venv/bin/activate"
- "uv pip install --upgrade build"
- "python3 -m build"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- label: "Generate and upload wheel indices"
depends_on: "build-wheel"
allow_dependency_failure: true
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
+1 -1
View File
@@ -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
+2 -17
View File
@@ -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:
@@ -1,88 +0,0 @@
#!/usr/bin/env bash
set -ex
# Generate and upload wheel indices for all vllm-omni wheels in the commit directory.
# This script should run once after all wheels have been built and uploaded.
# All paths are under the omni/ prefix in the vllm-wheels S3 bucket.
# ======== setup ========
BUCKET="vllm-wheels"
INDICES_OUTPUT_DIR="indices"
PYTHON="${PYTHON_PROG:-python3}"
SUBPATH="omni/$BUILDKITE_COMMIT"
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# detect if python3.12+ is available
has_new_python=$($PYTHON -c "print(1 if __import__('sys').version_info >= (3,12) else 0)")
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"
fi
echo "Using python interpreter: $PYTHON"
echo "Python version: $($PYTHON --version)"
# ======== generate and upload indices ========
# list all wheels in the commit directory
echo "Existing wheels on S3:"
aws s3 ls "$S3_COMMIT_PREFIX" || echo "(no objects found)"
obj_json="objects.json"
aws s3api list-objects-v2 --bucket "$BUCKET" --prefix "$SUBPATH/" --delimiter / --output json > "$obj_json"
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
# Generate indices -- the version is just the commit hash (not omni/{commit})
# because relative paths are computed between the index and wheel directories,
# both of which live under the omni/ prefix in S3.
$PYTHON .buildkite/scripts/generate-nightly-index.py \
--version "$BUILDKITE_COMMIT" \
--current-objects "$obj_json" \
--output-dir "$INDICES_OUTPUT_DIR" \
--comment "commit $BUILDKITE_COMMIT"
# copy indices to /omni/{commit}/ unconditionally
echo "Uploading indices to $S3_COMMIT_PREFIX"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "$S3_COMMIT_PREFIX"
# copy to /omni/nightly/ when NIGHTLY=1
if [[ "${NIGHTLY:-}" == "1" ]]; then
echo "Uploading indices to overwrite /omni/nightly/"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/nightly/"
fi
# detect version from any wheel in the commit directory
first_wheel_key=$($PYTHON -c "import json; obj=json.load(open('$obj_json')); print(next((c['Key'] for c in obj.get('Contents', []) if c['Key'].endswith('.whl')), ''))")
if [[ -z "$first_wheel_key" ]]; then
echo "Error: No wheels found in $S3_COMMIT_PREFIX"
exit 1
fi
first_wheel=$(basename "$first_wheel_key")
aws s3 cp "s3://$BUCKET/${first_wheel_key}" "/tmp/${first_wheel}"
version=$(unzip -p "/tmp/${first_wheel}" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
rm -f "/tmp/${first_wheel}"
echo "Version in wheel: $version"
pure_version="${version%%+*}"
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/$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" \
--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/"
fi
@@ -1,192 +0,0 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import re
import sys
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.parse import quote
def normalize_package_name(name: str) -> str:
"""Normalize package name per PEP 503."""
return re.sub(r"[-_.]+", "-", name).lower()
if not sys.version_info >= (3, 12):
raise RuntimeError("This script requires Python 3.12 or higher.")
INDEX_HTML_TEMPLATE = """<!DOCTYPE html>
<html>
<!-- {comment} -->
<meta name="pypi:repository-version" content="1.0">
<body>
{items}
</body>
</html>
"""
@dataclass
class WheelFileInfo:
package_name: str
version: str
build_tag: str | None
python_tag: str
abi_tag: str
platform_tag: str
filename: str
def parse_from_filename(file: str) -> WheelFileInfo:
"""
Parse wheel filename per PEP 427:
{package_name}-{version}(-{build_tag})?-{python_tag}-{abi_tag}-{platform_tag}.whl
"""
wheel_file_re = re.compile(
r"^(?P<package_name>.+)-(?P<version>[^-]+?)(-(?P<build_tag>[^-]+))?-(?P<python_tag>[^-]+)-(?P<abi_tag>[^-]+)-(?P<platform_tag>[^-]+)\.whl$"
)
match = wheel_file_re.match(file)
if not match:
raise ValueError(f"Invalid wheel file name: {file}")
return WheelFileInfo(
package_name=match.group("package_name"),
version=match.group("version"),
build_tag=match.group("build_tag"),
python_tag=match.group("python_tag"),
abi_tag=match.group("abi_tag"),
platform_tag=match.group("platform_tag"),
filename=file,
)
def generate_project_list(package_names: list[str], comment: str = "") -> str:
"""Generate top-level PEP 503 project list HTML."""
href_tags = []
for name in sorted(package_names):
href_tags.append(f' <a href="{name}/">{name}/</a><br/>')
return INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
def generate_package_index(
wheel_files: list[WheelFileInfo],
wheel_base_dir: Path,
index_base_dir: Path,
comment: str = "",
) -> tuple[str, str]:
"""Generate package index HTML and metadata JSON linking to wheel files."""
href_tags = []
metadata = []
for file in sorted(wheel_files, key=lambda x: x.filename):
relative_path = wheel_base_dir.relative_to(index_base_dir, walk_up=True) / file.filename
# handle '+' in URL; avoid double-encoding '/' and '%2B' (AWS S3 behavior)
file_path_quoted = quote(relative_path.as_posix(), safe=":%/")
href_tags.append(f' <a href="{file_path_quoted}">{file.filename}</a><br/>')
file_meta = asdict(file)
file_meta["path"] = file_path_quoted
metadata.append(file_meta)
index_str = INDEX_HTML_TEMPLATE.format(items="\n".join(href_tags), comment=comment)
metadata_str = json.dumps(metadata, indent=2)
return index_str, metadata_str
def generate_index(
whl_files: list[str],
wheel_base_dir: Path,
index_base_dir: Path,
comment: str = "",
):
"""
Generate PEP 503 index for all wheel files.
Output structure:
index_base_dir/
index.html # project list linking to vllm-omni/
vllm-omni/
index.html # package index linking to wheel files
metadata.json # machine-readable metadata
"""
parsed_files = [parse_from_filename(f) for f in whl_files]
if not parsed_files:
print("No wheel files found, skipping index generation.")
return
comment_str = f" ({comment})" if comment else ""
comment_tmpl = f"Generated on {datetime.now().isoformat()}{comment_str}"
# Group by normalized package name
packages: dict[str, list[WheelFileInfo]] = {}
for file in parsed_files:
name = normalize_package_name(file.package_name)
packages.setdefault(name, []).append(file)
print(f"Found packages: {list(packages.keys())}")
# Generate per-package index
for package, files in packages.items():
package_dir = index_base_dir / package
package_dir.mkdir(parents=True, exist_ok=True)
index_str, metadata_str = generate_package_index(files, wheel_base_dir, package_dir, comment)
with open(package_dir / "index.html", "w") as f:
f.write(index_str)
with open(package_dir / "metadata.json", "w") as f:
f.write(metadata_str)
# Generate top-level project list
project_list_str = generate_project_list(sorted(packages.keys()), comment_tmpl)
with open(index_base_dir / "index.html", "w") as f:
f.write(project_list_str)
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Generate PEP 503 wheel index from S3 object listing.")
parser.add_argument("--version", type=str, required=True, help="Version string (e.g., commit hash)")
parser.add_argument("--current-objects", type=str, required=True, help="Path to JSON from S3 list-objects-v2")
parser.add_argument("--output-dir", type=str, required=True, help="Directory to write index files")
parser.add_argument("--wheel-dir", type=str, default=None, help="Wheel directory (defaults to --version)")
parser.add_argument("--comment", type=str, default="", help="Comment for generated HTML")
args = parser.parse_args()
version = args.version
if "\\" in version or "/" in version:
raise ValueError("Version string must not contain slashes or backslashes.")
output_dir = Path(args.output_dir)
output_dir.mkdir(parents=True, exist_ok=True)
with open(args.current_objects) as f:
current_objects: dict[str, list[dict[str, Any]]] = json.load(f)
wheel_files = [
item["Key"].split("/")[-1] for item in current_objects.get("Contents", []) if item["Key"].endswith(".whl")
]
print(f"Found {len(wheel_files)} wheel files for version {version}: {wheel_files}")
# For release versions, filter to only matching non-dev wheels
PY_VERSION_RE = re.compile(r"^\d+\.\d+\.\d+([a-zA-Z0-9.+-]*)?$")
if PY_VERSION_RE.match(version):
wheel_files = [f for f in wheel_files if version in f and "dev" not in f]
print(f"Non-nightly version detected, wheel files used: {wheel_files}")
else:
print("Nightly version detected, keeping all wheel files.")
wheel_dir = (args.wheel_dir or version).strip().rstrip("/")
wheel_base_dir = Path(output_dir).parent / wheel_dir
index_base_dir = Path(output_dir)
generate_index(
whl_files=wheel_files,
wheel_base_dir=wheel_base_dir,
index_base_dir=index_base_dir,
comment=args.comment.strip(),
)
print(f"Successfully generated index in {output_dir}")
@@ -1,33 +0,0 @@
#!/usr/bin/env bash
set -ex
# Upload a single wheel to S3 under the omni/ prefix.
# Index generation is handled separately by generate-and-upload-nightly-index.sh.
BUCKET="vllm-wheels"
SUBPATH="omni/$BUILDKITE_COMMIT"
S3_COMMIT_PREFIX="s3://$BUCKET/$SUBPATH/"
# ========= collect & upload the wheel ==========
# python3 -m build outputs to dist/ by default
wheel_files=(dist/*.whl)
# Check that exactly one wheel is found
if [[ ${#wheel_files[@]} -ne 1 ]]; then
echo "Error: Expected exactly one wheel file in dist/, but found ${#wheel_files[@]}"
exit 1
fi
wheel="${wheel_files[0]}"
echo "Uploading wheel: $wheel"
# Extract the version from the wheel
version=$(unzip -p "$wheel" '**/METADATA' | grep '^Version: ' | cut -d' ' -f2)
echo "Version in wheel: $version"
# Upload wheel to S3
aws s3 cp "$wheel" "$S3_COMMIT_PREFIX"
echo "Wheel uploaded to $S3_COMMIT_PREFIX. Index generation is handled by a separate step."
@@ -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
+16 -15
View File
@@ -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 tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
@@ -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"
@@ -258,7 +260,9 @@ steps:
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 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/offline_inference/test_qwen3_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -344,7 +348,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 +393,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 +414,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:
-367
View File
@@ -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
+398 -186
View File
@@ -1,199 +1,235 @@
steps:
# Group: collapses under one heading in the Buildkite UI; child steps still run in parallel.
- 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"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni · Function 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/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
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Omni · Doc 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 omni 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: Omni · Doc 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/examples/ -m "advanced_model and omni 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: Omni · Perf Test"
key: nightly-omni-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export BENCHMARK_DIR=tests/dfx/perf/results
- 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
- label: ":full_moon: Omni Model Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
set +e
pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
EXIT1=$$?
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
EXIT2=$$?
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
EXIT3=$$?
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
EXIT4=$$?
exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
agents:
queue: "cpu_queue_premerge"
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: ":bar_chart: Testcase Statistics"
key: nightly-testcase-statistics
- 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
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Diffusion Model Test with H100"
timeout_in_minutes: 120
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion" -m "advanced_model and diffusion and H100" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion Model Wan22 completed Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py -m "advanced_model" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion Model Test with L4"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model"
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
- "HF_TOKEN"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Documentation Example Code Test with H100"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
- pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
key: nightly-omni-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export BENCHMARK_DIR=tests/dfx/perf/results
- 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"
- 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
+25 -87
View File
@@ -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
-534
View File
@@ -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
@@ -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 |
@@ -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)
```
@@ -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`.
@@ -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),
}
```
@@ -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
-284
View File
@@ -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)
-34
View File
@@ -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.
+2 -2
View File
@@ -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
+1 -16
View File
@@ -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/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
.claude/
# Codex
AGENTS.md
-3
View File
@@ -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
+2 -10
View File
@@ -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}" \
--task-type "${TASK_TYPE}"
--result-dir "${RESULT_DIR}"
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}" \
--task-type "${TASK_TYPE}"
--result-dir "${RESULT_DIR}"
# Allow GPU memory to settle
sleep 5
@@ -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)"
)
@@ -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(
+4 -2
View File
@@ -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
-22
View File
@@ -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 []
+13 -13
View File
@@ -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
# RUN git fetch origin --tags && git checkout v0.18.0
WORKDIR /vllm-workspace/vllm
RUN git fetch origin --tags && git checkout v0.18.0
# WORKDIR /vllm-workspace/vllm-ascend
# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# # Install vllm-ascend
# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
# source /usr/local/Ascend/nnal/atb/set_env.sh && \
# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
# python3 -m pip cache purge
WORKDIR /vllm-workspace/vllm-ascend
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# Install vllm-ascend
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip cache purge
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
+13 -13
View File
@@ -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
# RUN git fetch origin --tags && git checkout v0.18.0
WORKDIR /vllm-workspace/vllm
RUN git fetch origin --tags && git checkout v0.18.0
# WORKDIR /vllm-workspace/vllm-ascend
# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# # Install vllm-ascend
# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
# source /usr/local/Ascend/nnal/atb/set_env.sh && \
# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
# python3 -m pip cache purge
WORKDIR /vllm-workspace/vllm-ascend
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# Install vllm-ascend
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip cache purge
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
+1 -1
View File
@@ -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.
+1 -1
View File
@@ -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
-1
View File
@@ -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
+1
View File
@@ -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][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 96 KiB

After

Width:  |  Height:  |  Size: 99 KiB

+1 -3
View File
@@ -37,12 +37,10 @@ Committers have write access and merge rights. They typically have deep expertis
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
- [@wtomin](https://github.com/wtomin):
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
## Meetings
+2 -1
View File
@@ -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 ⬜
@@ -271,7 +272,7 @@ Before entering specific testing levels, the project establishes two common spec
L1 and L2 level testing form the foundation of the quality assurance system. L1 level testing focuses on verifying the internal logic correctness of code units (e.g., functions, classes), ensuring each independent component behaves as designed.
L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle.
L2 level testing builds upon L1 by introducing GPU resources and verifying that the end-to-end (E2E) process of the model in basic deployment scenarios is smooth. For example, it uses dummy models to confirm that core interfaces like the inference pipeline, output format, and streaming response work properly. The common goal of these two levels is to provide developers with rapid feedback, discovering and fixing issues early in the development cycle .
+1
View File
@@ -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 ⬜
@@ -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
-221
View File
@@ -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.
@@ -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_device: "cpu" # "cpu" for pinned memory (recommended), "cuda" for GPUDirect RDMA
memory_pool_size: 2147483648 # 2GB memory pool
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_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_size` | 1 GB | Total size of the RDMA-registered memory pool in bytes. |
| `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 |
|---|---|---|---|---|
| CPU Pinned | `memory_pool_device: "cpu"` | 4 GB | 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) |
| Mode | Config | Data Flow | Best For |
|---|---|---|---|
| CPU Pinned | `memory_pool_device: "cpu"` | GPU → CPU pool → RDMA → CPU pool → GPU | Most hardware topologies (recommended) |
| 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,
+18 -188
View File
@@ -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)
- [Encode Parallel Implementation](#encode-parallel-implementation)
- [Step-by-Step Implementation](#step-by-step-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
### 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) |
+ Accelerates decoding latency
### 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 |
|-------|------|-----------------|-----------------|
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | ✅ | ❌ |
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | ✅ | ✅ |
| **Qwen-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_qwenimage.py` | ✅ | ❌ |
| Model | Path | Notes |
|-------|------|-------|
| **Z-Image** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl.py` | Distributed AutoencoderKL |
| **Wan2.2** | `vllm_omni/diffusion/distributed/autoencoders/autoencoder_kl_wan.py` | Distributed AutoencoderKLWan |
| **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`
2. **Decode Parallel** - Refactor `tiled_decode` into `tile_split`/`tile_exec`/`tile_merge`
3. **Encode Parallel** (optional) - Implement `encode_tile_split`/`encode_tile_exec`/`encode_tile_merge` for I2V models
4. **Change VAE model in pipeline** - Use the distributed version
5. **Test** - Verify numerical consistency with `vae_patch_parallel_size=1` vs `N`
1. **Implement Distributed Vae** - mainly copy from `diffusers` tiled_decode, and refactor into split/exec/merge
2. **Change vae model in pipeline to Distributed Vae**
3. **Test** - Verify with `tensor_parallel_size=N` quality
@@ -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
+2 -2
View File
@@ -192,7 +192,7 @@ class _BaseScheduler(SchedulerInterface):
self._waiting = deque()
self._running = []
self._finished_req_ids = set()
self.max_num_running_reqs = 1
self._max_batch_size = 1
```
**Design Features**:
@@ -201,7 +201,7 @@ class _BaseScheduler(SchedulerInterface):
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
- **Current constraint**: `_max_batch_size` remains `1` because the current engine path is still synchronous request-mode execution.
#### 2.4 Current `RequestScheduler` Policy
@@ -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)
-20
View File
@@ -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"
@@ -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.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
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
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -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]
@@ -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]
@@ -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.18.0rc1
# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1-a3
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
VLLM_TARGET_DEVICE=empty pip install -v -e .
cd /vllm-workspace/vllm
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
+2 -2
View File
@@ -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
-1
View File
@@ -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,
-5
View File
@@ -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.
+3 -24
View File
@@ -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"],
"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"
}
]
"voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
}
```
`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,
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
"file_size": 1024000
}
}
```
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 "speaker_description=warm narrator"
-F "ref_text=The exact transcript of the audio sample."
```
## Streaming Text Input (WebSocket)
@@ -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.
@@ -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) |
|--------------|----------------|-----------|---------------------|-------------------|-----------------------------------|
| 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"` |
| Architecture | Example Models | DiT Class | Model-Level Offload | Layerwise Offload | Blocks Attr (Layerwise specific) |
|--------------|----------------|-----------|---------------------|-------------------|-------------|
| 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
@@ -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.
@@ -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"},
+49 -131
View File
@@ -12,10 +12,9 @@
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**
## Supported Features
@@ -44,12 +43,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
@@ -71,16 +64,6 @@ Extension methods add specialized capabilities to diffusion models beyond standa
| **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
### Execution Modes
Execution modes control how the diffusion pipeline processes denoise steps.
| Method | Description | Best For |
|--------|-------------|----------|
| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control |
**Note:** Step execution is currently supported by QwenImagePipeline only. See [Supported Models](#supported-models) for details.
### Quantization Methods
| Method | Configuration | Description | Best For |
@@ -104,50 +87,48 @@ The following tables show which models support each feature:
### ImageGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MagiHuman** | ❌ | ❌ | ❌ | ❓ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **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.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
### AudioGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ |
## Feature Compatibility
@@ -158,22 +139,21 @@ The following tables show which models support each feature:
- ❌: No support plan
- ❓: Not verified yet and Not Recommended
| | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference | 🔄Step Execution |
|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| **⚡TeaCache** | | | | | | | | | | | | | | |
| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | | |
| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | | |
| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | | |
| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | | |
| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | | |
| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | | |
| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | | |
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
| | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference |
|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| **⚡TeaCache** | | | | | | | | | | | | | |
| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | |
| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | |
| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | |
| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | |
| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | |
| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | |
| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | |
| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | |
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | |
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | |
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | |
!!! info
@@ -182,62 +162,8 @@ The following tables show which models support each feature:
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. CPU Offloading (Layerwise) supports single-card for now.
5. Using FP8-Quant as an example of qunatization methods.
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:**
@@ -259,14 +185,6 @@ Measured on NVIDIA H800:
- **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
**Execution Modes:**
- **[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
@@ -37,45 +37,12 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "input_reference=@/path/to/input.png" \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "negative_prompt=low quality, blurry, static" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42" \
-o sync_i2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
@@ -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",
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
"mime_type": "audio/wav"
}
]
}
```
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,
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
"file_size": 1024000
}
}
```
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 "speaker_description=warm narrator"
-F "ref_text=The exact transcript of the audio sample."
```
### 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"
@@ -37,43 +37,12 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42" \
-o sync_t2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
+8 -52
View File
@@ -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
@@ -122,7 +101,7 @@ def main():
if not prompts:
# Default prompt for text2img test if none provided
prompts = ["A cute cat"]
prompts = ["<|im_start|>A cute cat<|im_end|>"]
print(f"[Info] No prompts provided, using default: {prompts}")
omni_outputs = []
@@ -131,12 +110,8 @@ def main():
from vllm_omni.entrypoints.omni import Omni
omni_kwargs = {}
stage_configs_path = args.stage_configs_path
if args.think and stage_configs_path is None:
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
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.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"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>"
final_prompt_text = f"<|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"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
)
final_prompt_text = f"<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
prompt_dict = {
"prompt": final_prompt_text,
"multi_modal_data": {"image": loaded_image},
@@ -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"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>"
final_prompt_text = f"<|im_start|>user\n{p}<|im_end|>\n<|im_start|>assistant\n"
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
formatted_prompts.append(prompt_dict)
else:
think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>"
final_prompt_text = f"<|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
@@ -199,6 +168,7 @@ def main():
params_list = omni.default_sampling_params_list
if args.modality in ("text2img", "img2img"):
params_list[0].max_tokens = 1 # type: ignore
if len(params_list) > 1:
diffusion_params = params_list[1]
diffusion_params.num_inference_steps = args.steps # type: ignore
@@ -209,12 +179,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 +187,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 +195,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)
@@ -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).
File diff suppressed because it is too large. Load diff
@@ -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.
@@ -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:
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
use_ltx2_export = is_ltx2
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
if isinstance(audio_np, list):
audio_np = audio_np[0] if audio_np else None
if isinstance(audio_np, torch.Tensor):
audio_np = audio_np.detach().cpu().float().numpy()
if isinstance(audio_np, np.ndarray):
audio_np = np.squeeze(audio_np).astype(np.float32)
audio_out = None
if audio is not None:
if isinstance(audio, list):
audio = audio[0] if audio else None
if isinstance(audio, np.ndarray):
audio = torch.from_numpy(audio)
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(
frames_u8,
audio_np,
fps=float(fps),
audio_sample_rate=args.audio_sample_rate,
encode_video(
video_tensor,
fps=fps,
audio=audio_out,
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}")
@@ -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"
``````
@@ -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()
@@ -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.
@@ -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()
@@ -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:**
@@ -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
@@ -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
if isinstance(audio_np, list):
audio_np = audio_np[0] if audio_np else None
if isinstance(audio_np, torch.Tensor):
audio_np = audio_np.detach().cpu().float().numpy()
if isinstance(audio_np, np.ndarray):
audio_np = np.squeeze(audio_np).astype(np.float32)
audio_out = None
if audio is not None:
if isinstance(audio, list):
audio = audio[0] if audio else None
if isinstance(audio, np.ndarray):
audio = torch.from_numpy(audio)
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(
frames_u8,
audio_np,
fps=float(args.fps),
audio_sample_rate=args.audio_sample_rate,
encode_video(
video_tensor,
fps=args.fps,
audio=audio_out,
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}")
@@ -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.
@@ -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()
@@ -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()
@@ -24,9 +24,7 @@ dreamid_omni/
│ ├── models_t5_umt5-xxl-enc-bf16.pth
│ ├── Wan2.2_VAE.pth
│
├── model_index.json
└── transformer/
└── config.json # create by download_dreamid_omni.py
├── model_index.json # create by download_dreamid_omni.py
```
### Run the Inference
@@ -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_audio = output.images[0][1]
generated_video = output[0].images[0][0]
generated_audio = output[0].images[0][1]
try:
from dreamid_omni.utils.io_utils import save_video
except Exception as e:
@@ -125,7 +125,7 @@ def generate_image(
def main():
parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
parser.add_argument("--prompt", "-p", default="<|im_start|>A cute cat<|im_end|>", help="Text prompt")
parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
@@ -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.
@@ -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()
@@ -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) |
@@ -34,45 +34,12 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "input_reference=@/path/to/input.png" \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "negative_prompt=low quality, blurry, static" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42" \
-o sync_i2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
-131
View File
@@ -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
@@ -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
@@ -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()
@@ -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`
@@ -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)
+2 -8
View File
@@ -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,
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
"file_size": 1024000
}
}
```
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 "speaker_description=warm narrator"
-F "ref_text=The exact transcript of the audio sample."
```
### Endpoint
@@ -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,
}
@@ -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
@@ -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:
@@ -34,43 +34,12 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `POST /v1/videos`: create a video generation job
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42" \
-o sync_t2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
-5
View File
@@ -177,24 +177,20 @@ markers = [
"rocm: Tests that run on AMD/ROCm (auto-added)",
"xpu: Tests that run on XPU (auto-added)",
"npu: Tests that run on NPU/Ascend (auto-added)",
"musa: Tests that run on MUSA/Moore Threads (auto-added)",
# specified computation resources marks (auto-added)
"H100: Tests that require H100 GPU",
"L4: Tests that require L4 GPU",
"MI325: Tests that require MI325 GPU (AMD/ROCm)",
"S5000: Tests that require S5000 GPU (Moore Threads/MUSA)",
"A2: Tests that require A2 NPU",
"A3: Tests that require A3 NPU",
"distributed_cuda: Tests that require multi cards on CUDA platform",
"distributed_rocm: Tests that require multi cards on ROCm platform",
"distributed_xpu: Tests that require multi cards on XPU platform",
"distributed_npu: Tests that require multi cards on NPU platform",
"distributed_musa: Tests that require multi cards on MUSA platform",
"skipif_cuda: Skip if the num of CUDA cards is less than the required",
"skipif_rocm: Skip if the num of ROCm cards is less than the required",
"skipif_xpu: Skip if the num of XPU cards is less than the required",
"skipif_npu: Skip if the num of NPU cards is less than the required",
"skipif_musa: Skip if the num of MUSA cards is less than the required",
# more detailed markers
"slow: Slow tests (may skip in quick CI)",
"benchmark: Benchmark tests",
@@ -223,4 +219,3 @@ extend-ignore-identifiers-re = [
ue = "ue"
semantics = "semantics"
fullset = "fullset"
Vai = "Vai"
-1
View File
@@ -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
-4
View File
@@ -1,4 +0,0 @@
-r common.txt
# MUSA platform dependencies
torchada>=0.1.46
onnxruntime>=1.23.2
+4 -15
View File
@@ -3,7 +3,7 @@ Setup script for vLLM-Omni with hardware-dependent installation.
This setup.py implements platform-aware dependency routing so users can run
`pip install vllm-omni` and automatically receive the correct platform-specific
dependencies (CUDA/ROCm/CPU/XPU/NPU/MUSA) without requiring extras like `[cuda]`.
dependencies (CUDA/ROCm/CPU/XPU/NPU) without requiring extras like `[cuda]`.
"""
import os
@@ -46,16 +46,16 @@ def detect_target_device() -> str:
Priority order:
1. VLLM_OMNI_TARGET_DEVICE environment variable (highest priority)
2. Torch backend detection (cuda, rocm, npu, xpu, musa)
2. Torch backend detection (cuda, rocm, npu, xpu)
3. CPU fallback (default)
Returns:
str: Device name ('cuda', 'rocm', 'npu', 'xpu', 'musa', or 'cpu')
str: Device name ('cuda', 'rocm', 'npu', 'xpu', or 'cpu')
"""
# Priority 1: Explicit override via environment variable
target_device = os.environ.get("VLLM_OMNI_TARGET_DEVICE")
if target_device:
valid_devices = ["cuda", "rocm", "npu", "xpu", "musa", "cpu"]
valid_devices = ["cuda", "rocm", "npu", "xpu", "cpu"]
if target_device.lower() in valid_devices:
print(f"Using target device from VLLM_OMNI_TARGET_DEVICE: {target_device.lower()}")
return target_device.lower()
@@ -97,15 +97,6 @@ def detect_target_device() -> str:
except Exception:
pass
# Check for MUSA (Moore Threads)
if hasattr(torch, "musa"):
try:
if torch.musa.is_available():
print("Detected MUSA backend from torch")
return "musa"
except Exception:
pass
print("No GPU backend detected in torch, defaulting to CPU")
return "cpu"
@@ -161,8 +152,6 @@ def get_vllm_omni_version() -> str:
version += f"{sep}npu"
elif device == "xpu":
version += f"{sep}xpu"
elif device == "musa":
version += f"{sep}musa"
elif device == "cpu":
version += f"{sep}cpu"
else:
-67
View File
@@ -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"])
-54
View File
@@ -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"])
+52 -110
View File
@@ -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
@@ -49,11 +48,9 @@ from vllm.utils.network_utils import get_open_port
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.outputs import OmniRequestOutput
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 +336,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:
@@ -1068,7 +1065,7 @@ def convert_audio_to_text(audio_data):
Convert base64 encoded audio data to text using speech recognition.
"""
audio_data = base64.b64decode(audio_data)
output_path = f"./test_{uuid.uuid4().hex}.wav"
output_path = f"./test_{int(time.time())}.wav"
with open(output_path, "wb") as audio_file:
audio_file.write(audio_data)
@@ -1089,44 +1086,11 @@ 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
# Multi-GPU: use last visible device to avoid colliding with default device 0; single device uses 0.
device_index = None
if current_omni_platform.is_available():
n = current_omni_platform.get_device_count()
if n == 1:
device_index = 0
elif n > 1:
device_index = n - 1
if device_index is not None:
torch_device = current_omni_platform.get_torch_device(device_index)
current_omni_platform.set_device(torch_device)
device = str(torch_device)
use_accelerator = True
else:
use_accelerator = False
device = "cpu"
with _serialize_whisper_small_model_download():
model = whisper.load_model("small", device=device)
# Keep Whisper on CPU to avoid consuming GPU memory in tests.
model = whisper.load_model("small", device="cpu")
try:
text = model.transcribe(
output_path,
@@ -1137,15 +1101,13 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
finally:
del model
gc.collect()
if use_accelerator:
current_omni_platform.synchronize()
current_omni_platform.empty_cache()
return text or ""
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 +1126,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 +1722,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 +1751,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 +1817,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 +1829,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(str(kw).lower() in text_lower for kw in keywords), (
assert any(keyword in response.text_content.lower() for keyword in keywords), (
"The output does not contain any of the keywords."
)
else:
if keywords:
audio_lower = response.audio_content.lower()
assert any(str(kw).lower() in audio_lower for kw in keywords), (
assert any(keyword in response.audio_content.lower() for keyword in keywords), (
"The output does not contain any of the keywords."
)
@@ -1950,12 +1888,24 @@ 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(
response.audio_bytes,
request_config.get("voice"),
)
voice = (request_config.get("voice") or "").lower()
if voice and response.audio_bytes:
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}'"
)
def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None):
@@ -2071,11 +2021,7 @@ class OpenAIClientHandler:
if audio_data or text_content:
if audio_data:
merged_seg = _merge_base64_audio_to_segment(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)
audio_content = merge_base64_and_convert_to_text(audio_data)
if audio_content and text_content:
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
@@ -2130,8 +2076,7 @@ class OpenAIClientHandler:
if audio_data or text_content:
if audio_data:
result.audio_bytes = base64.b64decode(audio_data)
audio_content = convert_audio_bytes_to_text(result.audio_bytes)
audio_content = convert_audio_to_text(audio_data)
if audio_content and text_content:
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
@@ -2300,9 +2245,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.
Optional top-level ``speaker`` (str): Qwen3-Omni preset TTS speaker name; sent as
``extra_body["speaker"]`` to ``chat.completions.create``.
extraction (merged with any existing ``extra_body`` / ``mm_processor_kwargs``).
Optional ``extra_body`` (dict): passed through to ``chat.completions.create`` after merge.
request_num: Number of requests, defaults to 1 (single request)
Returns:
@@ -2314,8 +2258,9 @@ class OpenAIClientHandler:
modalities = request_config.get("modalities", ["text", "audio"])
extra_body: dict[str, Any] = {}
if "speaker" in request_config:
extra_body["speaker"] = request_config["speaker"]
raw_extra = request_config.get("extra_body")
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 +2292,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] = {
"model": request_config.get("model"),
"messages": request_config.get("messages"),
"modalities": modalities,
"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)
chat_completion = self.client.chat.completions.create(
model=request_config.get("model"),
messages=request_config.get("messages"),
modalities=modalities,
stream=stream,
)
if stream:
response = self._process_stream_omni_response(chat_completion)
else:
Loaded 100 of 392 files, more files were not shown because too many files have changed in this diff. Show more