Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
171826dd8b |
No files matched your search
@@ -90,18 +90,10 @@ upload_pipeline() {
|
||||
FAIL_FAST=$(fail_fast)
|
||||
|
||||
cd .buildkite
|
||||
|
||||
# Select test definition file: merge suite for main, ready suite for PRs
|
||||
if [[ $BUILDKITE_BRANCH == "main" ]]; then
|
||||
TEST_YAML="test-amd-merge.yml"
|
||||
else
|
||||
TEST_YAML="test-amd-ready.yaml"
|
||||
fi
|
||||
|
||||
(
|
||||
set -x
|
||||
# Output pipeline.yaml with all blank lines removed
|
||||
minijinja-cli test-template.j2 "$TEST_YAML" \
|
||||
minijinja-cli test-template.j2 test-amd.yaml \
|
||||
-D branch="$BUILDKITE_BRANCH" \
|
||||
-D list_file_diff="$LIST_FILE_DIFF" \
|
||||
-D run_all="$RUN_ALL" \
|
||||
|
||||
@@ -1,51 +0,0 @@
|
||||
#!/bin/bash
|
||||
# vllm-omni Intel bootstrap
|
||||
# Uses static pipeline-intel.yaml for Intel XPU tests
|
||||
# Based on: bootstrap-amd-omni.sh
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
if [[ -z "${DOCS_ONLY_DISABLE:-}" ]]; then
|
||||
DOCS_ONLY_DISABLE=0
|
||||
fi
|
||||
|
||||
upload_pipeline() {
|
||||
echo "--- 🛠 Preparing Intel pipeline"
|
||||
buildkite-agent pipeline upload .buildkite/pipeline-intel.yaml
|
||||
}
|
||||
|
||||
get_diff() {
|
||||
$(git add .)
|
||||
echo $(git diff --name-only --diff-filter=ACMDR $(git merge-base origin/main HEAD))
|
||||
}
|
||||
|
||||
get_diff_main() {
|
||||
$(git add .)
|
||||
echo $(git diff --name-only --diff-filter=ACMDR HEAD~1)
|
||||
}
|
||||
|
||||
file_diff=$(get_diff)
|
||||
if [[ $BUILDKITE_BRANCH == "main" ]]; then
|
||||
file_diff=$(get_diff_main)
|
||||
fi
|
||||
|
||||
# ----------------------------------------------------------------------
|
||||
# Early exit: skip pipeline if all changed files are under docs/
|
||||
# ----------------------------------------------------------------------
|
||||
if [[ "${DOCS_ONLY_DISABLE}" != "1" ]] && [[ -n "${file_diff:-}" ]]; then
|
||||
docs_only=1
|
||||
while IFS= read -r f; do
|
||||
[[ -z "$f" ]] && continue
|
||||
if [[ "$f" != docs/* ]]; then
|
||||
docs_only=0
|
||||
break
|
||||
fi
|
||||
done < <(printf '%s\n' "$file_diff" | tr ' ' '\n' | tr -d '\r')
|
||||
|
||||
if [[ "$docs_only" -eq 1 ]]; then
|
||||
buildkite-agent annotate ":memo: CI skipped — docs only" --style "info"
|
||||
exit 0
|
||||
fi
|
||||
fi
|
||||
|
||||
upload_pipeline
|
||||
@@ -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,16 +0,0 @@
|
||||
# .buildkite/pipeline-intel.yaml
|
||||
steps:
|
||||
- group: "INTEL Tests"
|
||||
steps:
|
||||
- label: "Intel: XPU Test"
|
||||
command: "bash .buildkite/scripts/hardware_ci/run-xpu-test.sh"
|
||||
agents:
|
||||
queue: "intel-gpu-omni"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.19.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: true
|
||||
+4
-19
@@ -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"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
@@ -38,17 +23,17 @@ steps:
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
if: __UPLOAD_MERGE_IF__
|
||||
if: build.branch == "main"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild)
|
||||
# L4 Test
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")'
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
|
||||
@@ -1,38 +0,0 @@
|
||||
steps:
|
||||
- label: ":docker: Build image"
|
||||
key: image-build
|
||||
commands:
|
||||
- "aws ecr-public get-login-password --region us-east-1 | docker login --username AWS --password-stdin public.ecr.aws/q9t5s3a7"
|
||||
- "docker build --progress=plain --file docker/Dockerfile.ci -t vllm-omni-ci ."
|
||||
- "docker tag vllm-omni-ci public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
- "docker push public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L2 Test
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L3 Test
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L4 Test
|
||||
- label: "Upload Nightly Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -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}")
|
||||
@@ -64,22 +64,15 @@ while true; do
|
||||
done
|
||||
|
||||
echo "--- Pulling container"
|
||||
## Temporary change to use AMD Docker Hub to store the vllm-ci image
|
||||
# to bypass the rate limit issue with ECR Public Gallery.
|
||||
# TODO: @tjtanaa point back to ECR Public Gallery
|
||||
# once the amd agents are configured to use ECR Public Gallery.
|
||||
# image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
|
||||
image_name="rocm/vllm-ci:${BUILDKITE_COMMIT}-rocm-omni"
|
||||
image_name="public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:${BUILDKITE_COMMIT}-rocm-omni"
|
||||
container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | head -c 10; echo)"
|
||||
|
||||
# TODO: @tjtanaa uncomment this once the amd agents are configured to use ECR Public Gallery.
|
||||
# # Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
|
||||
# sudo apt-get update && sudo apt-get install -y awscli
|
||||
## Use safe docker login helper to prevent race conditions
|
||||
# source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
|
||||
# safe_docker_login_ecr_public
|
||||
|
||||
## Pull the container from AMD Docker Hub
|
||||
# Install AWS CLI to authenticate to ECR Public Gallery to get higher rate limit for pulling images
|
||||
sudo apt-get update && sudo apt-get install -y awscli
|
||||
# Use safe docker login helper to prevent race conditions
|
||||
source "$(dirname "${BASH_SOURCE[0]}")/../docker_login_ecr_public.sh"
|
||||
safe_docker_login_ecr_public
|
||||
# Pull the container from ECR Public Gallery
|
||||
|
||||
docker pull "${image_name}"
|
||||
|
||||
@@ -94,11 +87,7 @@ HF_CACHE="$(realpath ~)/huggingface"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
if [[ -n "${TEST_COMMAND:-}" ]]; then
|
||||
commands="$TEST_COMMAND"
|
||||
else
|
||||
commands="$@"
|
||||
fi
|
||||
commands=$@
|
||||
echo "Commands:$commands"
|
||||
|
||||
PARALLEL_JOB_COUNT=8
|
||||
@@ -113,7 +102,6 @@ if [[ -z "$render_gid" ]]; then
|
||||
fi
|
||||
|
||||
# check if the command contains shard flag, we will run all shards in parallel because the host have 8 GPUs.
|
||||
# TODO: @tjtanaa reenable to run VLLM_ROCM_USE_AITER=1 when AITER is shipped with prebuilt kernels.
|
||||
if [[ $commands == *"--shard-id="* ]]; then
|
||||
# assign job count as the number of shards used
|
||||
commands=$(echo "$commands" | sed -E "s/--num-shards[[:blank:]]*=[[:blank:]]*[0-9]*/--num-shards=${PARALLEL_JOB_COUNT} /g" | sed 's/ \\ / /g')
|
||||
@@ -130,9 +118,11 @@ if [[ $commands == *"--shard-id="* ]]; then
|
||||
--rm \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-e VLLM_ROCM_USE_AITER=1 \
|
||||
-e HIP_VISIBLE_DEVICES="${GPU}" \
|
||||
-e HF_TOKEN \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
@@ -163,8 +153,10 @@ else
|
||||
--rm \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-e VLLM_ROCM_USE_AITER=1 \
|
||||
-e HF_TOKEN \
|
||||
-e AWS_ACCESS_KEY_ID \
|
||||
-e AWS_SECRET_ACCESS_KEY \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
-e "HF_HOME=${HF_MOUNT}" \
|
||||
-e "PYTHONPATH=${MYPYTHONPATH}" \
|
||||
|
||||
@@ -1,56 +0,0 @@
|
||||
#!/bin/bash
|
||||
|
||||
# This script build the XPU docker image and run the offline inference inside the container.
|
||||
set -ex
|
||||
|
||||
omni_source_dir=$(git rev-parse --show-toplevel)
|
||||
|
||||
base_image_name="xpu/vllm-omni-ci-base:${VLLM_VERSION:?VLLM_VERSION must be set}"
|
||||
image_name="xpu/vllm-omni-ci:${BUILDKITE_COMMIT:?BUILDKITE_COMMIT must be set}"
|
||||
container_name="xpu_${BUILDKITE_COMMIT}_$(
|
||||
tr -dc A-Za-z0-9 </dev/urandom | head -c 10
|
||||
echo
|
||||
)"
|
||||
|
||||
cd "${omni_source_dir}"
|
||||
if [ -z "$(docker images -q "${base_image_name}")" ]; then
|
||||
docker build --target vllm-base -t "${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -f docker/Dockerfile.xpu .
|
||||
fi
|
||||
|
||||
# Try building the docker image
|
||||
docker build --build-arg "VLLM_BASE=${base_image_name}" --build-arg "VLLM_VERSION=${VLLM_VERSION}" -t "${image_name}" -f docker/Dockerfile.xpu .
|
||||
|
||||
# Setup cleanup
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || true
|
||||
docker image rm -f "${image_name}" || true
|
||||
docker system prune -f || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
HF_CACHE="${HF_CACHE:-$(realpath ~)/.cache/huggingface}"
|
||||
mkdir -p "${HF_CACHE}"
|
||||
HF_MOUNT="/root/.cache/huggingface"
|
||||
|
||||
time timeout -k 30 30m docker run \
|
||||
--device /dev/dri:/dev/dri \
|
||||
--net=host \
|
||||
--ipc=host \
|
||||
-v /dev/dri/by-path:/dev/dri/by-path \
|
||||
-v "${HF_CACHE}:${HF_MOUNT}" \
|
||||
--security-opt seccomp=unconfined \
|
||||
--entrypoint="" \
|
||||
-e VLLM_LOGGING_LEVEL \
|
||||
-e VLLM_OMNI_LOGGING_LEVEL \
|
||||
-e HF_TOKEN \
|
||||
-e ZE_AFFINITY_MASK \
|
||||
--name "${container_name}" \
|
||||
"${image_name}" \
|
||||
bash -c '
|
||||
set -e
|
||||
echo $ZE_AFFINITY_MASK
|
||||
pip install tblib==3.1.0
|
||||
cd /workspace/vllm-omni
|
||||
pytest -v -s -m "core_model and xpu and B60"
|
||||
pytest -v -s -m "advanced_model and xpu and B60"
|
||||
'
|
||||
@@ -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
|
||||
@@ -1,215 +0,0 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export GPU_ARCHS=gfx942
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
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 -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
|
||||
'
|
||||
|
||||
- label: "Qwen3-TTS Base E2E Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
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 -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
|
||||
'
|
||||
|
||||
- label: "Diffusion Image Edit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
|
||||
# split Bagel Model Test with H100 (Real Weights) into three tests
|
||||
- label: "Bagel Text2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory" -k "rocm"
|
||||
|
||||
- label: "Bagel Img2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Online Serving Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
@@ -1,179 +0,0 @@
|
||||
steps:
|
||||
|
||||
- label: "Simple Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
# agent_pool: mi325_1
|
||||
# depends_on: amd-build
|
||||
# mirror_hardwares: [amdproduction]
|
||||
# grade: Blocking
|
||||
# commands:
|
||||
# - export GPU_ARCHS=gfx942
|
||||
# - export VLLM_LOGGING_LEVEL=DEBUG
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_model.py
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 17m pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- timeout 10m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "Diffusion Image Edit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/e2e/online_serving/test_image_gen_edit.py
|
||||
|
||||
- label: "Bagel Text2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Img2Img Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "Bagel Online Serving Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
- export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
- timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
@@ -1,7 +1,7 @@
|
||||
steps:
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 20
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -11,7 +11,7 @@ steps:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 20
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
@@ -101,7 +101,6 @@ steps:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py
|
||||
|
||||
|
||||
+53
-170
@@ -12,12 +12,11 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "advanced_model and diffusion" --run-level "advanced_model"
|
||||
@@ -30,12 +29,11 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
timeout_in_minutes: 30
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_images_generations_lora.py
|
||||
@@ -48,7 +46,6 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -56,8 +53,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:
|
||||
@@ -67,7 +64,6 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -85,7 +81,6 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -103,7 +98,6 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -111,7 +105,7 @@ 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
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
@@ -122,7 +116,6 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -141,7 +134,6 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -160,64 +152,24 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Engine Test"
|
||||
- label: "Benchmark & Engine Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 15m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 30
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
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"
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -229,24 +181,20 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Qwen3-TTS Base E2E Test"
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 15
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
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"
|
||||
'
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
@@ -258,7 +206,8 @@ 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"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -277,11 +226,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:
|
||||
@@ -317,11 +261,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:
|
||||
@@ -333,93 +272,37 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Bagel Model Test with H100 (Real Weights)"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 55m bash -c '
|
||||
set -e
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-merge-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# timeout_in_minutes: 30
|
||||
# depends_on: upload-merge-pipeline
|
||||
# commands:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
@@ -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
|
||||
+62
-218
@@ -1,201 +1,17 @@
|
||||
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: "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
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
- label: ":bar_chart: Testcase Statistics"
|
||||
key: nightly-testcase-statistics
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
- |
|
||||
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=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -214,11 +30,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:
|
||||
@@ -230,25 +41,58 @@ steps:
|
||||
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
|
||||
- label: "Omni Model Test"
|
||||
timeout_in_minutes: 60
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- pip install openpyxl
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-omni-performance
|
||||
- python tools/nightly/generate_nightly_perf_excel.py
|
||||
- python tools/nightly/generate_nightly_perf_html.py
|
||||
- python tools/nightly/send_nightly_email.py --report-file "tests/dfx/perf/results/*.xlsx, tests/dfx/perf/results/*.html"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.xlsx"
|
||||
- buildkite-agent artifact upload "tests/dfx/perf/results/*.html"
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/examples/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
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
|
||||
shm-size: "8gb"
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Perf Test"
|
||||
timeout_in_minutes: 120
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/perf/scripts/run_benchmark.py
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
+117
-401
@@ -5,39 +5,6 @@ steps:
|
||||
- "timeout 20m pytest -v -s -m 'core_model and cpu' --cov=vllm_omni --cov-branch --cov-report=term-missing --cov-report=html --cov-report=xml"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "CUDA Unit Test with single card"
|
||||
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
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "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
|
||||
@@ -51,41 +18,7 @@ steps:
|
||||
- label: "Diffusion Model Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion Batching Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_qwen_image_diffusion_batching.py -m "core_model and diffusion" --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion" --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
@@ -101,7 +34,15 @@ steps:
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- timeout 10m pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
|
||||
EXIT2=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
@@ -111,7 +52,6 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -128,7 +68,6 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -145,7 +84,6 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -163,21 +101,42 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: "Engine Test"
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 15m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
'
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.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
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
|
||||
- label: "Benchmark & Engine Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 15m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
set +e
|
||||
pytest -s -v tests/benchmarks/test_serve_cli.py
|
||||
EXIT1=$$?
|
||||
pytest -s -v tests/engine/test_async_omni_engine_abort.py
|
||||
EXIT2=$$?
|
||||
exit $$((EXIT1 | EXIT2))
|
||||
'
|
||||
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
|
||||
@@ -186,7 +145,6 @@ steps:
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
@@ -198,7 +156,7 @@ steps:
|
||||
timeout 17m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "core_model" --run-level "core_model"
|
||||
pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
@@ -209,89 +167,47 @@ steps:
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni Model Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "MiMo-Audio E2E Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 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: "Omni Model Test with H100"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - |
|
||||
# timeout 20m bash -c '
|
||||
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# export VLLM_TEST_CLEAN_GPU_MEMORY="1"
|
||||
# # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
# pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 2
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
@@ -300,35 +216,10 @@ steps:
|
||||
timeout 20m bash -c '
|
||||
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 -m "core_model" --run-level "core_model"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts.py
|
||||
'
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
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"
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -340,45 +231,6 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "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
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# - label: "Diffusion Image Edit Test with H100 (1 GPU)"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
@@ -416,175 +268,39 @@ steps:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
|
||||
- label: "Bagel Text2Img Model Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Bagel Img2Img Model Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Bagel Online Serving Test with H100"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- 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
|
||||
# - label: "Bagel Text2Img Model Test with H100"
|
||||
# depends_on: upload-ready-pipeline
|
||||
# commands:
|
||||
# - |
|
||||
# timeout 30m bash -c '
|
||||
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py
|
||||
# '
|
||||
# agents:
|
||||
# queue: "mithril-h100-pool"
|
||||
# plugins:
|
||||
# - kubernetes:
|
||||
# podSpec:
|
||||
# containers:
|
||||
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
# resources:
|
||||
# limits:
|
||||
# nvidia.com/gpu: 1
|
||||
# volumeMounts:
|
||||
# - name: devshm
|
||||
# mountPath: /dev/shm
|
||||
# - name: hf-cache
|
||||
# mountPath: /root/.cache/huggingface
|
||||
# env:
|
||||
# - name: HF_HOME
|
||||
# value: /root/.cache/huggingface
|
||||
# nodeSelector:
|
||||
# node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
# volumes:
|
||||
# - name: devshm
|
||||
# emptyDir:
|
||||
# medium: Memory
|
||||
# - name: hf-cache
|
||||
# hostPath:
|
||||
# path: /mnt/hf-cache
|
||||
# type: DirectoryOrCreate
|
||||
@@ -3,7 +3,7 @@
|
||||
Last synced: 2025-12-15
|
||||
Modifications: Removed unused CUDA/NVIDIA logic, keeping only AMD tests
|
||||
#}
|
||||
{% set docker_image_amd = "rocm/vllm-ci:$BUILDKITE_COMMIT-rocm-omni" %}
|
||||
{% set docker_image_amd = "public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT-rocm-omni" %}
|
||||
{% set default_working_dir = "/app/vllm-omni" %}
|
||||
|
||||
- group: "AMD Tests"
|
||||
@@ -13,6 +13,7 @@
|
||||
depends_on: ~
|
||||
soft_fail: false
|
||||
commands:
|
||||
- "source .buildkite/scripts/docker_login_ecr_public.sh && safe_docker_login_ecr_public"
|
||||
- "docker build -f docker/Dockerfile.rocm -t {{ docker_image_amd }} --target test --build-arg ARG_PYTORCH_ROCM_ARCH=gfx942 --progress plain ."
|
||||
- "docker push {{ docker_image_amd }}"
|
||||
key: "amd-build"
|
||||
@@ -21,15 +22,13 @@
|
||||
retry:
|
||||
automatic:
|
||||
- exit_status: -1 # Agent was lost
|
||||
limit: 2
|
||||
limit: 1
|
||||
- exit_status: -10 # Agent was lost
|
||||
limit: 2
|
||||
- exit_status: 128 # Git connectivity issues
|
||||
limit: 2
|
||||
limit: 1
|
||||
- exit_status: 1 # Machine occasionally fail
|
||||
limit: 1
|
||||
agents:
|
||||
queue: amd-cpu
|
||||
queue: cpu_queue_premerge
|
||||
|
||||
{% for step in steps %}
|
||||
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
|
||||
@@ -41,14 +40,9 @@
|
||||
{% else %}
|
||||
queue: amd_mi325_1
|
||||
{% endif %}
|
||||
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
|
||||
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh "(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }} ; {{ step.command or (step.commands | join(" && ")) | safe }}"
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
TEST_COMMAND: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{{ indented_cmd | safe }}
|
||||
priority: 100
|
||||
{% if step.grade and step.grade == "Blocking" %}
|
||||
soft_fail: false
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
@@ -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.
|
||||
@@ -74,21 +74,28 @@ body:
|
||||
If relevant, add a minimal example so that we can reproduce the error by running the code. It is very important for the snippet to be as succinct (minimal) as possible, so please take time to trim down any irrelevant code to help us debug efficiently. We are going to copy-paste your code and we expect to get the same result as you did: avoid any external data, and include the relevant imports, etc. For example:
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm import SamplingParams
|
||||
from vllm_omni import OmniLLM, create_ar_stage_config, create_dit_stage_config
|
||||
|
||||
omni = Omni(
|
||||
model="Qwen/Qwen-Image",
|
||||
stage_configs_path="/path/to/stage_configs.yaml",
|
||||
# Create stage configurations
|
||||
ar_config = create_ar_stage_config(
|
||||
stage_id=0,
|
||||
model_path="Qwen/Qwen3-0.6B",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["text"]
|
||||
)
|
||||
|
||||
prompts = [{"prompt": "A scenic watercolor painting of a lighthouse at sunset"}]
|
||||
sampling_params_list = [
|
||||
SamplingParams(max_tokens=1),
|
||||
OmniDiffusionSamplingParams(num_outputs_per_prompt=1),
|
||||
]
|
||||
outputs = omni.generate(prompts=prompts, sampling_params_list=sampling_params_list)
|
||||
dit_config = create_dit_stage_config(
|
||||
stage_id=1,
|
||||
model_path="stabilityai/stable-diffusion-2-1",
|
||||
input_modalities=["text"],
|
||||
output_modalities=["image"]
|
||||
)
|
||||
|
||||
# Initialize OmniLLM
|
||||
omni = OmniLLM([ar_config, dit_config])
|
||||
|
||||
# Generate
|
||||
outputs = omni.generate(prompt="A scenic watercolor painting of a lighthouse at sunset")
|
||||
```
|
||||
|
||||
If the code is too long (hopefully, it isn't), feel free to put it in a public gist and link it in the issue: https://gist.github.com.
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -39,10 +39,10 @@ jobs:
|
||||
python-version: ["3.11", "3.12"]
|
||||
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/checkout@v4
|
||||
|
||||
- name: Set up Python ${{ matrix.python-version }}
|
||||
uses: actions/setup-python@v6
|
||||
uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: ${{ matrix.python-version }}
|
||||
|
||||
@@ -57,7 +57,7 @@ jobs:
|
||||
bash scripts/build_wheel.sh --python python
|
||||
|
||||
- name: Upload Wheel Artifact
|
||||
uses: actions/upload-artifact@v7
|
||||
uses: actions/upload-artifact@v4
|
||||
with:
|
||||
name: vllm-omni-wheel-py${{ matrix.python-version }}
|
||||
path: dist/*.whl
|
||||
@@ -24,12 +24,12 @@ jobs:
|
||||
github.event.comment.author_association == 'OWNER')
|
||||
steps:
|
||||
- name: Checkout repository
|
||||
uses: actions/checkout@v6.0.2
|
||||
uses: actions/checkout@v4.2.2
|
||||
with:
|
||||
fetch-depth: 0
|
||||
|
||||
- name: Set up Python
|
||||
uses: actions/setup-python@v6.2.0
|
||||
uses: actions/setup-python@v5.3.0
|
||||
with:
|
||||
python-version: '3.11'
|
||||
cache: 'pip'
|
||||
@@ -54,7 +54,7 @@ jobs:
|
||||
|
||||
- name: Upload review logs
|
||||
if: always()
|
||||
uses: actions/upload-artifact@v7.0.0
|
||||
uses: actions/upload-artifact@v4.4.3
|
||||
with:
|
||||
name: pr-review-logs-${{ github.event.issue.number || github.event.pull_request.number }}
|
||||
path: /tmp/pr_review_*.log
|
||||
|
||||
@@ -18,8 +18,8 @@ jobs:
|
||||
pre-commit:
|
||||
runs-on: ubuntu-latest
|
||||
steps:
|
||||
- uses: actions/checkout@v6
|
||||
- uses: actions/setup-python@v6
|
||||
- uses: actions/checkout@v4
|
||||
- uses: actions/setup-python@v5
|
||||
with:
|
||||
python-version: "3.12"
|
||||
cache: 'pip'
|
||||
|
||||
+1
-16
@@ -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
|
||||
|
||||
@@ -8,7 +8,6 @@ repos:
|
||||
# list of supported hooks: https://pre-commit.com/hooks.html
|
||||
- id: check-yaml
|
||||
args: ["--unsafe"]
|
||||
exclude: "examples/online_serving/chart-helm/templates/"
|
||||
- id: debug-statements
|
||||
- id: end-of-file-fixer
|
||||
- id: mixed-line-ending
|
||||
|
||||
@@ -1,6 +1,6 @@
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" srcset="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png">
|
||||
<img alt="vllm-omni" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/docs/source/logos/vllm-omni-logo.png" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
@@ -9,16 +9,14 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
</h3>
|
||||
|
||||
<p align="center">
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1"><b>Slides</b></a> |
|
||||
| <a href="https://vllm-omni.readthedocs.io/en/latest/"><b>Documentation</b></a> | <a href="https://discuss.vllm.ai"><b>User Forum</b></a> | <a href="https://slack.vllm.ai"><b>Developer Slack</b></a> | <a href="docs/assets/WeChat.jpg"><b>WeChat</b></a> | <a href="https://arxiv.org/abs/2602.02204"><b>Paper</b></a> | <a href="https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true"><b>Slides</b></a> |
|
||||
</p>
|
||||
|
||||
|
||||
---
|
||||
|
||||
*Latest News* 🔥
|
||||
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
|
||||
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
|
||||
- [2026/03] **[vllm-omni-skills](https://github.com/hsliuustc0106/vllm-omni-skills)** is a community-driven collection of AI assistant skills that help developers work with vLLM-Omni more effectively. These skills can be used with popular agentic AI coding assistants like **Cursor IDE**, **Claude**, **Codex**, and more.
|
||||
|
||||
- [2026/02] We released [0.16.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.16.0) - A major alignment + capability release that rebases onto **upstream vLLM v0.16.0** and significantly expands performance, distributed execution, and production readiness across **Qwen3-Omni / Qwen3-TTS**, **Bagel**, **MiMo-Audio**, **GLM-Image** and the **Diffusion (DiT) image/video stack**—while also improving platform coverage (CUDA / ROCm / NPU / XPU), CI quality, and documentation.
|
||||
- [2026/02] We released [0.14.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0) - This is the first **stable release** of vLLM-Omni that expands Omni’s diffusion / image-video generation and audio / TTS stack, improves distributed execution and memory efficiency, and broadens platform/backend coverage (GPU/ROCm/NPU/XPU). It also brings meaningful upgrades to serving APIs, profiling & benchmarking, and overall stability. Please check our latest [paper](https://arxiv.org/abs/2602.02204) for architecture design and performance results.
|
||||
- [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm).
|
||||
|
||||
@@ -47,7 +47,6 @@ If no, check your shell running the ComfyUI process. There may be some error mes
|
||||
This extension offers the following nodes based on the output modalities (at **ComfyUI sidebar -> Node Library**):
|
||||
|
||||
- **Generate Image** for text-to-image and image-to-image tasks
|
||||
- **Generate Video** for text-to-video and image-to-video tasks
|
||||
- **Multimodality Understanding** for multimodality-to-text and multimodality-to-audio tasks
|
||||
- **TTS** and **TTS Voice Clone** for TTS tasks
|
||||
|
||||
@@ -60,7 +59,7 @@ To build a simple workflow yourself,
|
||||
|
||||
- Drag a generation node onto the canvas.
|
||||
- Depending on your need, grab built-in multimedia file loader nodes, such as **image->Load Image**, **image->video->Load Video**, **audio->Load Audio**
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**, **utils->Preview as Text**.
|
||||
- Depending on your need, grab built-in multimedia file preview nodes, such as **image->Preview Image**, **image->video->Save Video**, **audio->Preview Audio**. For text output, you can install [ComfyUI-Custom-Scripts plugin](https://github.com/pythongosssss/ComfyUI-Custom-Scripts/) and grab its **utils->Show Text 🐍** node.
|
||||
- If you want to tune sampling parameters, grab corresponding nodes from **vLLM-Omni-> Sampling Params**.
|
||||
- For multi-stage models, you can connect multiple **AR Sampling Params** and **Diffusion Sampling Params** nodes to a **Multi-Stage Sampling Params List** node, and connect this node to the generation node.
|
||||
- For some multi-stage models like BAGEL, [only one stage's sampling parameters are exposed and tunable via vLLM-Omni's online serving API](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/examples/online_serving/bagel/). Thus, these models are treated as single-stage ones. Please check the vLLM-Omni documentation on how to correctly set each model's sampling parameters.
|
||||
@@ -105,20 +104,6 @@ You can configure per-stage sampling parameters for multi-stage models.
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-image or image-to-image API endpoints depending on whether you connect an image input or not.
|
||||
|
||||
### Text-to-video and image-to-video generation (e.g., Wan)
|
||||
|
||||
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni Video Generation**)
|
||||
|
||||
<p align="center">
|
||||
<picture>
|
||||
<source media="(prefers-color-scheme: dark)" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg">
|
||||
<img alt="vLLM-Omni video generation" src="https://raw.githubusercontent.com/vllm-project/vllm-omni/refs/heads/main/apps/ComfyUI-vLLM-Omni/docs/images/comfyui-video-generation.jpg" width=55%>
|
||||
</picture>
|
||||
</p>
|
||||
|
||||
> [!TIP]
|
||||
> The node automatically choose text-to-video or image-to-video API endpoints depending on whether you connect an image input or not.
|
||||
|
||||
### TTS (e.g., Qwen TTS series)
|
||||
|
||||
(Also available at **ComfyUI sidebar->Template->vLLM-Omni->vLLM-Omni TTS**)
|
||||
|
||||
@@ -14,21 +14,17 @@ from .comfyui_vllm_omni.nodes import (
|
||||
VLLMOmniARSampling,
|
||||
VLLMOmniDiffusionSampling,
|
||||
VLLMOmniGenerateImage,
|
||||
VLLMOmniGenerateVideo,
|
||||
VLLMOmniQwenTTSParams,
|
||||
VLLMOmniRemoteLoRA,
|
||||
VLLMOmniSamplingParamsList,
|
||||
VLLMOmniTTS,
|
||||
VLLMOmniUnderstanding,
|
||||
VLLMOmniVoiceClone,
|
||||
VLLMOmniWanParams,
|
||||
)
|
||||
|
||||
# A dictionary that contains all nodes you want to export with their names
|
||||
NODE_CLASS_MAPPINGS = {
|
||||
# === Generation ===
|
||||
"VLLMOmniGenerateImage": VLLMOmniGenerateImage,
|
||||
"VLLMOmniGenerateVideo": VLLMOmniGenerateVideo,
|
||||
"VLLMOmniUnderstanding": VLLMOmniUnderstanding,
|
||||
"VLLMOmniTTS": VLLMOmniTTS,
|
||||
"VLLMOmniVoiceClone": VLLMOmniVoiceClone,
|
||||
@@ -36,16 +32,13 @@ NODE_CLASS_MAPPINGS = {
|
||||
"VLLMOmniARSampling": VLLMOmniARSampling,
|
||||
"VLLMOmniDiffusionSampling": VLLMOmniDiffusionSampling,
|
||||
"VLLMOmniSamplingParamsList": VLLMOmniSamplingParamsList,
|
||||
"VLLMOmniRemoteLoRA": VLLMOmniRemoteLoRA,
|
||||
"VLLMOmniQwenTTSParams": VLLMOmniQwenTTSParams,
|
||||
"VLLMOmniWanParams": VLLMOmniWanParams,
|
||||
}
|
||||
|
||||
# A dictionary that contains the friendly/humanly readable titles for the nodes
|
||||
NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
# === Generation ===
|
||||
"VLLMOmniGenerateImage": "Generate Image",
|
||||
"VLLMOmniGenerateVideo": "Generate Video",
|
||||
"VLLMOmniUnderstanding": "Multimodality Understanding",
|
||||
"VLLMOmniTTS": "TTS (Text to Speech)",
|
||||
"VLLMOmniVoiceClone": "TTS Voice Cloning",
|
||||
@@ -53,9 +46,7 @@ NODE_DISPLAY_NAME_MAPPINGS = {
|
||||
"VLLMOmniARSampling": "AR Sampling Params",
|
||||
"VLLMOmniDiffusionSampling": "Diffusion Sampling Params",
|
||||
"VLLMOmniSamplingParamsList": "Multi-Stage Sampling Params List",
|
||||
"VLLMOmniRemoteLoRA": "LoRA",
|
||||
"VLLMOmniQwenTTSParams": "Qwen TTS Params",
|
||||
"VLLMOmniWanParams": "Wan Video Params",
|
||||
}
|
||||
|
||||
WEB_DIRECTORY = "./web"
|
||||
@@ -11,7 +11,6 @@ from .utils.types import (
|
||||
AutoregressionSamplingParams,
|
||||
DiffusionSamplingParams,
|
||||
QwenTTSModelSpecificParams,
|
||||
WanModelSpecificParams,
|
||||
)
|
||||
from .utils.validators import (
|
||||
add_sampling_parameters_to_stage,
|
||||
@@ -57,7 +56,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
# "video": ("VIDEO",),
|
||||
# "audio": ("AUDIO",),
|
||||
"sampling_params": ("SAMPLING_PARAMS",),
|
||||
"lora": ("REMOTE_LORA",),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -78,11 +76,9 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
audio: AudioInput | None = None, # Hidden & unused
|
||||
video: VideoInput | None = None, # Hidden & unused
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
lora: dict | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
if image is None and mask is not None:
|
||||
@@ -109,7 +105,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
height=height,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
return (output,)
|
||||
elif image is not None and audio is None and video is None:
|
||||
@@ -124,7 +119,6 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
negative_prompt=negative_prompt,
|
||||
mask=mask,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
return (output,)
|
||||
|
||||
@@ -142,91 +136,11 @@ class VLLMOmniGenerateImage(_VLLMOmniGenerateBase):
|
||||
audio=audio,
|
||||
video=video,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
)
|
||||
|
||||
return (output,)
|
||||
|
||||
|
||||
class VLLMOmniGenerateVideo(_VLLMOmniGenerateBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"url": ("STRING", {"default": "http://localhost:8000/v1"}),
|
||||
"model": ("STRING", {"default": "Wan-AI/Wan2.2-T2V-A14B-Diffusers"}),
|
||||
"prompt": ("STRING", {"multiline": True}),
|
||||
"negative_prompt": ("STRING", {"multiline": True, "default": ""}),
|
||||
"width": ("INT", {"default": 832, "min": 1}),
|
||||
"height": ("INT", {"default": 480, "min": 1}),
|
||||
"fps": ("INT", {"default": 16, "min": 1}),
|
||||
"num_frames": ("INT", {"default": 41, "min": 1}),
|
||||
},
|
||||
"optional": {
|
||||
"image": ("IMAGE",),
|
||||
"sampling_params": ("SAMPLING_PARAMS",),
|
||||
"lora": ("REMOTE_LORA",),
|
||||
"model_params": ("VIDEO_PARAMS",),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO",)
|
||||
RETURN_NAMES = ("video",)
|
||||
FUNCTION = "generate"
|
||||
|
||||
async def generate(
|
||||
self,
|
||||
url: str,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
fps: int,
|
||||
num_frames: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
**kwargs,
|
||||
):
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
logger.debug("Got model params: %s", model_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
|
||||
# Currently, all video generation models are single-stage diffusion models
|
||||
if isinstance(sampling_params, list):
|
||||
if len(sampling_params) != 1:
|
||||
raise ValueError(
|
||||
"Video generation expects a single sampling params group. "
|
||||
"Please provide one Diffusion sampling node."
|
||||
)
|
||||
sampling_params = sampling_params[0]
|
||||
|
||||
if sampling_params is not None:
|
||||
sampling_params.pop("type", None) # internal fields
|
||||
if model_params is not None:
|
||||
model_params.pop("type", None) # internal fields
|
||||
|
||||
client = VLLMOmniClient(url)
|
||||
output = await client.generate_video(
|
||||
model=model,
|
||||
prompt=prompt,
|
||||
image=image, # image present => i2v, absent => t2v
|
||||
width=width,
|
||||
height=height,
|
||||
num_frames=num_frames,
|
||||
fps=fps,
|
||||
negative_prompt=negative_prompt,
|
||||
sampling_params=sampling_params,
|
||||
lora=lora,
|
||||
model_params=model_params,
|
||||
)
|
||||
return (output,)
|
||||
|
||||
|
||||
class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -274,8 +188,7 @@ class VLLMOmniUnderstanding(_VLLMOmniGenerateBase):
|
||||
use_audio_in_video: bool = True,
|
||||
**kwargs,
|
||||
) -> tuple[str, AudioInput]:
|
||||
if kwargs:
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.info("Uncaught kwargs: %s", kwargs)
|
||||
logger.debug("Got sampling params: %s", sampling_params)
|
||||
validate_model_and_sampling_params_types(model, sampling_params)
|
||||
|
||||
@@ -634,52 +547,6 @@ class VLLMOmniSamplingParamsList:
|
||||
return (params,)
|
||||
|
||||
|
||||
class VLLMOmniRemoteLoRA:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"local_path": ("STRING", {"default": ""}),
|
||||
"name": ("STRING", {"default": ""}),
|
||||
"scale": (
|
||||
"FLOAT",
|
||||
{"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.1},
|
||||
),
|
||||
"int_id": (
|
||||
"INT",
|
||||
{
|
||||
"default": 0,
|
||||
"min": 0,
|
||||
"step": 1,
|
||||
"tooltip": "0 means it is not set and the server can derive it.",
|
||||
},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("REMOTE_LORA",)
|
||||
RETURN_NAMES = ("lora",)
|
||||
FUNCTION = "get_lora"
|
||||
CATEGORY = "vLLM-Omni"
|
||||
|
||||
@classmethod
|
||||
def VALIDATE_INPUTS(cls, local_path, name) -> str | Literal[True]:
|
||||
if not local_path.strip() or not name.strip():
|
||||
return "Both local_path and name must be provided."
|
||||
return True
|
||||
|
||||
def get_lora(self, local_path: str, name: str, scale: float, int_id: int):
|
||||
local_path = local_path.strip()
|
||||
name = name.strip()
|
||||
lora = {
|
||||
"local_path": local_path or None,
|
||||
"name": name or None,
|
||||
"scale": float(scale),
|
||||
"int_id": int(int_id) if int_id > 0 else None,
|
||||
}
|
||||
return (lora,)
|
||||
|
||||
|
||||
class VLLMOmniQwenTTSParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -705,32 +572,3 @@ class VLLMOmniQwenTTSParams:
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
return (QwenTTSModelSpecificParams(kwargs),)
|
||||
|
||||
|
||||
class VLLMOmniWanParams:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"guidance_scale_2": (
|
||||
"FLOAT",
|
||||
{"default": 4.0, "min": 0.0, "max": 20.0, "step": 0.1},
|
||||
),
|
||||
"boundary_ratio": (
|
||||
"FLOAT",
|
||||
{"default": 0.875, "min": 0.0, "max": 1.0, "step": 0.001},
|
||||
),
|
||||
"flow_shift": (
|
||||
"FLOAT",
|
||||
{"default": 5.0, "min": 0.0, "max": 100.0, "step": 0.1},
|
||||
),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = ("VIDEO_PARAMS",)
|
||||
RETURN_NAMES = ("Wan video params",)
|
||||
FUNCTION = "get_params"
|
||||
CATEGORY = "vLLM-Omni/Video Params"
|
||||
|
||||
def get_params(self, **kwargs):
|
||||
return (WanModelSpecificParams(kwargs),)
|
||||
@@ -6,8 +6,6 @@ The image generation part is derived from dougbtv/comfyui-vllm-omni by Doug (@do
|
||||
Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed under the MIT License.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
from typing import Any
|
||||
|
||||
import aiohttp
|
||||
@@ -20,7 +18,6 @@ from .format import (
|
||||
base64_to_audio,
|
||||
base64_to_image_tensor,
|
||||
bytes_to_audio,
|
||||
bytes_to_video,
|
||||
image_tensor_to_base64,
|
||||
image_tensor_to_png_bytes,
|
||||
video_to_base64,
|
||||
@@ -32,43 +29,10 @@ from .types import AudioFormat
|
||||
logger = get_logger(__name__)
|
||||
|
||||
|
||||
async def url_json(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> dict[str, Any]:
|
||||
try:
|
||||
async with getattr(session, verb)(url, **kwargs) as response:
|
||||
if not response.ok:
|
||||
error_text = await response.text()
|
||||
raise (ValueError if response.status < 500 else RuntimeError)(
|
||||
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
||||
)
|
||||
try:
|
||||
return await response.json()
|
||||
except aiohttp.ContentTypeError as e:
|
||||
raise RuntimeError(f"Invalid JSON response from vLLM-Omni: {e}")
|
||||
except aiohttp.ClientError as e:
|
||||
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
||||
|
||||
|
||||
async def url_bytes(session: aiohttp.ClientSession, url: str, verb: str = "get", **kwargs) -> bytes:
|
||||
try:
|
||||
async with getattr(session, verb)(url, **kwargs) as response:
|
||||
if not response.ok:
|
||||
error_text = await response.text()
|
||||
raise (ValueError if response.status < 500 else RuntimeError)(
|
||||
f"vLLM-Omni API returned status {response.status}: {error_text}"
|
||||
)
|
||||
return await response.read()
|
||||
except aiohttp.ClientError as e:
|
||||
raise RuntimeError(f"Network error connecting to vLLM-Omni at {url}: {e}")
|
||||
|
||||
|
||||
class VLLMOmniClient:
|
||||
def __init__(
|
||||
self, base_url: str, timeout: float | None = None, poll_interval: float = 5.0, max_poll_duration: float = 60 * 5
|
||||
):
|
||||
def __init__(self, base_url: str, timeout: float = 300.0):
|
||||
self.base_url = base_url
|
||||
self.timeout = aiohttp.ClientTimeout(total=timeout)
|
||||
self.poll_interval = poll_interval
|
||||
self.max_poll_duration = max_poll_duration
|
||||
|
||||
async def generate_image(
|
||||
self,
|
||||
@@ -79,13 +43,12 @@ class VLLMOmniClient:
|
||||
height: int,
|
||||
negative_prompt: str | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run text-to-image generation via DALLE API"""
|
||||
await self._check_model_exist(model)
|
||||
|
||||
size = f"{width}x{height}"
|
||||
payload: dict[str, Any] = {
|
||||
payload = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"size": size,
|
||||
@@ -95,8 +58,6 @@ class VLLMOmniClient:
|
||||
payload["negative_prompt"] = negative_prompt
|
||||
if sampling_params is not None:
|
||||
payload.update(sampling_params)
|
||||
if lora is not None:
|
||||
payload["lora"] = lora
|
||||
logger.debug("img gen payload: %s", payload)
|
||||
|
||||
url = self.base_url + "/images/generations"
|
||||
@@ -149,7 +110,6 @@ class VLLMOmniClient:
|
||||
negative_prompt: str | None = None,
|
||||
mask: torch.Tensor | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Run image editing via DALLE API"""
|
||||
await self._check_model_exist(model)
|
||||
@@ -171,8 +131,6 @@ class VLLMOmniClient:
|
||||
if sampling_params is not None:
|
||||
for k, v in sampling_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if lora is not None:
|
||||
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
|
||||
if mask is not None:
|
||||
mask_filename = "mask.png"
|
||||
form.add_field(
|
||||
@@ -225,7 +183,6 @@ class VLLMOmniClient:
|
||||
audio: AudioInput | None = None,
|
||||
video: VideoInput | None = None,
|
||||
sampling_params: dict | list[dict] | None = None,
|
||||
lora: dict | None = None,
|
||||
) -> torch.Tensor:
|
||||
payload = VLLMOmniClient._prepare_chat_completion_messages(
|
||||
model=model,
|
||||
@@ -236,8 +193,6 @@ class VLLMOmniClient:
|
||||
video=video,
|
||||
sampling_params=sampling_params,
|
||||
modalities=["image"],
|
||||
# === below are additional `extra_body` fields, handled by **kwargs ===
|
||||
lora=lora,
|
||||
)
|
||||
choices = await self._generate_base_chat_completion(model, payload)
|
||||
|
||||
@@ -251,87 +206,6 @@ class VLLMOmniClient:
|
||||
|
||||
return torch.stack(image_tensors, dim=0)
|
||||
|
||||
async def generate_video(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
num_frames: int,
|
||||
fps: int,
|
||||
negative_prompt: str | None = None,
|
||||
image: torch.Tensor | None = None,
|
||||
sampling_params: dict | None = None,
|
||||
model_params: dict | None = None,
|
||||
lora: dict | None = None,
|
||||
**extra_body,
|
||||
) -> VideoInput:
|
||||
form = aiohttp.FormData()
|
||||
form.add_field("model", model)
|
||||
form.add_field("prompt", prompt)
|
||||
form.add_field("width", str(width))
|
||||
form.add_field("height", str(height))
|
||||
form.add_field("num_frames", str(num_frames))
|
||||
form.add_field("fps", str(fps))
|
||||
if negative_prompt:
|
||||
form.add_field("negative_prompt", negative_prompt)
|
||||
if sampling_params is not None:
|
||||
for k, v in sampling_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if model_params is not None:
|
||||
for k, v in model_params.items():
|
||||
form.add_field(k, str(v))
|
||||
if lora is not None:
|
||||
form.add_field("lora", json.dumps(lora, ensure_ascii=False))
|
||||
if extra_body:
|
||||
form.add_field("extra_body", json.dumps(extra_body, ensure_ascii=False))
|
||||
|
||||
if image is not None:
|
||||
image_filename = "image.png" # Required for multipart form
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_tensor_to_png_bytes(image, image_filename),
|
||||
filename=image_filename,
|
||||
content_type="image/png",
|
||||
)
|
||||
|
||||
async with aiohttp.ClientSession(timeout=self.timeout) as session:
|
||||
# Start the video generation job
|
||||
url = f"{self.base_url}/videos"
|
||||
data = await url_json(session, url, "post", data=form)
|
||||
if (job_id := data.get("id", None)) is None:
|
||||
raise RuntimeError("API response missing job 'id' field - expected OpenAI compliant format")
|
||||
if (job_status := data.get("status", None)) is None:
|
||||
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
|
||||
|
||||
# Poll for video generation job completion
|
||||
deadline = asyncio.get_running_loop().time() + self.max_poll_duration
|
||||
url = f"{self.base_url}/videos/{job_id}"
|
||||
while job_status not in {"completed", "failed"}:
|
||||
await asyncio.sleep(self.poll_interval)
|
||||
|
||||
data = await url_json(session, url)
|
||||
if (job_status := data.get("status", None)) is None:
|
||||
raise RuntimeError("API response missing job 'status' field - expected OpenAI compliant format")
|
||||
if asyncio.get_running_loop().time() >= deadline:
|
||||
raise RuntimeError(f"Timed out waiting for video job {job_id} to complete")
|
||||
|
||||
if job_status == "failed":
|
||||
raise RuntimeError(f"Video job failed: {data}")
|
||||
|
||||
# Retrieve completed content
|
||||
video_bytes = await url_bytes(session, f"{url}/content")
|
||||
|
||||
# Decode video and make a best effort at cleaning up server resources
|
||||
try:
|
||||
return bytes_to_video(video_bytes)
|
||||
finally:
|
||||
try:
|
||||
await url_json(session, url, "delete")
|
||||
except Exception as exc:
|
||||
logger.warning("Failed to clean up video job %s: %s", job_id, exc)
|
||||
|
||||
async def generate_understanding_chat_completion(
|
||||
self,
|
||||
*,
|
||||
|
||||
@@ -6,17 +6,12 @@ Original source at https://github.com/dougbtv/comfyui-vllm-omni, distributed und
|
||||
|
||||
import base64
|
||||
import mimetypes
|
||||
from fractions import Fraction
|
||||
from io import BytesIO
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import torch
|
||||
from av.audio.frame import AudioFrame
|
||||
from av.audio.resampler import AudioResampler
|
||||
from av.video.frame import VideoFrame
|
||||
from comfy_api.input import AudioInput, VideoInput
|
||||
from comfy_api.latest import InputImpl, Types
|
||||
from comfy_extras import nodes_audio
|
||||
from PIL import Image
|
||||
|
||||
@@ -143,75 +138,6 @@ def video_to_base64(video: VideoInput, filename: str = "video.mp4") -> str:
|
||||
return f"data:{mime_type};base64,{base64_str}"
|
||||
|
||||
|
||||
def bytes_to_video(video_bytes: bytes) -> VideoInput:
|
||||
video_buffer = BytesIO(video_bytes)
|
||||
|
||||
try:
|
||||
with av.open(video_buffer, mode="r") as container:
|
||||
video_stream = next((s for s in container.streams if s.type == "video"), None)
|
||||
if video_stream is None:
|
||||
raise ValueError("No video stream found in decoded payload.")
|
||||
|
||||
frames = []
|
||||
for frame in container.decode(video_stream):
|
||||
if not isinstance(frame, VideoFrame):
|
||||
continue
|
||||
image = frame.to_ndarray(format="rgb24")
|
||||
frames.append(torch.from_numpy(image).float() / 255.0)
|
||||
|
||||
if len(frames) == 0:
|
||||
raise ValueError("No video frames found in decoded payload.")
|
||||
|
||||
images = torch.stack(frames, dim=0)
|
||||
frame_rate = Fraction(video_stream.average_rate) if video_stream.average_rate else Fraction(1)
|
||||
|
||||
audio: AudioInput | None = None
|
||||
if len(container.streams.audio):
|
||||
audio_stream = container.streams.audio[-1]
|
||||
audio_frames = []
|
||||
resampler = AudioResampler(format="fltp")
|
||||
for frame in container.decode(audio_stream):
|
||||
if not isinstance(frame, AudioFrame):
|
||||
continue
|
||||
resampled = resampler.resample(frame)
|
||||
if not isinstance(resampled, list):
|
||||
resampled = [resampled]
|
||||
for audio_frame in resampled:
|
||||
if audio_frame is not None:
|
||||
audio_frames.append(audio_frame.to_ndarray())
|
||||
|
||||
if len(audio_frames) > 0:
|
||||
audio_data = np.concatenate(audio_frames, axis=1)
|
||||
sample_rate = int(audio_stream.sample_rate) if audio_stream.sample_rate else 1
|
||||
audio = {
|
||||
"waveform": torch.from_numpy(audio_data).unsqueeze(0),
|
||||
"sample_rate": sample_rate,
|
||||
}
|
||||
|
||||
components = Types.VideoComponents(
|
||||
images=images,
|
||||
frame_rate=frame_rate,
|
||||
audio=audio,
|
||||
metadata=container.metadata if container.metadata else None,
|
||||
)
|
||||
except Exception as e:
|
||||
raise RuntimeError(f"Failed to decode video: {e}")
|
||||
|
||||
return InputImpl.VideoFromComponents(components)
|
||||
|
||||
|
||||
def base64_to_video(base64_str: str) -> VideoInput:
|
||||
if base64_str.startswith("data:video"):
|
||||
_, base64_str = base64_str.split(",", 1)
|
||||
|
||||
try:
|
||||
video_bytes = base64.b64decode(base64_str)
|
||||
except Exception as e:
|
||||
raise ValueError(f"Invalid base64 string: {e}")
|
||||
|
||||
return bytes_to_video(video_bytes)
|
||||
|
||||
|
||||
def audio_to_bytes(audio: AudioInput, filename: str = "audio.mp3", quality: str = "128k") -> BytesIO:
|
||||
waveform = audio["waveform"][0] # Shape: (C, T)
|
||||
sample_rate = audio["sample_rate"]
|
||||
|
||||
@@ -23,10 +23,6 @@ class QwenTTSModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class WanModelSpecificParams(dict):
|
||||
pass
|
||||
|
||||
|
||||
class ModelMode(Enum):
|
||||
IMAGE_GENERATION = auto()
|
||||
VIDEO_GENERATION = auto()
|
||||
|
||||
@@ -16,7 +16,7 @@ def validate_model_and_sampling_params_types(
|
||||
# Skip if no spec or no sampling params
|
||||
pipeline_spec, _ = lookup_model_spec(model_name)
|
||||
if pipeline_spec is None:
|
||||
logger.info(f"skipping sampling params check because spec for {model_name} is not found")
|
||||
logger.info("skipping sampling params check because spec is not found")
|
||||
return
|
||||
if sampling_param_list is None:
|
||||
return
|
||||
@@ -70,17 +70,13 @@ def add_sampling_parameters_to_stage(
|
||||
|
||||
stages = pipeline_spec["stages"]
|
||||
if isinstance(sampling_param_list, dict):
|
||||
sampling_param_list = sampling_param_list.__class__(sampling_param_list)
|
||||
sampling_param_list.update(params_to_add)
|
||||
elif sampling_param_list is None:
|
||||
sampling_param_list = params_to_add.copy()
|
||||
else:
|
||||
for i, stage in enumerate(stages):
|
||||
if stage == stage_type:
|
||||
stage_param = sampling_param_list[i]
|
||||
stage_param = stage_param.__class__(stage_param)
|
||||
stage_param.update(params_to_add)
|
||||
sampling_param_list[i] = stage_param
|
||||
sampling_param_list[i].update(params_to_add)
|
||||
|
||||
return sampling_param_list
|
||||
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 129 KiB |
@@ -1,513 +0,0 @@
|
||||
{
|
||||
"id": "3aea1d08-7de5-4997-ba71-b2f66ad72a21",
|
||||
"revision": 0,
|
||||
"last_node_id": 14,
|
||||
"last_link_id": 28,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 10,
|
||||
"type": "MarkdownNote",
|
||||
"pos": [
|
||||
227.99306819575278,
|
||||
-231.24143306069843
|
||||
],
|
||||
"size": [
|
||||
240.3326478922474,
|
||||
136.3505820791881
|
||||
],
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"inputs": [],
|
||||
"outputs": [],
|
||||
"title": "Note: Task and Input",
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"vLLM-Omni nodes are categorized based on the output modality. The \"Generate Video\" node supports both text-to-video generation and image-to-video generation. The node will create corresponding payload depending on whether an input image is present or not."
|
||||
],
|
||||
"color": "#432",
|
||||
"bgcolor": "#000"
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
496.31859627609606,
|
||||
-229.71277089860084
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
314
|
||||
],
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "image"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "choose file to upload",
|
||||
"name": "upload",
|
||||
"type": "IMAGEUPLOAD",
|
||||
"widget": {
|
||||
"name": "upload"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "IMAGE",
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
25
|
||||
]
|
||||
},
|
||||
{
|
||||
"localized_name": "MASK",
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": null
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"cute-cat.jpg",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "VLLMOmniDiffusionSampling",
|
||||
"pos": [
|
||||
465.5140218220927,
|
||||
148.80072646828987
|
||||
],
|
||||
"size": [
|
||||
284.205078125,
|
||||
226
|
||||
],
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "n",
|
||||
"name": "n",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "n"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_inference_steps",
|
||||
"name": "num_inference_steps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_inference_steps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "guidance_scale",
|
||||
"name": "guidance_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "true_cfg_scale",
|
||||
"name": "true_cfg_scale",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "true_cfg_scale"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_slicing",
|
||||
"name": "vae_use_slicing",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_slicing"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "vae_use_tiling",
|
||||
"name": "vae_use_tiling",
|
||||
"type": "BOOLEAN",
|
||||
"widget": {
|
||||
"name": "vae_use_tiling"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "seed",
|
||||
"name": "seed",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "seed"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "diffusion sampling params",
|
||||
"name": "diffusion sampling params",
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"links": [
|
||||
27
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniDiffusionSampling"
|
||||
},
|
||||
"widgets_values": [
|
||||
1,
|
||||
50,
|
||||
1,
|
||||
1,
|
||||
false,
|
||||
false,
|
||||
1992,
|
||||
"randomize"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 14,
|
||||
"type": "VLLMOmniWanParams",
|
||||
"pos": [
|
||||
518.6005521845747,
|
||||
428.4455663719664
|
||||
],
|
||||
"size": [
|
||||
270,
|
||||
106
|
||||
],
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "guidance_scale_2",
|
||||
"name": "guidance_scale_2",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "guidance_scale_2"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "boundary_ratio",
|
||||
"name": "boundary_ratio",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "boundary_ratio"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "flow_shift",
|
||||
"name": "flow_shift",
|
||||
"type": "FLOAT",
|
||||
"widget": {
|
||||
"name": "flow_shift"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "Wan video params",
|
||||
"name": "Wan video params",
|
||||
"type": "VIDEO_PARAMS",
|
||||
"links": [
|
||||
28
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniWanParams"
|
||||
},
|
||||
"widgets_values": [
|
||||
4,
|
||||
0.875,
|
||||
5
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"type": "VLLMOmniGenerateVideo",
|
||||
"pos": [
|
||||
827.2566336087867,
|
||||
-73.77449831827545
|
||||
],
|
||||
"size": [
|
||||
400,
|
||||
346
|
||||
],
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "image",
|
||||
"name": "image",
|
||||
"shape": 7,
|
||||
"type": "IMAGE",
|
||||
"link": 25
|
||||
},
|
||||
{
|
||||
"localized_name": "sampling_params",
|
||||
"name": "sampling_params",
|
||||
"shape": 7,
|
||||
"type": "SAMPLING_PARAMS",
|
||||
"link": 27
|
||||
},
|
||||
{
|
||||
"localized_name": "lora",
|
||||
"name": "lora",
|
||||
"shape": 7,
|
||||
"type": "REMOTE_LORA",
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model_params",
|
||||
"name": "model_params",
|
||||
"shape": 7,
|
||||
"type": "VIDEO_PARAMS",
|
||||
"link": 28
|
||||
},
|
||||
{
|
||||
"localized_name": "url",
|
||||
"name": "url",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "url"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "model",
|
||||
"name": "model",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "model"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "prompt",
|
||||
"name": "prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "negative_prompt",
|
||||
"name": "negative_prompt",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "negative_prompt"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "width",
|
||||
"name": "width",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "width"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "height",
|
||||
"name": "height",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "height"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "fps",
|
||||
"name": "fps",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "fps"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "num_frames",
|
||||
"name": "num_frames",
|
||||
"type": "INT",
|
||||
"widget": {
|
||||
"name": "num_frames"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"links": [
|
||||
26
|
||||
]
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "VLLMOmniGenerateVideo"
|
||||
},
|
||||
"widgets_values": [
|
||||
"http://localhost:8000/v1",
|
||||
"Wan-AI/Wan2.2-TI2V-5B-Diffusers",
|
||||
"Make the cat in the attached video yawning. Eye closed and mouth wide open. Keep the original realistic style.",
|
||||
"Cartoonish",
|
||||
900,
|
||||
1200,
|
||||
12,
|
||||
41
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 13,
|
||||
"type": "SaveVideo",
|
||||
"pos": [
|
||||
1257.5441369501657,
|
||||
-44.12956394419863
|
||||
],
|
||||
"size": [
|
||||
374,
|
||||
372.34896826553796
|
||||
],
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"localized_name": "video",
|
||||
"name": "video",
|
||||
"type": "VIDEO",
|
||||
"link": 26
|
||||
},
|
||||
{
|
||||
"localized_name": "filename_prefix",
|
||||
"name": "filename_prefix",
|
||||
"type": "STRING",
|
||||
"widget": {
|
||||
"name": "filename_prefix"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "format",
|
||||
"name": "format",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "format"
|
||||
},
|
||||
"link": null
|
||||
},
|
||||
{
|
||||
"localized_name": "codec",
|
||||
"name": "codec",
|
||||
"type": "COMBO",
|
||||
"widget": {
|
||||
"name": "codec"
|
||||
},
|
||||
"link": null
|
||||
}
|
||||
],
|
||||
"outputs": [],
|
||||
"properties": {},
|
||||
"widgets_values": [
|
||||
"video/ComfyUI",
|
||||
"auto",
|
||||
"auto"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
25,
|
||||
4,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
26,
|
||||
12,
|
||||
0,
|
||||
13,
|
||||
0,
|
||||
"VIDEO"
|
||||
],
|
||||
[
|
||||
27,
|
||||
8,
|
||||
0,
|
||||
12,
|
||||
1,
|
||||
"SAMPLING_PARAMS"
|
||||
],
|
||||
[
|
||||
28,
|
||||
14,
|
||||
0,
|
||||
12,
|
||||
3,
|
||||
"VIDEO_PARAMS"
|
||||
]
|
||||
],
|
||||
"groups": [
|
||||
{
|
||||
"id": 1,
|
||||
"title": "Input",
|
||||
"bounding": [
|
||||
213.1706010087147,
|
||||
-313.20095750667554,
|
||||
560.8246471437027,
|
||||
407.4246532472182
|
||||
],
|
||||
"color": "#3f789e",
|
||||
"font_size": 24,
|
||||
"flags": {}
|
||||
}
|
||||
],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"workflowRendererVersion": "LG",
|
||||
"ds": {
|
||||
"scale": 1.1469075819486894,
|
||||
"offset": [
|
||||
101.7632950847089,
|
||||
492.29121889347806
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -15,8 +15,6 @@ This document explains the benchmark architecture across all benchmark assets in
|
||||
## Directory layout
|
||||
- `benchmarks/build_dataset/` — dataset prep utilities (e.g., SeedTTS top100).
|
||||
- `benchmarks/<model>/vllm_omni/` — vLLM-Omni pipeline benchmarks, logs, outputs.
|
||||
- `benchmarks/accuracy/` — accuracy benchmark integrations that adapt external
|
||||
benchmark suites to vLLM-Omni serving and evaluation flows.
|
||||
- Add new tasks under `benchmarks/<model>/...` with the same pattern: `transformers/`, `vllm_omni/`, task-specific README, and (optionally) dataset prep notes.
|
||||
|
||||
## Reference workflows
|
||||
|
||||
@@ -1 +0,0 @@
|
||||
"""Benchmark helpers and runnable benchmark entrypoints."""
|
||||
@@ -1,27 +0,0 @@
|
||||
# Accuracy Benchmarks
|
||||
|
||||
This directory hosts accuracy benchmark integrations that run entirely through a
|
||||
local `vllm-omni serve` deployment.
|
||||
|
||||
Current integrations:
|
||||
|
||||
- `text_to_image/`: GEBench generation + local judge scoring flow.
|
||||
- `image_to_image/`: GEdit-Bench generation + local VIEScore-style scoring flow.
|
||||
|
||||
Design notes:
|
||||
|
||||
- Generation is executed through the OpenAI-compatible endpoints exposed by
|
||||
`vllm-omni serve`.
|
||||
- Evaluation is also executed through a local OpenAI-compatible judge model
|
||||
served by `vllm-omni`.
|
||||
- Both generation and judge requests accept either `http://host:port` or
|
||||
`http://host:port/v1`.
|
||||
- Output directory layout intentionally stays close to the upstream repos.
|
||||
|
||||
Test guidance:
|
||||
|
||||
- Local static/self-checks live in `tests/benchmarks/test_accuracy_bench_utils.py`.
|
||||
- End-to-end generation/evaluation should be validated in a remote GPU
|
||||
environment. In the current repo marker system there is `L4` but no `L5`
|
||||
marker, so benchmark smoke tests should be wired as `advanced_model +
|
||||
benchmark + L4` when GPU capacity is available.
|
||||
@@ -1 +0,0 @@
|
||||
"""Accuracy benchmark integrations for vLLM-Omni."""
|
||||
@@ -1,199 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def ensure_dir(path: Path) -> Path:
|
||||
path.mkdir(parents=True, exist_ok=True)
|
||||
return path
|
||||
|
||||
|
||||
def load_json(path: Path) -> dict[str, Any]:
|
||||
with path.open("r", encoding="utf-8") as handle:
|
||||
return json.load(handle)
|
||||
|
||||
|
||||
def write_json(path: Path, payload: dict[str, Any]) -> None:
|
||||
ensure_dir(path.parent)
|
||||
with path.open("w", encoding="utf-8") as handle:
|
||||
json.dump(payload, handle, indent=2, ensure_ascii=False)
|
||||
|
||||
|
||||
def save_image(path: Path, image: Image.Image) -> None:
|
||||
ensure_dir(path.parent)
|
||||
image.save(path)
|
||||
|
||||
|
||||
def find_first_image(folder: Path, stem: str | None = None) -> Path | None:
|
||||
patterns = [f"{stem}.*"] if stem else ["*.png", "*.jpg", "*.jpeg", "*.webp"]
|
||||
for pattern in patterns:
|
||||
for candidate in sorted(folder.glob(pattern)):
|
||||
if candidate.suffix.lower() in {".png", ".jpg", ".jpeg", ".webp"}:
|
||||
return candidate
|
||||
return None
|
||||
|
||||
|
||||
def extract_json_object(raw_text: str) -> dict[str, Any]:
|
||||
raw_text = raw_text.strip()
|
||||
delimiter = "||V^=^V||"
|
||||
if raw_text.count(delimiter) >= 2:
|
||||
start = raw_text.find(delimiter) + len(delimiter)
|
||||
end = raw_text.rfind(delimiter)
|
||||
raw_text = raw_text[start:end].strip()
|
||||
|
||||
start = raw_text.find("{")
|
||||
end = raw_text.rfind("}")
|
||||
if start == -1 or end == -1 or end < start:
|
||||
raise ValueError(f"Could not find JSON object in: {raw_text[:200]}")
|
||||
return json.loads(raw_text[start : end + 1])
|
||||
|
||||
|
||||
def build_openai_url(base_url: str, api_path: str) -> str:
|
||||
base = base_url.rstrip("/")
|
||||
normalized_path = api_path if api_path.startswith("/") else f"/{api_path}"
|
||||
if base.endswith(normalized_path):
|
||||
return base
|
||||
if base.endswith("/v1"):
|
||||
return f"{base}{normalized_path}"
|
||||
return f"{base}/v1{normalized_path}"
|
||||
|
||||
|
||||
def pil_to_base64(image: Image.Image, image_format: str = "PNG") -> str:
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format=image_format)
|
||||
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def pil_to_data_url(image: Image.Image, image_format: str = "PNG") -> str:
|
||||
return f"data:image/{image_format.lower()};base64,{pil_to_base64(image, image_format=image_format)}"
|
||||
|
||||
|
||||
def decode_base64_image(encoded: str) -> Image.Image:
|
||||
image = Image.open(io.BytesIO(base64.b64decode(encoded)))
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
|
||||
|
||||
def pil_to_png_bytes(image: Image.Image) -> bytes:
|
||||
buffer = io.BytesIO()
|
||||
image.save(buffer, format="PNG")
|
||||
return buffer.getvalue()
|
||||
|
||||
|
||||
class VllmOmniImageClient:
|
||||
"""Thin OpenAI-compatible image client for vLLM-Omni serving."""
|
||||
|
||||
def __init__(self, base_url: str, api_key: str = "EMPTY", timeout: int = 600):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.timeout = timeout
|
||||
|
||||
@property
|
||||
def _headers(self) -> dict[str, str]:
|
||||
return {
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
}
|
||||
|
||||
def generate_text_to_image(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
width: int,
|
||||
height: int,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = None,
|
||||
output_compression: int | None = None,
|
||||
) -> Image.Image:
|
||||
payload: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"size": f"{width}x{height}",
|
||||
"response_format": "b64_json",
|
||||
"num_inference_steps": num_inference_steps,
|
||||
}
|
||||
if guidance_scale is not None:
|
||||
payload["guidance_scale"] = guidance_scale
|
||||
if seed is not None:
|
||||
payload["seed"] = seed
|
||||
if output_compression is not None:
|
||||
payload["output_compression"] = output_compression
|
||||
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, "/images/generations"),
|
||||
json=payload,
|
||||
headers=self._headers,
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
return decode_base64_image(response.json()["data"][0]["b64_json"])
|
||||
|
||||
def generate_image_edit(
|
||||
self,
|
||||
*,
|
||||
model: str,
|
||||
prompt: str,
|
||||
images: Image.Image | list[Image.Image],
|
||||
width: int,
|
||||
height: int,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
output_compression: int | None = None,
|
||||
) -> Image.Image:
|
||||
if not isinstance(images, list):
|
||||
images = [images]
|
||||
data: dict[str, Any] = {
|
||||
"model": model,
|
||||
"prompt": prompt,
|
||||
"n": 1,
|
||||
"size": f"{width}x{height}",
|
||||
"response_format": "b64_json",
|
||||
"num_inference_steps": str(num_inference_steps),
|
||||
}
|
||||
if guidance_scale is not None:
|
||||
data["guidance_scale"] = str(guidance_scale)
|
||||
if seed is not None:
|
||||
data["seed"] = str(seed)
|
||||
if negative_prompt:
|
||||
data["negative_prompt"] = negative_prompt
|
||||
if output_compression is not None:
|
||||
data["output_compression"] = str(output_compression)
|
||||
|
||||
files = [
|
||||
(
|
||||
"image[]" if len(images) > 1 else "image",
|
||||
(f"image_{index}.png", pil_to_png_bytes(image), "image/png"),
|
||||
)
|
||||
for index, image in enumerate(images)
|
||||
]
|
||||
|
||||
edit_paths = ["/images/edits", "/images/edit"]
|
||||
last_response: requests.Response | None = None
|
||||
for api_path in edit_paths:
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, api_path),
|
||||
data=data,
|
||||
files=files,
|
||||
headers={"Authorization": f"Bearer {self.api_key}"},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
last_response = response
|
||||
if response.status_code != 404:
|
||||
response.raise_for_status()
|
||||
return decode_base64_image(response.json()["data"][0]["b64_json"])
|
||||
|
||||
assert last_response is not None
|
||||
last_response.raise_for_status()
|
||||
raise ValueError("No image payload returned from image edit endpoint")
|
||||
@@ -1,103 +0,0 @@
|
||||
# GEdit-Bench on vLLM-Omni
|
||||
|
||||
This integration adapts the upstream `stepfun-ai/Step1X-Edit/GEdit-Bench`
|
||||
evaluation flow into `vllm-omni/benchmarks/accuracy/image_to_image`.
|
||||
|
||||
Upstream mapping:
|
||||
|
||||
- `run_gedit_score.py` -> `run_gedit_bench.py evaluate`
|
||||
- `calculate_statistics.py` -> `run_gedit_bench.py summarize`
|
||||
- upstream output layout under `results/<method>/fullset/...` is preserved
|
||||
|
||||
What changed:
|
||||
|
||||
- The upstream repo mainly ships evaluation scripts. This integration adds a
|
||||
generation runner that uses the local `vllm-omni` OpenAI-compatible
|
||||
`/v1/images/edits` endpoint to produce benchmark outputs in the expected
|
||||
directory structure.
|
||||
- The evaluator keeps the same VIEScore-style decomposition:
|
||||
- `semantics_score`
|
||||
- `quality_score`
|
||||
- `overall_score = sqrt(semantics_score * quality_score)`
|
||||
- Judge calls are routed to a local OpenAI-compatible model served by
|
||||
`vllm-omni`, not a remote provider.
|
||||
|
||||
Dataset:
|
||||
|
||||
- Default `--dataset-ref` is `stepfun-ai/GEdit-Bench`
|
||||
- You can also pass a local dataset directory previously saved with
|
||||
Hugging Face `datasets`
|
||||
|
||||
Example usage:
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/image_to_image/run_gedit_bench.py generate \
|
||||
--output-root benchmarks/accuracy/image_to_image/results \
|
||||
--base-url http://127.0.0.1:8000 \
|
||||
--model Qwen/Qwen-Image-Edit \
|
||||
--model-name qwen_image_edit \
|
||||
--dataset-ref stepfun-ai/GEdit-Bench \
|
||||
--task-type all \
|
||||
--instruction-language en
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/image_to_image/run_gedit_bench.py evaluate \
|
||||
--output-root benchmarks/accuracy/image_to_image/results \
|
||||
--model-name qwen_image_edit \
|
||||
--save-dir benchmarks/accuracy/image_to_image/scores \
|
||||
--dataset-ref stepfun-ai/GEdit-Bench \
|
||||
--judge-base-url http://127.0.0.1:8000 \
|
||||
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--judge-api-key EMPTY
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/image_to_image/run_gedit_bench.py summarize \
|
||||
--csv-path benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv \
|
||||
--language en
|
||||
```
|
||||
|
||||
Example summary output:
|
||||
|
||||
```json
|
||||
{
|
||||
"language": "all",
|
||||
"languages": {
|
||||
"en": {
|
||||
"overall": {"count": 110, "Q_SC": 6.58, "Q_PQ": 5.89, "Q_O": 5.86},
|
||||
"intersection": {"count": 78, "Q_SC": 6.50, "Q_PQ": 5.66, "Q_O": 5.65}
|
||||
},
|
||||
"cn": {
|
||||
"overall": {"count": 110, "Q_SC": 6.90, "Q_PQ": 5.78, "Q_O": 6.11},
|
||||
"intersection": {"count": 63, "Q_SC": 7.22, "Q_PQ": 5.59, "Q_O": 6.28}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Example generated images to inspect:
|
||||
|
||||
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/background_change/en/<sample>.png`
|
||||
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/text_change/en/<sample>.png`
|
||||
- `benchmarks/accuracy/image_to_image/results/qwen_image_edit/fullset/subject-replace/cn/<sample>.png`
|
||||
|
||||
Example score artifacts to inspect together with the images:
|
||||
|
||||
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_vie_score.csv`
|
||||
- `benchmarks/accuracy/image_to_image/scores/qwen_image_edit_all_all_summary.json`
|
||||
|
||||
What to expect:
|
||||
|
||||
- `Q_SC` measures instruction following and content preservation.
|
||||
- `Q_PQ` measures image naturalness and artifact quality.
|
||||
- `Q_O` is the combined overall score; higher is better.
|
||||
- `overall.count` is the number of evaluated samples for that language, while `intersection.count` is the subset with `intersection_exist == True`.
|
||||
|
||||
Notes:
|
||||
|
||||
- This flow requires the optional Hugging Face `datasets` package.
|
||||
- `generate` writes `generation_manifest.json` with local output coverage.
|
||||
- The current repo marker set exposes `L4` but not `L5`, so if you promote an
|
||||
end-to-end smoke test into CI, use the existing `advanced_model`, `benchmark`,
|
||||
and `L4` markers or introduce a new repo-wide marker explicitly first.
|
||||
@@ -1 +0,0 @@
|
||||
"""GEdit-Bench integration for vLLM-Omni."""
|
||||
@@ -1,718 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import statistics
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from datetime import datetime, timezone
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from benchmarks.accuracy.common import (
|
||||
VllmOmniImageClient,
|
||||
build_openai_url,
|
||||
extract_json_object,
|
||||
write_json,
|
||||
)
|
||||
|
||||
GROUPS = [
|
||||
"background_change",
|
||||
"color_alter",
|
||||
"material_alter",
|
||||
"motion_change",
|
||||
"ps_human",
|
||||
"style_change",
|
||||
"subject-add",
|
||||
"subject-remove",
|
||||
"subject-replace",
|
||||
"text_change",
|
||||
"tone_transfer",
|
||||
]
|
||||
DEFAULT_SAMPLES_PER_GROUP = 10
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
def infer_model_name(model: str) -> str:
|
||||
normalized = model.rstrip("/\\")
|
||||
name = Path(normalized).name
|
||||
return name or normalized
|
||||
|
||||
|
||||
def resolve_model_name(*, model_name: str | None, model: str | None = None, output_root: Path | None = None) -> str:
|
||||
if model_name:
|
||||
return model_name
|
||||
if model:
|
||||
return infer_model_name(model)
|
||||
if output_root is not None:
|
||||
candidates = sorted(path.name for path in output_root.iterdir() if path.is_dir())
|
||||
if len(candidates) == 1:
|
||||
return candidates[0]
|
||||
if not candidates:
|
||||
raise ValueError(f"Could not infer model-name from empty output root: {output_root}")
|
||||
raise ValueError(
|
||||
f"Could not infer model-name from output root {output_root}; multiple candidates found: {candidates}"
|
||||
)
|
||||
raise ValueError("model-name is required when it cannot be inferred from model or output-root")
|
||||
|
||||
|
||||
def parse_score_payload(raw_text: str) -> dict[str, Any]:
|
||||
try:
|
||||
parsed = extract_json_object(raw_text)
|
||||
except Exception:
|
||||
stripped = raw_text.strip()
|
||||
if stripped.startswith("[") and stripped.endswith("]"):
|
||||
parsed = {"score": json.loads(stripped), "reasoning": ""}
|
||||
else:
|
||||
raise
|
||||
score = parsed.get("score", [])
|
||||
if not isinstance(score, list):
|
||||
score = [score]
|
||||
parsed["score"] = [int(value) for value in score]
|
||||
return parsed
|
||||
|
||||
|
||||
def _utc_timestamp() -> str:
|
||||
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
||||
|
||||
|
||||
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
by_task: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for record in records:
|
||||
by_task[record["task_type"]].append(record)
|
||||
by_language[record["instruction_language"]].append(record)
|
||||
|
||||
return {
|
||||
"count": len(records),
|
||||
"by_task": {
|
||||
task_type: {
|
||||
"count": len(rows),
|
||||
"samples": sorted(row["key"] for row in rows),
|
||||
}
|
||||
for task_type, rows in sorted(by_task.items())
|
||||
},
|
||||
"by_language": {
|
||||
language: {
|
||||
"count": len(rows),
|
||||
"samples": sorted(row["key"] for row in rows),
|
||||
}
|
||||
for language, rows in sorted(by_language.items())
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def select_balanced_gedit_rows(
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
task_type: str = "all",
|
||||
instruction_language: str = "all",
|
||||
samples_per_group: int | None,
|
||||
) -> list[dict[str, Any]]:
|
||||
filtered_rows = []
|
||||
for row in rows:
|
||||
if task_type != "all" and row["task_type"] != task_type:
|
||||
continue
|
||||
if instruction_language != "all" and row["instruction_language"] != instruction_language:
|
||||
continue
|
||||
filtered_rows.append(row)
|
||||
|
||||
if samples_per_group is None:
|
||||
return filtered_rows
|
||||
|
||||
def _select_rows_for_group(group_rows: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
if instruction_language != "all":
|
||||
return group_rows[:samples_per_group]
|
||||
|
||||
rows_by_language: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for row in group_rows:
|
||||
rows_by_language[str(row.get("instruction_language", "")).strip()].append(row)
|
||||
|
||||
ordered_languages = ["en", "cn"]
|
||||
per_language_quota = samples_per_group // max(len(ordered_languages), 1)
|
||||
remainder = samples_per_group % max(len(ordered_languages), 1)
|
||||
|
||||
selected_rows: list[dict[str, Any]] = []
|
||||
leftovers: list[dict[str, Any]] = []
|
||||
for index, language in enumerate(ordered_languages):
|
||||
language_rows = rows_by_language.get(language, [])
|
||||
quota = per_language_quota + (1 if index < remainder else 0)
|
||||
selected_rows.extend(language_rows[:quota])
|
||||
leftovers.extend(language_rows[quota:])
|
||||
|
||||
if len(selected_rows) < samples_per_group:
|
||||
selected_rows.extend(leftovers[: samples_per_group - len(selected_rows)])
|
||||
|
||||
return selected_rows
|
||||
|
||||
if task_type != "all":
|
||||
return _select_rows_for_group(filtered_rows)
|
||||
|
||||
grouped_rows: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for row in filtered_rows:
|
||||
grouped_rows[row["task_type"]].append(row)
|
||||
|
||||
selected: list[dict[str, Any]] = []
|
||||
for group in GROUPS:
|
||||
selected.extend(_select_rows_for_group(grouped_rows.get(group, [])))
|
||||
return selected
|
||||
|
||||
|
||||
def summarize_gedit_rows(rows: list[dict[str, Any]], language: str = "all") -> dict[str, Any]:
|
||||
return summarize_gedit_rows_with_backbone(rows, language=language)
|
||||
|
||||
|
||||
def _mean_or_none(values: list[float]) -> float | None:
|
||||
return statistics.fmean(values) if values else None
|
||||
|
||||
|
||||
def _to_q_metrics(section: dict[str, float | None]) -> dict[str, float | None]:
|
||||
return {
|
||||
"count": section.get("count"),
|
||||
"Q_SC": section.get("avg_semantics"),
|
||||
"Q_PQ": section.get("avg_quality"),
|
||||
"Q_O": section.get("avg_overall"),
|
||||
}
|
||||
|
||||
|
||||
def _summarize_gedit_rows_single_language(
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
language: str,
|
||||
) -> dict[str, Any]:
|
||||
filtered_rows = [row for row in rows if str(row.get("instruction_language", "")).strip() == language]
|
||||
|
||||
def _to_bool(value: Any) -> bool:
|
||||
if isinstance(value, bool):
|
||||
return value
|
||||
return str(value).strip().lower() in {"1", "true", "yes"}
|
||||
|
||||
per_group: dict[str, dict[str, float | None]] = {}
|
||||
intersection_group: dict[str, dict[str, float | None]] = {}
|
||||
|
||||
for group in GROUPS:
|
||||
group_rows = [row for row in filtered_rows if row.get("task_type") == group]
|
||||
semantics = [float(row["semantics_score"]) for row in group_rows]
|
||||
quality = [float(row["quality_score"]) for row in group_rows]
|
||||
overall = [math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in group_rows]
|
||||
|
||||
per_group[group] = {
|
||||
"count": len(group_rows),
|
||||
"avg_semantics": _mean_or_none(semantics),
|
||||
"avg_quality": _mean_or_none(quality),
|
||||
"avg_overall": _mean_or_none(overall),
|
||||
}
|
||||
|
||||
intersection_rows = [row for row in group_rows if _to_bool(row.get("intersection_exist", False))]
|
||||
intersection_semantics = [float(row["semantics_score"]) for row in intersection_rows]
|
||||
intersection_quality = [float(row["quality_score"]) for row in intersection_rows]
|
||||
intersection_overall = [
|
||||
math.sqrt(float(row["semantics_score"]) * float(row["quality_score"])) for row in intersection_rows
|
||||
]
|
||||
intersection_group[group] = {
|
||||
"count": len(intersection_rows),
|
||||
"avg_semantics": _mean_or_none(intersection_semantics),
|
||||
"avg_quality": _mean_or_none(intersection_quality),
|
||||
"avg_overall": _mean_or_none(intersection_overall),
|
||||
}
|
||||
|
||||
overall_section = {
|
||||
"count": len(filtered_rows),
|
||||
"avg_semantics": _mean_or_none(
|
||||
[score["avg_semantics"] for score in per_group.values() if score["avg_semantics"] is not None]
|
||||
),
|
||||
"avg_quality": _mean_or_none(
|
||||
[score["avg_quality"] for score in per_group.values() if score["avg_quality"] is not None]
|
||||
),
|
||||
"avg_overall": _mean_or_none(
|
||||
[score["avg_overall"] for score in per_group.values() if score["avg_overall"] is not None]
|
||||
),
|
||||
}
|
||||
intersection_section = {
|
||||
"count": sum(score["count"] for score in intersection_group.values()),
|
||||
"avg_semantics": _mean_or_none(
|
||||
[score["avg_semantics"] for score in intersection_group.values() if score["avg_semantics"] is not None]
|
||||
),
|
||||
"avg_quality": _mean_or_none(
|
||||
[score["avg_quality"] for score in intersection_group.values() if score["avg_quality"] is not None]
|
||||
),
|
||||
"avg_overall": _mean_or_none(
|
||||
[score["avg_overall"] for score in intersection_group.values() if score["avg_overall"] is not None]
|
||||
),
|
||||
}
|
||||
|
||||
return {
|
||||
"language": language,
|
||||
"by_group": {group: _to_q_metrics(section) for group, section in per_group.items()},
|
||||
"overall": _to_q_metrics(overall_section),
|
||||
"intersection": _to_q_metrics(intersection_section),
|
||||
}
|
||||
|
||||
|
||||
def summarize_gedit_rows_with_backbone(
|
||||
rows: list[dict[str, Any]],
|
||||
*,
|
||||
language: str = "all",
|
||||
) -> dict[str, Any]:
|
||||
if language == "all":
|
||||
return {
|
||||
"language": "all",
|
||||
"languages": {
|
||||
single_language: _summarize_gedit_rows_single_language(
|
||||
rows,
|
||||
language=single_language,
|
||||
)
|
||||
for single_language in ["en", "cn"]
|
||||
},
|
||||
}
|
||||
|
||||
return _summarize_gedit_rows_single_language(rows, language=language)
|
||||
|
||||
|
||||
def _require_datasets():
|
||||
try:
|
||||
from datasets import load_dataset, load_from_disk
|
||||
except ImportError as exc:
|
||||
raise ImportError("GEdit-Bench requires the optional `datasets` package.") from exc
|
||||
return load_dataset, load_from_disk
|
||||
|
||||
|
||||
def _load_gedit_dataset(dataset_ref: str):
|
||||
load_dataset, load_from_disk = _require_datasets()
|
||||
dataset_path = Path(dataset_ref)
|
||||
if dataset_path.exists():
|
||||
if (dataset_path / "state.json").exists() and (dataset_path / "dataset_info.json").exists():
|
||||
return load_from_disk(str(dataset_path))
|
||||
return load_dataset(str(dataset_path))
|
||||
return load_dataset(dataset_ref)
|
||||
|
||||
|
||||
def _resolve_gedit_split(dataset_obj: Any) -> Any:
|
||||
if isinstance(dataset_obj, dict):
|
||||
if "train" in dataset_obj:
|
||||
return dataset_obj["train"]
|
||||
return dataset_obj
|
||||
try:
|
||||
return dataset_obj["train"]
|
||||
except Exception:
|
||||
return dataset_obj
|
||||
|
||||
|
||||
def _to_pil_image(value: Any) -> Image.Image:
|
||||
if isinstance(value, Image.Image):
|
||||
return value.convert("RGB")
|
||||
if isinstance(value, dict) and "bytes" in value:
|
||||
image = Image.open(BytesIO(value["bytes"]))
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
raise TypeError(f"Unsupported image payload type: {type(value)!r}")
|
||||
|
||||
|
||||
class LocalVIEScorer:
|
||||
def __init__(self, *, base_url: str, api_key: str, model: str, timeout: int = 600):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
self.sc_prompt = (
|
||||
"You are evaluating image editing quality.\n"
|
||||
"Two images are provided: the source image and the edited image.\n"
|
||||
'Return JSON only in the format {"score": [edit_success, content_preservation], '
|
||||
'"reasoning": "..."}.\n'
|
||||
"Each score must be an integer from 0 to 10.\n"
|
||||
"Editing instruction: <instruction>"
|
||||
)
|
||||
self.pq_prompt = (
|
||||
"You are evaluating image quality.\n"
|
||||
'Return JSON only in the format {"score": [naturalness, artifact_free], "reasoning": "..."}.\n'
|
||||
"Each score must be an integer from 0 to 10."
|
||||
)
|
||||
|
||||
def _request(self, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
|
||||
from benchmarks.accuracy.common import pil_to_data_url
|
||||
|
||||
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||||
for image in images:
|
||||
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
|
||||
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, "/chat/completions"),
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"temperature": 0,
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
message_content = response.json()["choices"][0]["message"]["content"]
|
||||
if isinstance(message_content, list):
|
||||
text = "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
|
||||
else:
|
||||
text = str(message_content)
|
||||
return parse_score_payload(text)
|
||||
|
||||
def evaluate(self, source_image: Image.Image, edited_image: Image.Image, instruction: str) -> dict[str, float]:
|
||||
sc_payload = self._request(self.sc_prompt.replace("<instruction>", instruction), [source_image, edited_image])
|
||||
pq_payload = self._request(self.pq_prompt, [edited_image])
|
||||
semantics = float(min(sc_payload["score"])) if sc_payload["score"] else 0.0
|
||||
quality = float(min(pq_payload["score"])) if pq_payload["score"] else 0.0
|
||||
overall = math.sqrt(semantics * quality)
|
||||
return {
|
||||
"semantics_score": semantics,
|
||||
"quality_score": quality,
|
||||
"overall_score": overall,
|
||||
}
|
||||
|
||||
|
||||
class GEditBenchRunner:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dataset_ref: str,
|
||||
output_root: Path,
|
||||
base_url: str,
|
||||
model: str,
|
||||
api_key: str = "EMPTY",
|
||||
width: int = 512,
|
||||
height: int = 512,
|
||||
num_inference_steps: int = 20,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = 42,
|
||||
):
|
||||
self.dataset_ref = dataset_ref
|
||||
self.output_root = output_root
|
||||
self.model = model
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.guidance_scale = guidance_scale
|
||||
self.seed = seed
|
||||
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
*,
|
||||
model_name: str,
|
||||
task_type: str = "all",
|
||||
instruction_language: str = "all",
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_group: int | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
|
||||
rows = select_balanced_gedit_rows(
|
||||
list(dataset),
|
||||
task_type=task_type,
|
||||
instruction_language=instruction_language,
|
||||
samples_per_group=samples_per_group,
|
||||
)
|
||||
if max_samples is not None:
|
||||
rows = rows[:max_samples]
|
||||
|
||||
outputs: list[dict[str, Any]] = []
|
||||
total = len(rows)
|
||||
if workers <= 1:
|
||||
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
|
||||
for item in rows:
|
||||
result = self._safe_generate_one(model_name, item)
|
||||
if result:
|
||||
outputs.append(result)
|
||||
progress.update(1)
|
||||
return outputs
|
||||
|
||||
with tqdm(total=total, desc="GEdit generate", unit="sample") as progress:
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [executor.submit(self._safe_generate_one, model_name, item) for item in rows]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
outputs.append(result)
|
||||
progress.update(1)
|
||||
return outputs
|
||||
|
||||
def _safe_generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
try:
|
||||
return self._generate_one(model_name, item)
|
||||
except Exception:
|
||||
logger.exception("Failed to generate GEdit-Bench sample %s", item.get("key", "<unknown>"))
|
||||
return None
|
||||
|
||||
def _generate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
output_path = (
|
||||
self.output_root
|
||||
/ model_name
|
||||
/ "fullset"
|
||||
/ item["task_type"]
|
||||
/ item["instruction_language"]
|
||||
/ f"{item['key']}.png"
|
||||
)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
if output_path.exists():
|
||||
return {
|
||||
"task_type": item["task_type"],
|
||||
"instruction_language": item["instruction_language"],
|
||||
"key": item["key"],
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
source_image = _to_pil_image(item["input_image_raw"])
|
||||
edited_image = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=item["instruction"],
|
||||
images=source_image,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
edited_image.save(output_path)
|
||||
return {
|
||||
"task_type": item["task_type"],
|
||||
"instruction_language": item["instruction_language"],
|
||||
"key": item["key"],
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
|
||||
class GEditBenchEvaluator:
|
||||
def __init__(self, *, dataset_ref: str, output_root: Path, scorer: LocalVIEScorer):
|
||||
self.dataset_ref = dataset_ref
|
||||
self.output_root = output_root
|
||||
self.scorer = scorer
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
*,
|
||||
model_name: str,
|
||||
save_dir: Path,
|
||||
task_type: str = "all",
|
||||
instruction_language: str = "all",
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_group: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
dataset = _resolve_gedit_split(_load_gedit_dataset(self.dataset_ref))
|
||||
rows = select_balanced_gedit_rows(
|
||||
list(dataset),
|
||||
task_type=task_type,
|
||||
instruction_language=instruction_language,
|
||||
samples_per_group=samples_per_group,
|
||||
)
|
||||
if max_samples is not None:
|
||||
rows = rows[:max_samples]
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
total = len(rows)
|
||||
if workers <= 1:
|
||||
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
|
||||
for item in rows:
|
||||
result = self._safe_evaluate_one(model_name, item)
|
||||
if result:
|
||||
results.append(result)
|
||||
progress.update(1)
|
||||
else:
|
||||
with tqdm(total=total, desc="GEdit evaluate", unit="sample") as progress:
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [executor.submit(self._safe_evaluate_one, model_name, item) for item in rows]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
results.append(result)
|
||||
progress.update(1)
|
||||
|
||||
save_dir.mkdir(parents=True, exist_ok=True)
|
||||
base_name = f"{model_name}_{task_type}_{instruction_language}"
|
||||
timestamp = _utc_timestamp()
|
||||
csv_path = save_dir / f"{base_name}_vie_score.csv"
|
||||
timestamped_csv_path = save_dir / f"{base_name}_vie_score_{timestamp}.csv"
|
||||
|
||||
def _write_csv(path: Path) -> None:
|
||||
with path.open("w", encoding="utf-8", newline="") as handle:
|
||||
writer = csv.DictWriter(
|
||||
handle,
|
||||
fieldnames=[
|
||||
"key",
|
||||
"task_type",
|
||||
"edited_image",
|
||||
"instruction",
|
||||
"semantics_score",
|
||||
"quality_score",
|
||||
"overall_score",
|
||||
"intersection_exist",
|
||||
"instruction_language",
|
||||
],
|
||||
)
|
||||
writer.writeheader()
|
||||
for row in results:
|
||||
writer.writerow(row)
|
||||
|
||||
_write_csv(csv_path)
|
||||
_write_csv(timestamped_csv_path)
|
||||
|
||||
summary = summarize_gedit_rows_with_backbone(
|
||||
results,
|
||||
language=instruction_language,
|
||||
)
|
||||
summary_path = save_dir / f"{base_name}_summary.json"
|
||||
timestamped_summary_path = save_dir / f"{base_name}_summary_{timestamp}.json"
|
||||
write_json(summary_path, summary)
|
||||
write_json(timestamped_summary_path, summary)
|
||||
return {
|
||||
"results": results,
|
||||
"summary": summary,
|
||||
"csv_path": str(csv_path),
|
||||
"summary_path": str(summary_path),
|
||||
"timestamped_csv_path": str(timestamped_csv_path),
|
||||
"timestamped_summary_path": str(timestamped_summary_path),
|
||||
}
|
||||
|
||||
def _safe_evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
try:
|
||||
return self._evaluate_one(model_name, item)
|
||||
except Exception:
|
||||
logger.exception("Failed to evaluate GEdit-Bench sample %s", item.get("key", "<unknown>"))
|
||||
return None
|
||||
|
||||
def _evaluate_one(self, model_name: str, item: dict[str, Any]) -> dict[str, Any] | None:
|
||||
edited_image_path = (
|
||||
self.output_root
|
||||
/ model_name
|
||||
/ "fullset"
|
||||
/ item["task_type"]
|
||||
/ item["instruction_language"]
|
||||
/ f"{item['key']}.png"
|
||||
)
|
||||
if not edited_image_path.exists():
|
||||
return None
|
||||
|
||||
source_image = _to_pil_image(item["input_image_raw"])
|
||||
edited_image = Image.open(edited_image_path).convert("RGB")
|
||||
scores = self.scorer.evaluate(source_image, edited_image, item["instruction"])
|
||||
return {
|
||||
"key": item["key"],
|
||||
"task_type": item["task_type"],
|
||||
"edited_image": str(edited_image_path),
|
||||
"instruction": item["instruction"],
|
||||
"semantics_score": scores["semantics_score"],
|
||||
"quality_score": scores["quality_score"],
|
||||
"overall_score": scores["overall_score"],
|
||||
"intersection_exist": item.get("Intersection_exist", False),
|
||||
"instruction_language": item["instruction_language"],
|
||||
}
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="Run the GEdit-Bench integration against a local vLLM-Omni server.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
generate = subparsers.add_parser("generate")
|
||||
generate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
|
||||
generate.add_argument("--output-root", type=Path, required=True)
|
||||
generate.add_argument("--base-url", type=str, required=True)
|
||||
generate.add_argument("--model", type=str, required=True)
|
||||
generate.add_argument("--model-name", type=str, default=None)
|
||||
generate.add_argument("--api-key", type=str, default="EMPTY")
|
||||
generate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
|
||||
generate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
|
||||
generate.add_argument("--width", type=int, default=512)
|
||||
generate.add_argument("--height", type=int, default=512)
|
||||
generate.add_argument("--num-inference-steps", type=int, default=20)
|
||||
generate.add_argument("--guidance-scale", type=float, default=None)
|
||||
generate.add_argument("--seed", type=int, default=42)
|
||||
generate.add_argument("--workers", type=int, default=1)
|
||||
generate.add_argument("--max-samples", type=int, default=None)
|
||||
generate.add_argument("--samples-per-group", type=int, default=None)
|
||||
|
||||
evaluate = subparsers.add_parser("evaluate")
|
||||
evaluate.add_argument("--dataset-ref", type=str, default="stepfun-ai/GEdit-Bench")
|
||||
evaluate.add_argument("--output-root", type=Path, required=True)
|
||||
evaluate.add_argument("--model-name", type=str, default=None)
|
||||
evaluate.add_argument("--save-dir", type=Path, required=True)
|
||||
evaluate.add_argument("--task-type", choices=["all", *GROUPS], default="all")
|
||||
evaluate.add_argument("--instruction-language", choices=["all", "en", "cn"], default="all")
|
||||
evaluate.add_argument("--judge-base-url", type=str, required=True)
|
||||
evaluate.add_argument("--judge-model", type=str, required=True)
|
||||
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
|
||||
evaluate.add_argument("--workers", type=int, default=1)
|
||||
evaluate.add_argument("--max-samples", type=int, default=None)
|
||||
evaluate.add_argument("--samples-per-group", type=int, default=None)
|
||||
|
||||
summarize = subparsers.add_parser("summarize")
|
||||
summarize.add_argument("--csv-path", type=Path, required=True)
|
||||
summarize.add_argument("--language", choices=["all", "en", "cn"], default="all")
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.command == "generate":
|
||||
model_name = resolve_model_name(model_name=args.model_name, model=args.model)
|
||||
runner = GEditBenchRunner(
|
||||
dataset_ref=args.dataset_ref,
|
||||
output_root=args.output_root,
|
||||
base_url=args.base_url,
|
||||
model=args.model,
|
||||
api_key=args.api_key,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
)
|
||||
records = runner.generate(
|
||||
model_name=model_name,
|
||||
task_type=args.task_type,
|
||||
instruction_language=args.instruction_language,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_group=args.samples_per_group,
|
||||
)
|
||||
payload = {"records": records, "summary": summarize_generated_records(records)}
|
||||
write_json(args.output_root / model_name / "generation_manifest.json", payload)
|
||||
return 0
|
||||
|
||||
if args.command == "evaluate":
|
||||
model_name = resolve_model_name(model_name=args.model_name, output_root=args.output_root)
|
||||
scorer = LocalVIEScorer(
|
||||
base_url=args.judge_base_url,
|
||||
api_key=args.judge_api_key,
|
||||
model=args.judge_model,
|
||||
)
|
||||
evaluator = GEditBenchEvaluator(dataset_ref=args.dataset_ref, output_root=args.output_root, scorer=scorer)
|
||||
evaluator.evaluate(
|
||||
model_name=model_name,
|
||||
save_dir=args.save_dir,
|
||||
task_type=args.task_type,
|
||||
instruction_language=args.instruction_language,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_group=args.samples_per_group,
|
||||
)
|
||||
return 0
|
||||
|
||||
if args.command == "summarize":
|
||||
with args.csv_path.open("r", encoding="utf-8", newline="") as handle:
|
||||
rows = list(csv.DictReader(handle))
|
||||
summary = summarize_gedit_rows_with_backbone(rows, language=args.language)
|
||||
print(json.dumps(summary, indent=2, ensure_ascii=False))
|
||||
return 0
|
||||
|
||||
parser.error(f"Unknown command: {args.command}")
|
||||
return 1
|
||||
@@ -1,13 +0,0 @@
|
||||
# ruff: noqa: E402, I001
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from benchmarks.accuracy.image_to_image.gedit_bench import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,103 +0,0 @@
|
||||
# GEBench on vLLM-Omni
|
||||
|
||||
This integration adapts the upstream `stepfun-ai/GEBench` scripts into
|
||||
`vllm-omni/benchmarks/accuracy/text_to_image`.
|
||||
|
||||
Upstream mapping:
|
||||
|
||||
- `scripts/generate.py` -> `run_gebench.py generate`
|
||||
- `scripts/evaluate.py` -> `run_gebench.py evaluate`
|
||||
- upstream prompt / judge logic -> `gbench.py`
|
||||
|
||||
What changed:
|
||||
|
||||
- Generation calls the local OpenAI-compatible `vllm-omni` endpoints:
|
||||
- `/v1/images/generations` for text-only frame generation
|
||||
- `/v1/images/edits` for image-conditioned GUI transition generation
|
||||
- Evaluation still keeps the GEBench scoring dimensions:
|
||||
- `goal`
|
||||
- `logic`
|
||||
- `cons`
|
||||
- `ui`
|
||||
- `qual`
|
||||
- Judge calls are also routed to a local OpenAI-compatible model served by
|
||||
`vllm-omni` instead of a remote service.
|
||||
|
||||
Expected dataset layout:
|
||||
|
||||
- Clone the benchmark dataset from Hugging Face into a local directory:
|
||||
|
||||
```bash
|
||||
git clone https://huggingface.co/datasets/stepfun-ai/GEBench /path/to/GEBench
|
||||
```
|
||||
|
||||
Example usage:
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/text_to_image/run_gebench.py generate \
|
||||
--dataset-root /path/to/GEBench \
|
||||
--output-root benchmarks/accuracy/text_to_image/outputs \
|
||||
--base-url http://127.0.0.1:8000 \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--data-type type3
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/text_to_image/run_gebench.py evaluate \
|
||||
--dataset-root /path/to/GEBench \
|
||||
--output-root benchmarks/accuracy/text_to_image/outputs \
|
||||
--data-type type3 \
|
||||
--judge-base-url http://127.0.0.1:8000 \
|
||||
--judge-model Qwen/Qwen2.5-VL-7B-Instruct \
|
||||
--judge-api-key EMPTY
|
||||
```
|
||||
|
||||
```bash
|
||||
python benchmarks/accuracy/text_to_image/run_gebench.py summarize \
|
||||
--output-root benchmarks/accuracy/text_to_image/outputs
|
||||
```
|
||||
|
||||
Example summary output:
|
||||
|
||||
```json
|
||||
{
|
||||
"generation": {
|
||||
"count": 20,
|
||||
"by_type": {
|
||||
"type3": {"count": 10},
|
||||
"type4": {"count": 10}
|
||||
}
|
||||
},
|
||||
"evaluation": {
|
||||
"count": 20,
|
||||
"overall_mean": 0.52,
|
||||
"by_type": {
|
||||
"type3": {"count": 10, "overall_mean": 0.50, "overall_mean_100": 50.0},
|
||||
"type4": {"count": 10, "overall_mean": 0.54, "overall_mean_100": 54.0}
|
||||
}
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Example generated images to inspect:
|
||||
|
||||
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame0.png`
|
||||
- `benchmarks/accuracy/text_to_image/outputs/03_trajectory_text_fictionalapp/<lang_device>/<sample>/frame5.png`
|
||||
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame0.png`
|
||||
- `benchmarks/accuracy/text_to_image/outputs/04_trajectory_text_realapp/<lang_device>/<sample>/frame5.png`
|
||||
|
||||
What to expect:
|
||||
|
||||
- `overall_mean` is normalized to `0.0 ~ 1.0`; higher is better.
|
||||
- `frame0.png` is the initial GUI frame; `frame5.png` is the final trajectory frame most often used for quick inspection.
|
||||
- For full debugging, inspect the whole `frame0.png` ... `frame5.png` sequence for one sample directory.
|
||||
|
||||
Notes:
|
||||
|
||||
- GEBench upstream leaves type3/type4 generation unfinished. This integration
|
||||
fills that gap with a trajectory runner that generates `frame0.png` followed
|
||||
by `frame1.png` ... `frame5.png`.
|
||||
- Type1/2/5 require an image-edit capable model exposed through
|
||||
`vllm-omni serve`.
|
||||
- `summarize` will report both generated coverage and any existing evaluation
|
||||
summary files.
|
||||
@@ -1 +0,0 @@
|
||||
"""GEBench integration for vLLM-Omni."""
|
||||
@@ -1,927 +0,0 @@
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import statistics
|
||||
from collections import defaultdict
|
||||
from concurrent.futures import ThreadPoolExecutor, as_completed
|
||||
from dataclasses import dataclass
|
||||
from datetime import datetime, timezone
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import requests
|
||||
from PIL import Image
|
||||
|
||||
from benchmarks.accuracy.common import (
|
||||
VllmOmniImageClient,
|
||||
build_openai_url,
|
||||
ensure_dir,
|
||||
extract_json_object,
|
||||
find_first_image,
|
||||
load_json,
|
||||
pil_to_data_url,
|
||||
save_image,
|
||||
write_json,
|
||||
)
|
||||
|
||||
TYPE_TO_FOLDER = {
|
||||
"type1": "01_single_step",
|
||||
"type2": "02_multi_step",
|
||||
"type3": "03_trajectory_text_fictionalapp",
|
||||
"type4": "04_trajectory_text_realapp",
|
||||
"type5": "05_grounding_data",
|
||||
}
|
||||
SCORE_KEYS = ("goal", "logic", "cons", "ui", "qual")
|
||||
DEFAULT_SAMPLES_PER_TYPE = 10
|
||||
|
||||
|
||||
def _utc_timestamp() -> str:
|
||||
return datetime.now(timezone.utc).strftime("%Y%m%dT%H%M%SZ")
|
||||
|
||||
|
||||
def _write_json_with_timestamp(path: Path, payload: dict[str, Any]) -> Path:
|
||||
write_json(path, payload)
|
||||
timestamped_path = path.with_name(f"{path.stem}_{_utc_timestamp()}{path.suffix}")
|
||||
write_json(timestamped_path, payload)
|
||||
return timestamped_path
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class GEBenchSampleSpec:
|
||||
sample_path: Path
|
||||
metadata: dict[str, Any]
|
||||
sample_name: str
|
||||
lang_device: str
|
||||
|
||||
|
||||
def summarize_generated_records(records: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for record in records:
|
||||
by_type[record["data_type"]].append(record)
|
||||
|
||||
return {
|
||||
"count": len(records),
|
||||
"by_type": {
|
||||
data_type: {
|
||||
"count": len(rows),
|
||||
}
|
||||
for data_type, rows in sorted(by_type.items())
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def summarize_gebench_results(results: list[dict[str, Any]]) -> dict[str, Any]:
|
||||
by_type: dict[str, list[dict[str, Any]]] = defaultdict(list)
|
||||
for result in results:
|
||||
by_type[result["data_type"]].append(result)
|
||||
|
||||
summary: dict[str, Any] = {
|
||||
"count": len(results),
|
||||
"overall_mean": statistics.fmean(r["overall"] for r in results) if results else 0.0,
|
||||
"by_type": {},
|
||||
}
|
||||
for data_type, rows in by_type.items():
|
||||
score_means: dict[str, float] = {}
|
||||
all_score_keys = {key for row in rows for key in row.get("scores", {}).keys()}
|
||||
for score_key in all_score_keys:
|
||||
values = [row["scores"][score_key] for row in rows if score_key in row.get("scores", {})]
|
||||
score_means[score_key] = statistics.fmean(values) if values else 0.0
|
||||
overall_mean = statistics.fmean(row["overall"] for row in rows)
|
||||
summary["by_type"][data_type] = {
|
||||
"count": len(rows),
|
||||
"overall_mean": overall_mean,
|
||||
"overall_mean_100": overall_mean * 100.0,
|
||||
"score_means": score_means,
|
||||
}
|
||||
return summary
|
||||
|
||||
|
||||
def select_balanced_gebench_samples(
|
||||
sample_paths_by_type: dict[str, list[Any]],
|
||||
*,
|
||||
samples_per_type: int | None,
|
||||
) -> dict[str, list[Any]]:
|
||||
if samples_per_type is None:
|
||||
return {data_type: list(paths) for data_type, paths in sample_paths_by_type.items()}
|
||||
return {data_type: list(paths)[:samples_per_type] for data_type, paths in sample_paths_by_type.items()}
|
||||
|
||||
|
||||
def collect_gebench_generation_summary(output_root: Path) -> dict[str, Any]:
|
||||
records: list[dict[str, Any]] = []
|
||||
for data_type, folder_name in TYPE_TO_FOLDER.items():
|
||||
type_root = output_root / folder_name
|
||||
if not type_root.exists():
|
||||
continue
|
||||
for lang_dir in sorted(path for path in type_root.iterdir() if path.is_dir()):
|
||||
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir()):
|
||||
expected = sample_dir / "frame5.png" if data_type in {"type2", "type3", "type4"} else None
|
||||
if expected is None:
|
||||
expected = find_first_image(sample_dir)
|
||||
elif not expected.exists():
|
||||
expected = None
|
||||
if expected is None:
|
||||
continue
|
||||
records.append(
|
||||
{
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_dir.name}/{sample_dir.name}",
|
||||
"output_path": str(expected),
|
||||
}
|
||||
)
|
||||
return summarize_generated_records(records)
|
||||
|
||||
|
||||
def _normalize_score_key(key: str) -> str:
|
||||
mapping = {
|
||||
"goal": "goal",
|
||||
"logic": "logic",
|
||||
"cons": "cons",
|
||||
"consistency": "cons",
|
||||
"ui": "ui",
|
||||
"qual": "qual",
|
||||
"quality": "qual",
|
||||
}
|
||||
return mapping.get(key.lower(), key.lower())
|
||||
|
||||
|
||||
def _normalize_scores(raw_scores: dict[str, Any]) -> dict[str, int]:
|
||||
scores: dict[str, int] = {}
|
||||
for key, value in raw_scores.items():
|
||||
normalized = _normalize_score_key(key)
|
||||
if normalized not in SCORE_KEYS:
|
||||
continue
|
||||
scalar = value.get("s", 0) if isinstance(value, dict) else value
|
||||
try:
|
||||
scores[normalized] = int(scalar)
|
||||
except (TypeError, ValueError):
|
||||
scores[normalized] = 0
|
||||
for key in SCORE_KEYS:
|
||||
scores.setdefault(key, 0)
|
||||
return scores
|
||||
|
||||
|
||||
def _compute_overall(scores: dict[str, int]) -> float:
|
||||
return sum(scores.values()) / (len(SCORE_KEYS) * 5.0)
|
||||
|
||||
|
||||
def _iter_sample_paths(dataset_root: Path, data_type: str) -> list[Path]:
|
||||
data_dir = dataset_root / TYPE_TO_FOLDER[data_type]
|
||||
if not data_dir.exists():
|
||||
raise FileNotFoundError(f"Dataset path not found: {data_dir}")
|
||||
|
||||
samples: list[Path] = []
|
||||
for lang_dir in sorted(path for path in data_dir.iterdir() if path.is_dir()):
|
||||
for child in sorted(lang_dir.iterdir()):
|
||||
if child.is_dir():
|
||||
samples.append(child)
|
||||
elif child.suffix.lower() == ".json":
|
||||
samples.append(child)
|
||||
return samples
|
||||
|
||||
|
||||
def _sample_name_from_metadata(metadata: dict[str, Any], sample_path: Path, item_index: int | None = None) -> str:
|
||||
sample_id = metadata.get("id") or metadata.get("sample_id") or metadata.get("name")
|
||||
if sample_id:
|
||||
return str(sample_id)
|
||||
if item_index is not None:
|
||||
return f"{sample_path.stem}_{item_index:04d}"
|
||||
return sample_path.stem if sample_path.is_file() else sample_path.name
|
||||
|
||||
|
||||
def _expand_sample_path(sample_path: Path) -> list[GEBenchSampleSpec]:
|
||||
if sample_path.is_dir():
|
||||
metadata = _load_metadata(sample_path)
|
||||
return [
|
||||
GEBenchSampleSpec(
|
||||
sample_path=sample_path,
|
||||
metadata=metadata,
|
||||
sample_name=_sample_name(sample_path),
|
||||
lang_device=_lang_device(sample_path, metadata),
|
||||
)
|
||||
]
|
||||
|
||||
payload = load_json(sample_path)
|
||||
if isinstance(payload, dict):
|
||||
return [
|
||||
GEBenchSampleSpec(
|
||||
sample_path=sample_path,
|
||||
metadata=payload,
|
||||
sample_name=_sample_name_from_metadata(payload, sample_path),
|
||||
lang_device=_lang_device(sample_path, payload),
|
||||
)
|
||||
]
|
||||
|
||||
if isinstance(payload, list):
|
||||
specs: list[GEBenchSampleSpec] = []
|
||||
for item_index, item in enumerate(payload):
|
||||
if not isinstance(item, dict):
|
||||
continue
|
||||
specs.append(
|
||||
GEBenchSampleSpec(
|
||||
sample_path=sample_path,
|
||||
metadata=item,
|
||||
sample_name=_sample_name_from_metadata(item, sample_path, item_index),
|
||||
lang_device=_lang_device(sample_path, item),
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
raise TypeError(f"Unsupported metadata payload type for sample {sample_path}: {type(payload)!r}")
|
||||
|
||||
|
||||
def _iter_sample_specs(dataset_root: Path, data_type: str) -> list[GEBenchSampleSpec]:
|
||||
sample_specs: list[GEBenchSampleSpec] = []
|
||||
for sample_path in _iter_sample_paths(dataset_root, data_type):
|
||||
sample_specs.extend(_expand_sample_path(sample_path))
|
||||
return sample_specs
|
||||
|
||||
|
||||
def _load_metadata(sample_path: Path) -> dict[str, Any]:
|
||||
if sample_path.is_file():
|
||||
return load_json(sample_path)
|
||||
for candidate in ("meta_data.json", "metadata.json"):
|
||||
meta_path = sample_path / candidate
|
||||
if meta_path.exists():
|
||||
return load_json(meta_path)
|
||||
raise FileNotFoundError(f"Metadata not found for sample: {sample_path}")
|
||||
|
||||
|
||||
def _sample_name(sample_path: Path) -> str:
|
||||
return sample_path.stem if sample_path.is_file() else sample_path.name
|
||||
|
||||
|
||||
def _lang_device(sample_path: Path, metadata: dict[str, Any]) -> str:
|
||||
return str(metadata.get("lang_device") or sample_path.parent.name)
|
||||
|
||||
|
||||
def _resolve_referenced_image(
|
||||
*,
|
||||
metadata: dict[str, Any],
|
||||
sample_path: Path,
|
||||
dataset_root: Path,
|
||||
data_type: str,
|
||||
) -> Image.Image | None:
|
||||
for key in ("image", "input_image", "initial_image", "reference_image"):
|
||||
image_ref = metadata.get(key)
|
||||
if not image_ref:
|
||||
continue
|
||||
candidate = dataset_root / TYPE_TO_FOLDER[data_type] / str(image_ref)
|
||||
if candidate.exists():
|
||||
image = Image.open(candidate)
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
if sample_path.is_dir():
|
||||
local_image = find_first_image(sample_path)
|
||||
if local_image:
|
||||
image = Image.open(local_image)
|
||||
image.load()
|
||||
return image.convert("RGB")
|
||||
return None
|
||||
|
||||
|
||||
def _trajectory_steps(metadata: dict[str, Any]) -> list[dict[str, Any]]:
|
||||
for key in ("trajectory", "steps", "frames"):
|
||||
value = metadata.get(key)
|
||||
if isinstance(value, list):
|
||||
return [step for step in value if isinstance(step, dict)]
|
||||
extracted: list[dict[str, Any]] = []
|
||||
for index in range(1, 6):
|
||||
value = metadata.get(f"step{index}") or metadata.get(str(index))
|
||||
if isinstance(value, dict):
|
||||
extracted.append(value)
|
||||
return extracted
|
||||
|
||||
|
||||
def _text_or_default(value: Any, default: str = "") -> str:
|
||||
return str(value).strip() if value is not None else default
|
||||
|
||||
|
||||
def _type1_prompt(metadata: dict[str, Any]) -> str:
|
||||
caption = _text_or_default(metadata.get("caption") or metadata.get("instruction"), "Transform the reference GUI.")
|
||||
return (
|
||||
"Using the reference GUI screenshot, generate the next GUI state after the requested interaction.\n\n"
|
||||
f"Requested change:\n{caption}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Preserve layout, visual identity, and unrelated regions.\n"
|
||||
"- Only apply the requested state change.\n"
|
||||
"- Keep all text and controls readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type2_prompt(goal: str, step_num: int) -> str:
|
||||
return (
|
||||
"Generate the next GUI state for a multi-step task.\n\n"
|
||||
f"Overall goal: {goal}\n"
|
||||
f"Current progress step: {step_num}/5\n\n"
|
||||
"Requirements:\n"
|
||||
"- The change should be incremental and plausible.\n"
|
||||
"- Preserve layout and visual identity.\n"
|
||||
"- Make text/buttons readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type34_initial_prompt(metadata: dict[str, Any], first_step: dict[str, Any]) -> str:
|
||||
app_name = _text_or_default(metadata.get("app_name"), "App")
|
||||
final_goal = _text_or_default(metadata.get("final_goal") or metadata.get("instruction"), "Complete the task.")
|
||||
visual_description = _text_or_default(
|
||||
metadata.get("visual_description") or first_step.get("visual_description") or first_step.get("description"),
|
||||
"A clean product-quality app home screen.",
|
||||
)
|
||||
return (
|
||||
"Generate the first GUI frame for a task trajectory.\n\n"
|
||||
f"App name: {app_name}\n"
|
||||
f"Final goal: {final_goal}\n"
|
||||
f"Visual description:\n{visual_description}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Generate a production-looking UI screenshot only.\n"
|
||||
"- Keep the layout coherent and readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type34_next_prompt(step_num: int, step_info: dict[str, Any]) -> str:
|
||||
action = _text_or_default(step_info.get("action") or step_info.get("instruction"), "Continue the task.")
|
||||
visual_description = _text_or_default(
|
||||
step_info.get("visual_description") or step_info.get("description"),
|
||||
"Reflect the expected next GUI state.",
|
||||
)
|
||||
return (
|
||||
"Using the previous frame as reference, generate the next GUI frame.\n\n"
|
||||
f"Step {step_num} action: {action}\n"
|
||||
f"Expected visual state:\n{visual_description}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Only change UI regions affected by this action.\n"
|
||||
"- Preserve persistent bars, layout, and style.\n"
|
||||
"- Keep text and icons readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _type5_prompt(metadata: dict[str, Any]) -> str:
|
||||
grounding = metadata.get("grounding") or {}
|
||||
explanation = _text_or_default(
|
||||
metadata.get("grounding_explanation") or grounding.get("effect") or grounding.get("description"),
|
||||
"Predict the immediate GUI reaction to the indicated target.",
|
||||
)
|
||||
return (
|
||||
"Using the reference GUI screenshot, predict the immediate GUI state after the grounded interaction.\n\n"
|
||||
f"Expected effect: {explanation}\n"
|
||||
f"Grounding metadata: {json.dumps(grounding, ensure_ascii=False)}\n\n"
|
||||
"Requirements:\n"
|
||||
"- Apply only the interaction-triggered change.\n"
|
||||
"- Preserve unrelated regions.\n"
|
||||
"- Keep the UI realistic and readable.\n"
|
||||
)
|
||||
|
||||
|
||||
def _make_storyboard_image(
|
||||
frames: list[Image.Image],
|
||||
*,
|
||||
columns: int = 3,
|
||||
background_color: tuple[int, int, int] = (255, 255, 255),
|
||||
) -> Image.Image:
|
||||
if not frames:
|
||||
raise ValueError("Expected at least one frame to build a storyboard image.")
|
||||
|
||||
normalized = [frame.convert("RGB") for frame in frames]
|
||||
frame_width = max(frame.width for frame in normalized)
|
||||
frame_height = max(frame.height for frame in normalized)
|
||||
rows = (len(normalized) + columns - 1) // columns
|
||||
storyboard = Image.new("RGB", (frame_width * columns, frame_height * rows), color=background_color)
|
||||
|
||||
for index, frame in enumerate(normalized):
|
||||
x_offset = (index % columns) * frame_width
|
||||
y_offset = (index // columns) * frame_height
|
||||
if frame.size != (frame_width, frame_height):
|
||||
frame = frame.resize((frame_width, frame_height))
|
||||
storyboard.paste(frame, (x_offset, y_offset))
|
||||
return storyboard
|
||||
|
||||
|
||||
def _trajectory_judge_payload(frames: list[Image.Image]) -> tuple[str, list[Image.Image]]:
|
||||
storyboard = _make_storyboard_image(frames, columns=3)
|
||||
prompt_suffix = (
|
||||
"The attached image is a storyboard containing six frames arranged left-to-right, "
|
||||
"top-to-bottom as frame0, frame1, frame2, frame3, frame4, frame5."
|
||||
)
|
||||
return prompt_suffix, [storyboard]
|
||||
|
||||
|
||||
class LocalJudgeClient:
|
||||
def __init__(self, base_url: str, api_key: str, model: str, timeout: int = 600):
|
||||
self.base_url = base_url.rstrip("/")
|
||||
self.api_key = api_key
|
||||
self.model = model
|
||||
self.timeout = timeout
|
||||
|
||||
def _build_scoring_prompt(self, task_prompt: str) -> str:
|
||||
return (
|
||||
"You are an expert evaluator for GUI image editing and GUI trajectory generation.\n"
|
||||
"Evaluate whether the generated image(s) satisfy the task.\n\n"
|
||||
"Score these five dimensions from 0 to 5:\n"
|
||||
"- goal: whether the user goal is completed correctly\n"
|
||||
"- logic: whether the transition/state change is logically correct\n"
|
||||
"- cons: whether unrelated regions remain consistent\n"
|
||||
"- ui: whether the UI layout/components remain realistic and coherent\n"
|
||||
"- qual: whether the images are visually clear and artifact-free\n\n"
|
||||
"Return JSON only. Do not add any prose outside JSON.\n"
|
||||
"Use exactly this schema:\n"
|
||||
"{\n"
|
||||
' "goal": 0,\n'
|
||||
' "logic": 0,\n'
|
||||
' "cons": 0,\n'
|
||||
' "ui": 0,\n'
|
||||
' "qual": 0,\n'
|
||||
' "reasoning": "short explanation"\n'
|
||||
"}\n\n"
|
||||
"Scoring task:\n"
|
||||
f"{task_prompt}"
|
||||
)
|
||||
|
||||
def _request_text(self, prompt: str, images: list[Image.Image]) -> str:
|
||||
content: list[dict[str, Any]] = [{"type": "text", "text": prompt}]
|
||||
for image in images:
|
||||
content.append({"type": "image_url", "image_url": {"url": pil_to_data_url(image)}})
|
||||
|
||||
response = requests.post(
|
||||
build_openai_url(self.base_url, "/chat/completions"),
|
||||
json={
|
||||
"model": self.model,
|
||||
"messages": [{"role": "user", "content": content}],
|
||||
"temperature": 0,
|
||||
},
|
||||
headers={
|
||||
"Authorization": f"Bearer {self.api_key}",
|
||||
"Content-Type": "application/json",
|
||||
},
|
||||
timeout=self.timeout,
|
||||
)
|
||||
response.raise_for_status()
|
||||
message_content = response.json()["choices"][0]["message"]["content"]
|
||||
if isinstance(message_content, list):
|
||||
return "\n".join(part.get("text", "") for part in message_content if part.get("type") == "text")
|
||||
return str(message_content)
|
||||
|
||||
def evaluate(self, *, prompt: str, images: list[Image.Image]) -> dict[str, Any]:
|
||||
primary_prompt = self._build_scoring_prompt(prompt)
|
||||
text = self._request_text(primary_prompt, images)
|
||||
try:
|
||||
return extract_json_object(text)
|
||||
except ValueError:
|
||||
retry_prompt = (
|
||||
self._build_scoring_prompt(prompt) + "\n\nYour previous response was not valid JSON. "
|
||||
"Return only the JSON object with integer scores."
|
||||
)
|
||||
retry_text = self._request_text(retry_prompt, images)
|
||||
try:
|
||||
return extract_json_object(retry_text)
|
||||
except ValueError:
|
||||
return {
|
||||
"goal": 0,
|
||||
"logic": 0,
|
||||
"cons": 0,
|
||||
"ui": 0,
|
||||
"qual": 0,
|
||||
"reasoning": retry_text.strip() or text.strip() or "Judge response was not valid JSON.",
|
||||
}
|
||||
|
||||
|
||||
class GEBenchRunner:
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
dataset_root: Path,
|
||||
output_root: Path,
|
||||
base_url: str,
|
||||
model: str,
|
||||
api_key: str = "EMPTY",
|
||||
width: int = 768,
|
||||
height: int = 576,
|
||||
num_inference_steps: int = 8,
|
||||
output_compression: int | None = 98,
|
||||
guidance_scale: float | None = None,
|
||||
seed: int | None = 42,
|
||||
):
|
||||
self.dataset_root = dataset_root
|
||||
self.output_root = output_root
|
||||
self.model = model
|
||||
self.width = width
|
||||
self.height = height
|
||||
self.num_inference_steps = num_inference_steps
|
||||
self.output_compression = output_compression
|
||||
self.guidance_scale = guidance_scale
|
||||
self.seed = seed
|
||||
self.client = VllmOmniImageClient(base_url=base_url, api_key=api_key)
|
||||
|
||||
def generate(
|
||||
self,
|
||||
*,
|
||||
data_type: str,
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_type: int | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
sample_specs = select_balanced_gebench_samples(
|
||||
{data_type: _iter_sample_specs(self.dataset_root, data_type)},
|
||||
samples_per_type=samples_per_type,
|
||||
)[data_type]
|
||||
if max_samples is not None:
|
||||
sample_specs = sample_specs[:max_samples]
|
||||
|
||||
results: list[dict[str, Any]] = []
|
||||
if workers <= 1:
|
||||
for sample_spec in sample_specs:
|
||||
result = self._generate_one(data_type, sample_spec)
|
||||
if result:
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [executor.submit(self._generate_one, data_type, sample_spec) for sample_spec in sample_specs]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
results.append(result)
|
||||
return results
|
||||
|
||||
def _generate_one(self, data_type: str, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
|
||||
sample_path = sample_spec.sample_path
|
||||
metadata = sample_spec.metadata
|
||||
lang_device = sample_spec.lang_device
|
||||
sample_name = sample_spec.sample_name
|
||||
output_dir = ensure_dir(self.output_root / TYPE_TO_FOLDER[data_type] / lang_device / sample_name)
|
||||
|
||||
if data_type == "type1":
|
||||
output_path = output_dir / "generated.png"
|
||||
if output_path.exists():
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
if source is None:
|
||||
return None
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type1_prompt(metadata),
|
||||
images=source,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(output_path, generated)
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
if data_type == "type2":
|
||||
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
if source is None:
|
||||
return None
|
||||
frame0_path = output_dir / "frame0.png"
|
||||
if not frame0_path.exists():
|
||||
save_image(frame0_path, source)
|
||||
previous = source
|
||||
for step_num in range(1, 6):
|
||||
frame_path = output_dir / f"frame{step_num}.png"
|
||||
if frame_path.exists():
|
||||
previous = Image.open(frame_path).convert("RGB")
|
||||
continue
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type2_prompt(goal, step_num),
|
||||
images=previous,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(frame_path, generated)
|
||||
previous = generated
|
||||
output_path = output_dir / "frame5.png"
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
if data_type in {"type3", "type4"}:
|
||||
steps = _trajectory_steps(metadata)
|
||||
frame0_path = output_dir / "frame0.png"
|
||||
if frame0_path.exists():
|
||||
previous = Image.open(frame0_path).convert("RGB")
|
||||
else:
|
||||
previous = self.client.generate_text_to_image(
|
||||
model=self.model,
|
||||
prompt=_type34_initial_prompt(metadata, steps[0] if steps else {}),
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(frame0_path, previous)
|
||||
|
||||
for step_num in range(1, 6):
|
||||
frame_path = output_dir / f"frame{step_num}.png"
|
||||
if frame_path.exists():
|
||||
previous = Image.open(frame_path).convert("RGB")
|
||||
continue
|
||||
step_info = steps[step_num - 1] if step_num - 1 < len(steps) else {}
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type34_next_prompt(step_num, step_info),
|
||||
images=previous,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(frame_path, generated)
|
||||
previous = generated
|
||||
output_path = output_dir / "frame5.png"
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
if data_type == "type5":
|
||||
output_path = output_dir / "generated.png"
|
||||
if output_path.exists():
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=sample_path, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
if source is None:
|
||||
return None
|
||||
generated = self.client.generate_image_edit(
|
||||
model=self.model,
|
||||
prompt=_type5_prompt(metadata),
|
||||
images=source,
|
||||
width=self.width,
|
||||
height=self.height,
|
||||
num_inference_steps=self.num_inference_steps,
|
||||
output_compression=self.output_compression,
|
||||
guidance_scale=self.guidance_scale,
|
||||
seed=self.seed,
|
||||
)
|
||||
save_image(output_path, generated)
|
||||
return {
|
||||
"data_type": data_type,
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"output_path": str(output_path),
|
||||
}
|
||||
|
||||
raise ValueError(f"Unsupported data type: {data_type}")
|
||||
|
||||
|
||||
class GEBenchEvaluator:
|
||||
def __init__(self, *, dataset_root: Path, output_root: Path, judge: LocalJudgeClient):
|
||||
self.dataset_root = dataset_root
|
||||
self.output_root = output_root
|
||||
self.judge = judge
|
||||
|
||||
def evaluate(
|
||||
self,
|
||||
*,
|
||||
data_type: str,
|
||||
workers: int = 1,
|
||||
max_samples: int | None = None,
|
||||
samples_per_type: int | None = None,
|
||||
) -> dict[str, Any]:
|
||||
output_type_dir = self.output_root / TYPE_TO_FOLDER[data_type]
|
||||
sample_specs_by_name = {
|
||||
(spec.lang_device, spec.sample_name): spec for spec in _iter_sample_specs(self.dataset_root, data_type)
|
||||
}
|
||||
if not output_type_dir.exists():
|
||||
payload = {"data_type": data_type, "results": [], "summary": summarize_gebench_results([])}
|
||||
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
|
||||
return payload
|
||||
sample_dirs = [
|
||||
sample_dir
|
||||
for lang_dir in sorted(path for path in output_type_dir.iterdir() if path.is_dir())
|
||||
for sample_dir in sorted(path for path in lang_dir.iterdir() if path.is_dir())
|
||||
if (lang_dir.name, sample_dir.name) in sample_specs_by_name
|
||||
]
|
||||
sample_dirs = select_balanced_gebench_samples(
|
||||
{data_type: sample_dirs},
|
||||
samples_per_type=samples_per_type,
|
||||
)[data_type]
|
||||
if max_samples is not None:
|
||||
sample_dirs = sample_dirs[:max_samples]
|
||||
results: list[dict[str, Any]] = []
|
||||
if workers <= 1:
|
||||
for sample_dir in sample_dirs:
|
||||
result = self._evaluate_one(
|
||||
data_type,
|
||||
sample_dir,
|
||||
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
|
||||
)
|
||||
if result:
|
||||
results.append(result)
|
||||
else:
|
||||
with ThreadPoolExecutor(max_workers=workers) as executor:
|
||||
futures = [
|
||||
executor.submit(
|
||||
self._evaluate_one,
|
||||
data_type,
|
||||
sample_dir,
|
||||
sample_specs_by_name[(sample_dir.parent.name, sample_dir.name)],
|
||||
)
|
||||
for sample_dir in sample_dirs
|
||||
]
|
||||
for future in as_completed(futures):
|
||||
result = future.result()
|
||||
if result:
|
||||
results.append(result)
|
||||
|
||||
payload = {"data_type": data_type, "results": results, "summary": summarize_gebench_results(results)}
|
||||
write_json(self.output_root / "evaluations" / f"{data_type}.json", payload)
|
||||
return payload
|
||||
|
||||
def _evaluate_one(self, data_type: str, sample_dir: Path, sample_spec: GEBenchSampleSpec) -> dict[str, Any] | None:
|
||||
lang_device = sample_dir.parent.name
|
||||
sample_name = sample_dir.name
|
||||
dataset_sample = sample_spec.sample_path
|
||||
metadata = sample_spec.metadata
|
||||
|
||||
if data_type == "type1":
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
generated_path = find_first_image(sample_dir)
|
||||
if source is None or generated_path is None:
|
||||
return None
|
||||
generated = Image.open(generated_path).convert("RGB")
|
||||
raw_scores = self.judge.evaluate(prompt=_type1_prompt(metadata), images=[source, generated])
|
||||
elif data_type == "type2":
|
||||
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
|
||||
goal = _text_or_default(metadata.get("question") or metadata.get("caption"), "Complete the task.")
|
||||
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
|
||||
raw_scores = self.judge.evaluate(
|
||||
prompt=f"Evaluate a six-frame GUI trajectory.\nTask: {goal}\n{prompt_suffix}",
|
||||
images=judge_images,
|
||||
)
|
||||
elif data_type in {"type3", "type4"}:
|
||||
frames = [Image.open(sample_dir / f"frame{i}.png").convert("RGB") for i in range(6)]
|
||||
instruction = _text_or_default(metadata.get("instruction") or metadata.get("caption"), "Complete the task.")
|
||||
prompt_suffix, judge_images = _trajectory_judge_payload(frames)
|
||||
raw_scores = self.judge.evaluate(
|
||||
prompt=f"Evaluate a six-frame GUI trajectory.\nInstruction: {instruction}\n{prompt_suffix}",
|
||||
images=judge_images,
|
||||
)
|
||||
elif data_type == "type5":
|
||||
source = _resolve_referenced_image(
|
||||
metadata=metadata, sample_path=dataset_sample, dataset_root=self.dataset_root, data_type=data_type
|
||||
)
|
||||
generated_path = find_first_image(sample_dir)
|
||||
if source is None or generated_path is None:
|
||||
return None
|
||||
generated = Image.open(generated_path).convert("RGB")
|
||||
raw_scores = self.judge.evaluate(prompt=_type5_prompt(metadata), images=[source, generated])
|
||||
else:
|
||||
raise ValueError(f"Unsupported data type: {data_type}")
|
||||
|
||||
scores = _normalize_scores(raw_scores)
|
||||
return {
|
||||
"sample_name": f"{lang_device}/{sample_name}",
|
||||
"data_type": data_type,
|
||||
"scores": scores,
|
||||
"overall": _compute_overall(scores),
|
||||
"raw_scores": raw_scores,
|
||||
}
|
||||
|
||||
|
||||
def _data_types_arg(value: str) -> list[str]:
|
||||
return list(TYPE_TO_FOLDER.keys()) if value == "all" else [value]
|
||||
|
||||
|
||||
def build_parser() -> argparse.ArgumentParser:
|
||||
parser = argparse.ArgumentParser(description="Run local GEBench generation and scoring against vLLM-Omni.")
|
||||
subparsers = parser.add_subparsers(dest="command", required=True)
|
||||
|
||||
generate = subparsers.add_parser("generate")
|
||||
generate.add_argument("--dataset-root", type=Path, required=True)
|
||||
generate.add_argument("--output-root", type=Path, required=True)
|
||||
generate.add_argument("--base-url", type=str, required=True)
|
||||
generate.add_argument("--model", type=str, required=True)
|
||||
generate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
|
||||
generate.add_argument("--api-key", type=str, default="EMPTY")
|
||||
generate.add_argument("--width", type=int, default=768)
|
||||
generate.add_argument("--height", type=int, default=576)
|
||||
generate.add_argument("--num-inference-steps", type=int, default=8)
|
||||
generate.add_argument("--output-compression", type=int, default=98)
|
||||
generate.add_argument("--guidance-scale", type=float, default=None)
|
||||
generate.add_argument("--seed", type=int, default=42)
|
||||
generate.add_argument("--workers", type=int, default=1)
|
||||
generate.add_argument("--max-samples", type=int, default=None)
|
||||
generate.add_argument("--samples-per-type", type=int, default=None)
|
||||
|
||||
evaluate = subparsers.add_parser("evaluate")
|
||||
evaluate.add_argument("--dataset-root", type=Path, required=True)
|
||||
evaluate.add_argument("--output-root", type=Path, required=True)
|
||||
evaluate.add_argument("--data-type", choices=["all", *TYPE_TO_FOLDER.keys()], default="all")
|
||||
evaluate.add_argument("--judge-base-url", type=str, required=True)
|
||||
evaluate.add_argument("--judge-model", type=str, required=True)
|
||||
evaluate.add_argument("--judge-api-key", type=str, default="EMPTY")
|
||||
evaluate.add_argument("--workers", type=int, default=1)
|
||||
evaluate.add_argument("--max-samples", type=int, default=None)
|
||||
evaluate.add_argument("--samples-per-type", type=int, default=None)
|
||||
|
||||
summarize = subparsers.add_parser("summarize")
|
||||
summarize.add_argument("--output-root", type=Path, required=True)
|
||||
|
||||
return parser
|
||||
|
||||
|
||||
def main(argv: list[str] | None = None) -> int:
|
||||
parser = build_parser()
|
||||
args = parser.parse_args(argv)
|
||||
|
||||
if args.command == "generate":
|
||||
runner = GEBenchRunner(
|
||||
dataset_root=args.dataset_root,
|
||||
output_root=args.output_root,
|
||||
base_url=args.base_url,
|
||||
model=args.model,
|
||||
api_key=args.api_key,
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
output_compression=args.output_compression,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
)
|
||||
records: list[dict[str, Any]] = []
|
||||
for data_type in _data_types_arg(args.data_type):
|
||||
records.extend(
|
||||
runner.generate(
|
||||
data_type=data_type,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_type=args.samples_per_type,
|
||||
)
|
||||
)
|
||||
payload = {"records": records, "summary": summarize_generated_records(records)}
|
||||
write_json(args.output_root / "generation_manifest.json", payload)
|
||||
return 0
|
||||
|
||||
if args.command == "evaluate":
|
||||
judge = LocalJudgeClient(
|
||||
base_url=args.judge_base_url,
|
||||
api_key=args.judge_api_key,
|
||||
model=args.judge_model,
|
||||
)
|
||||
evaluator = GEBenchEvaluator(dataset_root=args.dataset_root, output_root=args.output_root, judge=judge)
|
||||
combined_results: list[dict[str, Any]] = []
|
||||
for data_type in _data_types_arg(args.data_type):
|
||||
payload = evaluator.evaluate(
|
||||
data_type=data_type,
|
||||
workers=args.workers,
|
||||
max_samples=args.max_samples,
|
||||
samples_per_type=args.samples_per_type,
|
||||
)
|
||||
combined_results.extend(payload["results"])
|
||||
_write_json_with_timestamp(
|
||||
args.output_root / "evaluations" / "summary.json",
|
||||
{"summary": summarize_gebench_results(combined_results)},
|
||||
)
|
||||
return 0
|
||||
|
||||
if args.command == "summarize":
|
||||
generation_summary = collect_gebench_generation_summary(args.output_root)
|
||||
evaluation_dir = args.output_root / "evaluations"
|
||||
result_records: list[dict[str, Any]] = []
|
||||
if evaluation_dir.exists():
|
||||
for file_path in sorted(evaluation_dir.glob("type*.json")):
|
||||
payload = load_json(file_path)
|
||||
result_records.extend(payload.get("results", []))
|
||||
payload: dict[str, Any] = {"generation": generation_summary}
|
||||
if result_records:
|
||||
payload["evaluation"] = summarize_gebench_results(result_records)
|
||||
_write_json_with_timestamp(args.output_root / "summary.json", payload)
|
||||
print(json.dumps(payload, indent=2, ensure_ascii=False))
|
||||
return 0
|
||||
|
||||
parser.error(f"Unknown command: {args.command}")
|
||||
return 1
|
||||
@@ -1,13 +0,0 @@
|
||||
# ruff: noqa: E402, I001
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[3]
|
||||
if str(REPO_ROOT) not in sys.path:
|
||||
sys.path.insert(0, str(REPO_ROOT))
|
||||
|
||||
from benchmarks.accuracy.text_to_image.gbench import main
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -43,27 +43,6 @@ The benchmark supports three dataset modes via `--dataset`:
|
||||
|
||||
### VBench dataset
|
||||
|
||||
`vbench` only provides prompt data (and image paths for i2v/i2i); it does not carry
|
||||
per-request generation fields. In this mode, all requests share CLI values:
|
||||
`--width --height --num-frames --fps --num-inference-steps`
|
||||
(pass `--width` and `--height` together).
|
||||
|
||||
Example (`t2v`):
|
||||
|
||||
```bash
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--base-url http://localhost:8099 \
|
||||
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--task t2v \
|
||||
--dataset vbench \
|
||||
--num-prompts 50 \
|
||||
--width 640 --height 480 \
|
||||
--num-frames 81 --fps 16 \
|
||||
--num-inference-steps 40
|
||||
```
|
||||
|
||||
Note: `vbench` can also be used for other tasks such as `t2i` / `i2v` (and `i2i`). For `t2i`, the loader reuses VBench t2v text prompts; for `i2v` / `i2i`, it loads the VBench i2v dataset (with image paths).
|
||||
|
||||
If you use i2v/i2i bench datasets and need auto-download support, you may need:
|
||||
|
||||
```bash
|
||||
@@ -111,7 +90,7 @@ Common optional flags:
|
||||
Related flags: `--width`, `--height`, `--num-frames`, `--fps`, `--num-inference-steps`.
|
||||
|
||||
- For `vbench` / `random`: these CLI flags act as global defaults for all generated requests.
|
||||
- For `trace`: requests can carry their own fields (e.g., `width/height/num_frames/num_inference_steps`), with overrides/fallbacks as below.
|
||||
- For `trace`: each request can carry its own fields (e.g., `width/height/num_frames/num_inference_steps`).
|
||||
|
||||
Precedence rules for `trace` (i.e., what actually gets sent):
|
||||
|
||||
|
||||
@@ -1,4 +1,3 @@
|
||||
import asyncio
|
||||
import base64
|
||||
import mimetypes
|
||||
import os
|
||||
@@ -36,7 +35,6 @@ class RequestFuncOutput:
|
||||
error: str = ""
|
||||
start_time: float = 0.0
|
||||
response_body: dict[str, Any] = field(default_factory=dict)
|
||||
stage_durations: dict[str, float] = field(default_factory=dict)
|
||||
peak_memory_mb: float = 0.0
|
||||
slo_achieved: bool | None = None
|
||||
|
||||
@@ -57,7 +55,6 @@ async def async_request_chat_completions(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: tqdm | None = None,
|
||||
enable_diffusion_pipeline_profiler: bool = False,
|
||||
) -> RequestFuncOutput:
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
@@ -109,19 +106,8 @@ async def async_request_chat_completions(
|
||||
resp_json = await response.json()
|
||||
output.response_body = resp_json
|
||||
output.success = True
|
||||
try:
|
||||
choices = resp_json.get("choices", [])
|
||||
if choices and isinstance(choices, list):
|
||||
msg = choices[0].get("message", {})
|
||||
if isinstance(msg, dict):
|
||||
content = msg.get("content", [])
|
||||
if content and isinstance(content, list) and len(content) > 0:
|
||||
first_item = content[0]
|
||||
if isinstance(first_item, dict):
|
||||
output.stage_durations = first_item.get("stage_durations") or {}
|
||||
output.peak_memory_mb = first_item.get("peak_memory_mb", 0.0)
|
||||
except (IndexError, TypeError, AttributeError):
|
||||
pass
|
||||
if "peak_memory_mb" in resp_json:
|
||||
output.peak_memory_mb = resp_json["peak_memory_mb"]
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
@@ -206,140 +192,7 @@ async def async_request_openai_images(
|
||||
return output
|
||||
|
||||
|
||||
async def async_request_v1_videos(
|
||||
input: RequestFuncInput,
|
||||
session: aiohttp.ClientSession,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestFuncOutput:
|
||||
output = RequestFuncOutput()
|
||||
output.start_time = time.perf_counter()
|
||||
|
||||
files = dict(input.extra_body)
|
||||
if input.prompt:
|
||||
files.setdefault("prompt", input.prompt)
|
||||
if input.width and input.height:
|
||||
files.setdefault("height", input.height)
|
||||
files.setdefault("width", input.width)
|
||||
if input.num_frames:
|
||||
files.setdefault("num_frames", input.num_frames)
|
||||
if input.num_inference_steps:
|
||||
files.setdefault("num_inference_steps", input.num_inference_steps)
|
||||
if input.seed is not None:
|
||||
files.setdefault("seed", input.seed)
|
||||
if input.fps:
|
||||
files.setdefault("fps", input.fps)
|
||||
|
||||
form = aiohttp.FormData()
|
||||
for k, v in files.items():
|
||||
form.add_field(k, str(v))
|
||||
|
||||
image_file = None
|
||||
if input.image_paths and len(input.image_paths) > 0:
|
||||
image_path = input.image_paths[0]
|
||||
image_file = open(image_path, "rb")
|
||||
form.add_field(
|
||||
"input_reference",
|
||||
image_file,
|
||||
filename=os.path.basename(image_path),
|
||||
content_type="application/octet-stream",
|
||||
)
|
||||
|
||||
job_id = None
|
||||
job_status = None
|
||||
poll_json = {}
|
||||
resp_json = {}
|
||||
|
||||
try:
|
||||
# invoke a post request (POST /v1/videos)
|
||||
async with session.post(input.api_url, data=form) as response:
|
||||
if response.status == 200:
|
||||
resp_json = await response.json()
|
||||
job_id = resp_json.get("id")
|
||||
job_status = resp_json.get("status")
|
||||
if not job_id or not job_status:
|
||||
output.error = "API response missing job 'id' or 'status' field."
|
||||
output.success = False
|
||||
return output
|
||||
else:
|
||||
output.error = f"HTTP {response.status}: {await response.text()}"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
# invoke a poll request (GET /v1/videos/{video_id})
|
||||
poll_interval = 2.0 # Unit(s)
|
||||
timeout_seconds = 600.0
|
||||
deadline = time.perf_counter() + timeout_seconds
|
||||
job_url = f"{input.api_url}/{job_id}"
|
||||
|
||||
while job_status not in {"completed", "failed"}:
|
||||
await asyncio.sleep(poll_interval)
|
||||
|
||||
async with session.get(job_url) as poll_response:
|
||||
if poll_response.status != 200:
|
||||
output.error = f"Polling failed HTTP {poll_response.status}: {await poll_response.text()}"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
poll_json = await poll_response.json()
|
||||
job_status = poll_json.get("status")
|
||||
|
||||
if time.perf_counter() >= deadline:
|
||||
output.error = f"Timed out waiting for video job {job_id} to complete."
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
if job_status == "failed":
|
||||
output.error = f"Video job failed: {poll_json}"
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
# invoke a get request (GET /v1/videos/{video_id}/content)
|
||||
content_url = f"{job_url}/content"
|
||||
async with session.get(content_url) as content_response:
|
||||
if content_response.status != 200:
|
||||
output.error = (
|
||||
f"Content retrieval failed HTTP {content_response.status}: {await content_response.text()}"
|
||||
)
|
||||
output.success = False
|
||||
return output
|
||||
|
||||
video_bytes = await content_response.read()
|
||||
output.response_body = video_bytes
|
||||
output.success = True
|
||||
if "peak_memory_mb" in poll_json:
|
||||
output.peak_memory_mb = poll_json["peak_memory_mb"]
|
||||
elif "peak_memory_mb" in resp_json:
|
||||
output.peak_memory_mb = resp_json["peak_memory_mb"]
|
||||
except Exception as e:
|
||||
output.error = str(e)
|
||||
output.success = False
|
||||
finally:
|
||||
if image_file is not None:
|
||||
image_file.close()
|
||||
|
||||
if job_id is not None:
|
||||
try:
|
||||
async with session.delete(f"{input.api_url}/{job_id}") as _:
|
||||
pass
|
||||
except Exception as e:
|
||||
print(f"Failed to clean up video job {job_id}: {e}")
|
||||
|
||||
output.latency = time.perf_counter() - output.start_time
|
||||
|
||||
if output.success and input.slo_ms is not None:
|
||||
output.slo_achieved = (output.latency * 1000.0) <= float(input.slo_ms)
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return output
|
||||
|
||||
|
||||
backends_function_mapping = {
|
||||
"2i": {
|
||||
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
|
||||
"openai": (async_request_openai_images, "/v1/images/generations"),
|
||||
},
|
||||
"2v": {
|
||||
"v1/videos": (async_request_v1_videos, "/v1/videos"),
|
||||
},
|
||||
"vllm-omni": (async_request_chat_completions, "/v1/chat/completions"),
|
||||
"openai": (async_request_openai_images, "/v1/images/generations"),
|
||||
}
|
||||
@@ -1,4 +1,4 @@
|
||||
# adapted from fastvideo
|
||||
# adapted from sglang and fastvideo
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
@@ -9,7 +9,6 @@ If you want to use i2v, i2i dataset, you should `uv pip install gdown` first
|
||||
Supports multiple backends:
|
||||
- vllm-omni: Uses /v1/chat/completions endpoint (default)
|
||||
- openai: Uses /v1/images/generations endpoint
|
||||
- v1/videos: Use /v1/videos endpoint
|
||||
|
||||
Usage:
|
||||
# Video (vllm-omni backend)
|
||||
@@ -29,16 +28,6 @@ Usage:
|
||||
--backend vllm-omni --dataset vbench --task t2i --num-prompts 10 \
|
||||
--height 1024 --width 1024
|
||||
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset random --task t2i --num-prompts 1 \
|
||||
--max-concurrency 1 --enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
|
||||
{"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
|
||||
{"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
|
||||
]'
|
||||
|
||||
i2i:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni --dataset vbench --task i2i --num-prompts 10
|
||||
@@ -49,16 +38,6 @@ Usage:
|
||||
--backend openai --dataset vbench --task t2i --num-prompts 10 \
|
||||
--height 1024 --width 1024 --port 3000
|
||||
|
||||
# Video (v1/vedeos)
|
||||
t2v:
|
||||
python3 benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos --dataset random --task t2v --num-prompts 1 \
|
||||
--max-concurrency 1 --enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":854,"height":480,"num_inference_steps":18,"num_frames":120,"fps":24,"weight":1}
|
||||
]'
|
||||
|
||||
|
||||
"""
|
||||
|
||||
import argparse
|
||||
@@ -66,10 +45,7 @@ import ast
|
||||
import asyncio
|
||||
import glob
|
||||
import json
|
||||
import logging
|
||||
import os
|
||||
import random
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from abc import ABC, abstractmethod
|
||||
@@ -81,11 +57,8 @@ import aiohttp
|
||||
import numpy as np
|
||||
import requests
|
||||
from backends import RequestFuncInput, RequestFuncOutput, backends_function_mapping
|
||||
from PIL import Image
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
class BaseDataset(ABC):
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
@@ -289,9 +262,6 @@ class VBenchDataset(BaseDataset):
|
||||
|
||||
def _resize_data(self, data: list[dict[str, Any]]) -> list[dict[str, Any]]:
|
||||
"""Resize data to match num_prompts."""
|
||||
if not data:
|
||||
raise ValueError("No benchmark data available. Install Pillow or provide --dataset-path.")
|
||||
|
||||
if not self.args.num_prompts:
|
||||
return data
|
||||
|
||||
@@ -515,14 +485,6 @@ class TraceDataset(BaseDataset):
|
||||
timestamp = self._coerce_float(row.get("timestamp"))
|
||||
slo_ms = self._coerce_float(row.get("slo_ms"))
|
||||
image_paths = row.get("image_paths")
|
||||
if not image_paths:
|
||||
single = row.get("image_path")
|
||||
image_paths = [single] if single else None
|
||||
|
||||
if not image_paths and self.args.task in ["i2v", "i2i", "ti2v", "ti2i"]:
|
||||
raise ValueError(
|
||||
f"Task {self.args.task} requires image input, but no image_path or image_paths found in trace row."
|
||||
)
|
||||
|
||||
override_w = self.args.width
|
||||
override_h = self.args.height
|
||||
@@ -554,66 +516,26 @@ class TraceDataset(BaseDataset):
|
||||
|
||||
|
||||
class RandomDataset(BaseDataset):
|
||||
def __init__(self, args, api_url: str, model: str, enable_negative_prompt: bool = False):
|
||||
super().__init__(args, api_url, model)
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
self.args = args
|
||||
self.api_url = api_url
|
||||
self.model = model
|
||||
self.num_prompts = args.num_prompts
|
||||
self.enable_negative_prompt = enable_negative_prompt
|
||||
self.random_request_config = getattr(args, "random_request_config", None)
|
||||
if self.random_request_config:
|
||||
self.random_request_config = json.loads(self.random_request_config)
|
||||
self._weights = [p["weight"] for p in self.random_request_config]
|
||||
|
||||
self.random_request_config = [
|
||||
{k: v for k, v in p.items() if k != "weight"} for p in self.random_request_config
|
||||
]
|
||||
|
||||
seed = getattr(args, "random_request_seed", 42)
|
||||
self._rng = random.Random(seed)
|
||||
|
||||
self._sampled_requests = self._rng.choices(
|
||||
self.random_request_config,
|
||||
weights=self._weights,
|
||||
k=self.num_prompts,
|
||||
)
|
||||
else:
|
||||
self._sampled_requests = None
|
||||
|
||||
# Random image generate
|
||||
if self.args.task in ["i2v", "ti2v", "ti2i", "i2i"]:
|
||||
img = Image.new("RGB", (512, 512), (255, 255, 255))
|
||||
|
||||
image_path = os.path.join(tempfile.gettempdir(), "diffusion_benchmark_random_image.png")
|
||||
self._random_image_path = [image_path]
|
||||
img.save(image_path)
|
||||
else:
|
||||
self._random_image_path = None
|
||||
|
||||
def __len__(self) -> int:
|
||||
return self.num_prompts
|
||||
|
||||
def __getitem__(self, idx: int) -> RequestFuncInput:
|
||||
extra_body = {}
|
||||
if self.enable_negative_prompt:
|
||||
extra_body["negative_prompt"] = f"Negative prompt {idx} for benchmarking diffusion models"
|
||||
|
||||
params = {
|
||||
"width": self.args.width,
|
||||
"height": self.args.height,
|
||||
"num_frames": self.args.num_frames,
|
||||
"num_inference_steps": self.args.num_inference_steps,
|
||||
"fps": self.args.fps,
|
||||
}
|
||||
if self._sampled_requests:
|
||||
profile = self._sampled_requests[idx]
|
||||
params.update(profile)
|
||||
return RequestFuncInput(
|
||||
prompt=f"Random prompt {idx} for benchmarking diffusion models",
|
||||
api_url=self.api_url,
|
||||
model=self.model,
|
||||
width=self.args.width,
|
||||
height=self.args.height,
|
||||
num_frames=self.args.num_frames,
|
||||
num_inference_steps=self.args.num_inference_steps,
|
||||
seed=self.args.seed,
|
||||
extra_body=extra_body,
|
||||
image_paths=self._random_image_path,
|
||||
**params,
|
||||
fps=self.args.fps,
|
||||
)
|
||||
|
||||
def get_requests(self) -> list[RequestFuncInput]:
|
||||
@@ -725,19 +647,19 @@ async def iter_requests(
|
||||
requests_list: list[RequestFuncInput],
|
||||
request_rate: float,
|
||||
) -> AsyncGenerator[RequestFuncInput, None]:
|
||||
"""Yield requests using a Poisson process if request_rate is set.
|
||||
"""Yield requests using a fixed interval if request_rate is set.
|
||||
|
||||
- If request_rate is inf, all requests are yielded immediately (no sleep).
|
||||
- Otherwise, inter-arrival times follow an exponential distribution.
|
||||
- Otherwise, requests are emitted at a fixed cadence of 1 / request_rate seconds.
|
||||
"""
|
||||
|
||||
if request_rate != float("inf"):
|
||||
if request_rate <= 0:
|
||||
raise ValueError(f"request_rate must be positive or inf, got {request_rate}.")
|
||||
interval_s = 1.0 / float(request_rate)
|
||||
|
||||
for i, req in enumerate(requests_list):
|
||||
if request_rate != float("inf") and i > 0:
|
||||
interval_s = random.expovariate(request_rate)
|
||||
await asyncio.sleep(interval_s)
|
||||
yield req
|
||||
|
||||
@@ -756,15 +678,6 @@ def calculate_metrics(
|
||||
latencies = [o.latency for o in success_outputs]
|
||||
peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb > 0]
|
||||
|
||||
# Aggregate per-stage durations across all successful requests that reported them.
|
||||
stage_duration_lists: dict[str, list[float]] = {}
|
||||
for o in success_outputs:
|
||||
for stage, duration in (o.stage_durations or {}).items():
|
||||
stage_duration_lists.setdefault(stage, []).append(duration)
|
||||
stage_durations_mean = {s: float(np.mean(v)) for s, v in stage_duration_lists.items()}
|
||||
stage_durations_p50 = {s: float(np.percentile(v, 50)) for s, v in stage_duration_lists.items()}
|
||||
stage_durations_p99 = {s: float(np.percentile(v, 99)) for s, v in stage_duration_lists.items()}
|
||||
|
||||
metrics = {
|
||||
"duration": total_duration,
|
||||
"completed_requests": num_success,
|
||||
@@ -773,14 +686,10 @@ def calculate_metrics(
|
||||
"latency_mean": np.mean(latencies) if latencies else 0,
|
||||
"latency_median": np.median(latencies) if latencies else 0,
|
||||
"latency_p99": np.percentile(latencies, 99) if latencies else 0,
|
||||
"latency_p95": np.percentile(latencies, 95) if latencies else 0,
|
||||
"latency_p50": np.percentile(latencies, 50) if latencies else 0,
|
||||
"peak_memory_mb_max": max(peak_memories) if peak_memories else 0,
|
||||
"peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0,
|
||||
"peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0,
|
||||
"stage_durations_mean": stage_durations_mean,
|
||||
"stage_durations_p50": stage_durations_p50,
|
||||
"stage_durations_p99": stage_durations_p99,
|
||||
}
|
||||
|
||||
if slo_enabled:
|
||||
@@ -833,32 +742,8 @@ async def benchmark(args):
|
||||
if args.base_url is None:
|
||||
args.base_url = f"http://{args.host}:{args.port}"
|
||||
|
||||
VIDEO_TASKS = {"t2v", "i2v", "ti2v"}
|
||||
IMAGE_TASKS = {"t2i", "i2i", "ti2i"}
|
||||
|
||||
if args.task in VIDEO_TASKS:
|
||||
task_type = "2v"
|
||||
elif args.task in IMAGE_TASKS:
|
||||
task_type = "2i"
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Unsupported task: '{args.task}'. "
|
||||
f"Valid video tasks: {sorted(VIDEO_TASKS)}, "
|
||||
f"Valid image tasks: {sorted(IMAGE_TASKS)}"
|
||||
)
|
||||
|
||||
valid_backends = sorted(backends_function_mapping[task_type].keys())
|
||||
|
||||
if args.backend not in valid_backends:
|
||||
logger.error(
|
||||
f"Invalid backend '{args.backend}' for task '{args.task}' (task type: '{task_type}').\n"
|
||||
f"Valid backends for this task type: {valid_backends}\n"
|
||||
f"Example usage: --task {args.task} --backend {valid_backends[0]}"
|
||||
)
|
||||
raise ValueError("Backend validation failed. See log above for valid options.")
|
||||
|
||||
# Setup API URL and request function based on backend
|
||||
request_func, api_url = backends_function_mapping[task_type][args.backend]
|
||||
request_func, api_url = backends_function_mapping[args.backend]
|
||||
api_url = f"{args.base_url}{api_url}"
|
||||
|
||||
if args.dataset == "vbench":
|
||||
@@ -866,7 +751,7 @@ async def benchmark(args):
|
||||
elif args.dataset == "trace":
|
||||
dataset = TraceDataset(args, api_url, args.model)
|
||||
elif args.dataset == "random":
|
||||
dataset = RandomDataset(args, api_url, args.model, args.enable_negative_prompt)
|
||||
dataset = RandomDataset(args, api_url, args.model)
|
||||
else:
|
||||
raise ValueError(f"Unknown dataset: {args.dataset}")
|
||||
|
||||
@@ -904,8 +789,6 @@ async def benchmark(args):
|
||||
warm_req,
|
||||
num_inference_steps=args.warmup_num_inference_steps,
|
||||
)
|
||||
if args.task == "t2v":
|
||||
warm_req = replace(warm_req, num_frames=1)
|
||||
warm_out = await limited_request_func(warm_req, session, None)
|
||||
warmup_pairs.append((warm_req, warm_out))
|
||||
|
||||
@@ -964,7 +847,6 @@ async def benchmark(args):
|
||||
print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"]))
|
||||
print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"]))
|
||||
print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"]))
|
||||
print("{:<40} {:<15.4f}".format("Latency P95 (s):", metrics["latency_p95"]))
|
||||
|
||||
if args.slo:
|
||||
print(f"{'-' * 50}")
|
||||
@@ -978,12 +860,6 @@ async def benchmark(args):
|
||||
print("{:<40} {:<15.2f}".format("Peak Memory Mean (MB):", metrics["peak_memory_mb_mean"]))
|
||||
print("{:<40} {:<15.2f}".format("Peak Memory Median (MB):", metrics["peak_memory_mb_median"]))
|
||||
|
||||
if metrics["stage_durations_mean"]:
|
||||
print(f"{'-' * 50}")
|
||||
print("Stage Durations Mean (s):")
|
||||
for stage, val in metrics["stage_durations_mean"].items():
|
||||
print("{:<40} {:<15.4f}".format(f" {stage}:", val))
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
|
||||
if args.output_file:
|
||||
@@ -1007,7 +883,7 @@ if __name__ == "__main__":
|
||||
"--backend",
|
||||
type=str,
|
||||
default="vllm-omni",
|
||||
choices=["vllm-omni", "openai", "v1/videos"],
|
||||
choices=["vllm-omni", "openai"],
|
||||
help="Backend to target the benchmark to.",
|
||||
)
|
||||
parser.add_argument(
|
||||
@@ -1096,25 +972,6 @@ if __name__ == "__main__":
|
||||
help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
|
||||
)
|
||||
parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
|
||||
parser.add_argument(
|
||||
"--enable-negative-prompt",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Generate negative prompts when using the random dataset.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--random-request-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help=(
|
||||
"JSON string defining random request profiles. "
|
||||
"Each profile may contain: width, height, num_inference_steps, etc. "
|
||||
"The 'weight' field controls sampling probability (relative weight). "
|
||||
"Example: "
|
||||
'[{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},'
|
||||
'{"width":768,"height":768,"num_inference_steps":20,"weight":0.85}]'
|
||||
),
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
|
||||
@@ -1,169 +0,0 @@
|
||||
# Qwen-Image Serving Performance Dashboard
|
||||
|
||||
This document describes how to deploy and benchmark **Qwen-Image** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
|
||||
|
||||
---
|
||||
|
||||
# 1. Overview
|
||||
|
||||
Qwen-Image is a multimodal text-to-image generation model served through the vLLM-Omni infrastructure.
|
||||
|
||||
This document covers:
|
||||
|
||||
* Service launch configuration (including acceleration options)
|
||||
* Benchmark scripts and usage
|
||||
* Dataset and workload settings
|
||||
* Performance measurement results
|
||||
* Reproducibility guidelines
|
||||
|
||||
---
|
||||
|
||||
# 2. Test Environment
|
||||
| Component | Specification |
|
||||
|------------|----------------|
|
||||
| GPU | NVIDIA A100-SXM4-80GB |
|
||||
| Diffusion Attention Backend | FlashAttention |
|
||||
|
||||
# 3. Service Launch Configuration
|
||||
|
||||
## 3.1 Basic Serving Command
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen-Image --omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
## 3.2 Key Parameters
|
||||
|
||||
| Parameter | Description |
|
||||
| --------------------- | ------------------------ |
|
||||
| `--cfg-parallel-size` | CFG parallelism degree |
|
||||
| `--ulysses-degree` | Ulysses parallel degree |
|
||||
| `--vae-patch-parallel-size` | VAE parallel degree |
|
||||
| `--tensor-parallel-size` | Tensor parallelism degree |
|
||||
|
||||
Record these parameters when reporting performance results.
|
||||
|
||||
---
|
||||
|
||||
# 4. Benchmark Script
|
||||
|
||||
## 4.1 Benchmark Entry
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni \
|
||||
--dataset <DATASET_NAME> \
|
||||
--task t2i \
|
||||
--num-prompts <N> \
|
||||
--max-concurrency <C> \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config <CFG>
|
||||
```
|
||||
|
||||
## 4.2 Key Benchmark Arguments
|
||||
|
||||
| Parameter | Description |
|
||||
| ---------------------- | --------------------------------- |
|
||||
| `--backend` | Serving backend (use `vllm-omni`) |
|
||||
| `--dataset` | Dataset name (`random` or custom) |
|
||||
| `--task` | Task type (e.g., `t2i`) |
|
||||
| `--num-prompts` | Total number of requests |
|
||||
| `--max-concurrency` | Client-side concurrency |
|
||||
| `--random-request-config`| JSON string defining random request |
|
||||
|
||||
---
|
||||
|
||||
# 5. Dataset & Workload Settings
|
||||
|
||||
## 5.1 Recommended Evaluation Configurations
|
||||
|
||||
### Dataset A ( 512 Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2i
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":1}
|
||||
]
|
||||
```
|
||||
|
||||
### Dataset B (1536 Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2i
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":1}
|
||||
]
|
||||
```
|
||||
|
||||
### Dataset C (Mix Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2i
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":0.15},
|
||||
{"width":768,"height":768,"num_inference_steps":20,"weight":0.25},
|
||||
{"width":1024,"height":1024,"num_inference_steps":25,"weight":0.45},
|
||||
{"width":1536,"height":1536,"num_inference_steps":35,"weight":0.15}
|
||||
]
|
||||
```
|
||||
---
|
||||
|
||||
## 5.2 Example Benchmark Command
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend vllm-omni \
|
||||
--dataset random \
|
||||
--task t2i \
|
||||
--num-prompts 1 \
|
||||
--max-concurrency 1 \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":512,"height":512,"num_inference_steps":20,"weight":1}
|
||||
]'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
# 6. Performance Metrics
|
||||
|
||||
The following metrics are collected during benchmarking:
|
||||
|
||||
| Metric | Description | Unit |
|
||||
| ------------------ | ----------------------------- | ------- |
|
||||
| Mean Latency | Mean of latency | seconds |
|
||||
| P99 Latency | P99 of latency | seconds |
|
||||
|
||||
---
|
||||
|
||||
# 7. Performance Results
|
||||
|
||||
| Dataset Configuration | Max Concur. | CFG | Usp | Tp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
|
||||
|-----------------------|-----|-----|-----|----|--------------|------------------|------------------|
|
||||
| Dataset A | 1 | 2 | 2 | Off | Off | 2.2087 | 2.2087 |
|
||||
| Dataset B | 1 | 2 | 2 | Off | Off | 19.6739 | 19.6739 |
|
||||
| Dataset C | 1 | 2 | 2 | Off | Off | 5.67259 | 18.6234 |
|
||||
---
|
||||
|
||||
# 8. Reproducibility Checklist
|
||||
|
||||
To ensure consistent and comparable benchmark results:
|
||||
|
||||
* Record GPU type
|
||||
* Record parallel configuration
|
||||
* Record benchmark parameters (resolution, concurrency, number of prompts)
|
||||
* Ensure no background workload on GPUs during testing
|
||||
|
||||
---
|
||||
|
||||
This document serves as the official Qwen-Image serving performance reference under vLLM-Omni.
|
||||
@@ -1,170 +0,0 @@
|
||||
# Wan2.2 Serving Performance Dashboard
|
||||
|
||||
This document describes how to deploy and benchmark **Wan-AI/Wan2.2-T2V-A14B-Diffusers** using vLLM-Omni. It includes service startup configuration, acceleration-related options, benchmark methodology, dataset settings, and performance results.
|
||||
|
||||
---
|
||||
|
||||
# 1. Overview
|
||||
|
||||
Wan-AI/Wan2.2-T2V-A14B-Diffusers is a multimodal text-to-video generation model served through the vLLM-Omni infrastructure.
|
||||
|
||||
This document covers:
|
||||
|
||||
* Service launch configuration (including acceleration options)
|
||||
* Benchmark scripts and usage
|
||||
* Dataset and workload settings
|
||||
* Performance measurement results
|
||||
* Reproducibility guidelines
|
||||
|
||||
---
|
||||
|
||||
# 2. Test Environment
|
||||
| Component | Specification |
|
||||
|------------|----------------|
|
||||
| GPU | NVIDIA A100-SXM4-80GB |
|
||||
| Diffusion Attention Backend | FlashAttention |
|
||||
|
||||
# 3. Service Launch Configuration
|
||||
|
||||
## 3.1 Basic Serving Command
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
## 3.2 Key Parameters
|
||||
|
||||
| Parameter | Description |
|
||||
| --------------------- | ------------------------ |
|
||||
| `--cfg-parallel-size` | CFG parallelism degree |
|
||||
| `--ulysses-degree` | Ulysses parallel degree |
|
||||
| `--vae-patch-parallel-size` | VAE parallel degree |
|
||||
| `--tensor-parallel-size` | Tensor parallelism degree |
|
||||
| `--use-hsdp` | Enable Hybrid Sharded Data Parallel to shard model weights across GPUs |
|
||||
|
||||
Record these parameters when reporting performance results.
|
||||
|
||||
---
|
||||
|
||||
# 4. Benchmark Script
|
||||
|
||||
## 4.1 Benchmark Entry
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos \
|
||||
--dataset <DATASET_NAME> \
|
||||
--task t2v \
|
||||
--num-prompts <N> \
|
||||
--max-concurrency <C> \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config <CFG>
|
||||
```
|
||||
|
||||
## 4.2 Key Benchmark Arguments
|
||||
|
||||
| Parameter | Description |
|
||||
| ---------------------- | --------------------------------- |
|
||||
| `--backend` | Serving backend (use `v1/videos`) |
|
||||
| `--dataset` | Dataset name (`random` or custom) |
|
||||
| `--task` | Task type (e.g., `t2v`) |
|
||||
| `--num-prompts` | Total number of requests |
|
||||
| `--max-concurrency` | Client-side concurrency |
|
||||
| `--random-request-config`| JSON string defining random request |
|
||||
|
||||
---
|
||||
|
||||
# 5. Dataset & Workload Settings
|
||||
|
||||
## 5.1 Recommended Evaluation Configurations
|
||||
|
||||
### Dataset A (480p)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2v
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":1}
|
||||
]
|
||||
```
|
||||
### Dataset B (720p)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2v
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":1}
|
||||
]
|
||||
```
|
||||
### Dataset C (Mix Resolution)
|
||||
|
||||
* Dataset: `random`
|
||||
* Task: t2v
|
||||
* Concurrency: 1
|
||||
* Mix Resolution
|
||||
```
|
||||
[
|
||||
{"width":854,"height":480,"num_inference_steps":3,"num_frames":80,"fps":16,"weight":0.15},
|
||||
{"width":854,"height":480,"num_inference_steps":4,"num_frames":120,"fps":24,"weight":0.25},
|
||||
{"width":1280,"height":720,"num_inference_steps":6,"num_frames":80,"fps":16,"weight":0.6}
|
||||
]
|
||||
```
|
||||
---
|
||||
|
||||
## 5.2 Example Benchmark Command
|
||||
|
||||
```bash
|
||||
python benchmarks/diffusion/diffusion_benchmark_serving.py \
|
||||
--backend v1/videos \
|
||||
--dataset random \
|
||||
--task t2v \
|
||||
--num-prompts 1 \
|
||||
--max-concurrency 1 \
|
||||
--enable-negative-prompt \
|
||||
--random-request-config '[
|
||||
{"width":854,"height":480,"num_inference_steps":18,"num_frames": 33,"fps":16",weight":1}
|
||||
]'
|
||||
```
|
||||
|
||||
---
|
||||
|
||||
# 6. Performance Metrics
|
||||
|
||||
The following metrics are collected during benchmarking:
|
||||
|
||||
| Metric | Description | Unit |
|
||||
| ------------------ | ----------------------------- | ------- |
|
||||
| Mean Latency | Mean of latency | seconds |
|
||||
| P99 Latency | P99 of latency | seconds |
|
||||
|
||||
---
|
||||
|
||||
# 7. Performance Results
|
||||
|
||||
| Dataset Configuration | Max Concur. | CFG | Usp | Tp | Hsdp | VAE Parallel | Mean Latency (s) | P99 Latency (s) |
|
||||
|-----------------------|-----|-----|-----|-----|----|--------------|------------------|------------------|
|
||||
| Dataset A | 1 | 2 | 2 | 1 | On | 1 | 24.6766 | 24.6766 |
|
||||
| Dataset A | 1 | 2 | 2 | 1 | On | 4 | 21.6810 | 21.6810 |
|
||||
| Dataset B | 1 | 2 | 2 | 1 | On | 1 | 124.6639 | 124.6639 |
|
||||
| Dataset B | 1 | 2 | 2 | 1 | On | 4 | 117.44 | 117.44 |
|
||||
| Dataset C | 1 | 2 | 2 | 1 | On | 1 | 79.2175 | 124.2565 |
|
||||
| Dataset C | 1 | 2 | 2 | 1 | On | 4 | 74.4977 | 117.710 |
|
||||
---
|
||||
|
||||
# 8. Reproducibility Checklist
|
||||
|
||||
To ensure consistent and comparable benchmark results:
|
||||
|
||||
* Record GPU type
|
||||
* Record parallel configuration
|
||||
* Record benchmark parameters (resolution, concurrency, number of prompts)
|
||||
* Ensure no background workload on GPUs during testing
|
||||
|
||||
---
|
||||
|
||||
This document serves as the official Wan2.2 serving performance reference under vLLM-Omni.
|
||||
@@ -1,460 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Benchmark quantization quality loss for diffusion models (image & video).
|
||||
|
||||
Generates outputs with BF16 (baseline) and a quantized config using the same
|
||||
seed, then computes LPIPS perceptual distance between them. Results are printed
|
||||
as a Markdown table ready to paste into a PR description.
|
||||
|
||||
Requirements:
|
||||
pip install lpips Pillow numpy
|
||||
|
||||
Image example (text-to-image):
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--task t2i \
|
||||
--quantization fp8 \
|
||||
--prompts \
|
||||
"an aerial view of a coral reef with crystal clear turquoise water" \
|
||||
"a campfire in a dark forest with sparks rising into a starry sky" \
|
||||
"a gourmet dessert plate with chocolate mousse and gold leaf" \
|
||||
--height 1024 --width 1024 \
|
||||
--num-inference-steps 50 --seed 42
|
||||
|
||||
Video example (text-to-video):
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Wan-AI/Wan2.2-T2V-A14B-Diffusers \
|
||||
--task t2v \
|
||||
--quantization fp8 \
|
||||
--prompts \
|
||||
"A serene lakeside sunrise with mist over the water" \
|
||||
"A cat walking across a wooden bridge in autumn" \
|
||||
--height 720 --width 1280 \
|
||||
--num-frames 81 --num-inference-steps 40 --seed 42
|
||||
|
||||
Multiple quantization methods:
|
||||
python benchmarks/diffusion/quantization_quality.py \
|
||||
--model Tongyi-MAI/Z-Image-Turbo \
|
||||
--task t2i \
|
||||
--quantization fp8 int8 bitsandbytes \
|
||||
--prompts "a cup of coffee on the table" \
|
||||
--height 1024 --width 1024 \
|
||||
--num-inference-steps 50 --seed 42
|
||||
|
||||
Output directory structure (--output-dir, default: ./quant_bench_output):
|
||||
quant_bench_output/
|
||||
baseline/ # BF16 outputs
|
||||
<method>/ # Quantized outputs per method
|
||||
results.md # Markdown table
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import gc
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
|
||||
def compute_lpips_images(
|
||||
baseline_images: list,
|
||||
quantized_images: list,
|
||||
net: str = "alex",
|
||||
) -> list[float]:
|
||||
"""Compute LPIPS between paired lists of PIL images."""
|
||||
import lpips
|
||||
from torchvision import transforms
|
||||
|
||||
loss_fn = lpips.LPIPS(net=net).eval()
|
||||
if torch.cuda.is_available():
|
||||
loss_fn = loss_fn.cuda()
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((256, 256)),
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
|
||||
scores = []
|
||||
for img_bl, img_qt in zip(baseline_images, quantized_images):
|
||||
t_bl = transform(img_bl.convert("RGB")).unsqueeze(0)
|
||||
t_qt = transform(img_qt.convert("RGB")).unsqueeze(0)
|
||||
if torch.cuda.is_available():
|
||||
t_bl, t_qt = t_bl.cuda(), t_qt.cuda()
|
||||
with torch.no_grad():
|
||||
score = loss_fn(t_bl, t_qt).item()
|
||||
scores.append(score)
|
||||
return scores
|
||||
|
||||
|
||||
def compute_lpips_video(
|
||||
baseline_frames: np.ndarray,
|
||||
quantized_frames: np.ndarray,
|
||||
net: str = "alex",
|
||||
) -> float:
|
||||
"""Compute mean per-frame LPIPS for a video pair.
|
||||
|
||||
Args:
|
||||
baseline_frames: (F, H, W, C) float array in [0, 1].
|
||||
quantized_frames: same shape.
|
||||
|
||||
Returns:
|
||||
Mean LPIPS across all frames.
|
||||
"""
|
||||
import lpips
|
||||
|
||||
loss_fn = lpips.LPIPS(net=net).eval()
|
||||
if torch.cuda.is_available():
|
||||
loss_fn = loss_fn.cuda()
|
||||
|
||||
num_frames = min(len(baseline_frames), len(quantized_frames))
|
||||
scores = []
|
||||
for i in range(num_frames):
|
||||
# Convert (H, W, C) float [0,1] -> (1, C, H, W) float [-1, 1]
|
||||
f_bl = torch.from_numpy(baseline_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
|
||||
f_qt = torch.from_numpy(quantized_frames[i]).permute(2, 0, 1).unsqueeze(0).float() * 2 - 1
|
||||
if torch.cuda.is_available():
|
||||
f_bl, f_qt = f_bl.cuda(), f_qt.cuda()
|
||||
with torch.no_grad():
|
||||
score = loss_fn(f_bl, f_qt).item()
|
||||
scores.append(score)
|
||||
return float(np.mean(scores))
|
||||
|
||||
|
||||
def _build_omni_kwargs(args, quantization=None):
|
||||
"""Build kwargs dict for Omni() constructor."""
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
)
|
||||
kwargs = {
|
||||
"model": args.model,
|
||||
"parallel_config": parallel_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
}
|
||||
if quantization:
|
||||
kwargs["quantization_config"] = quantization
|
||||
return kwargs
|
||||
|
||||
|
||||
def _generate_image(omni, args, prompt, seed):
|
||||
"""Generate a single image and return (PIL.Image, time_seconds, memory_gib)."""
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
{"prompt": prompt},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
req_out = first.request_output[0] if hasattr(first, "request_output") else first
|
||||
img = req_out.images[0]
|
||||
return img, elapsed, peak_mem
|
||||
|
||||
|
||||
def _generate_video(omni, args, prompt, seed):
|
||||
"""Generate a video and return (np.ndarray [F,H,W,C], time_seconds, memory_gib)."""
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(seed)
|
||||
torch.cuda.reset_peak_memory_stats()
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
{"prompt": prompt, "negative_prompt": ""},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
peak_mem = torch.cuda.max_memory_allocated() / (1024**3)
|
||||
|
||||
first = outputs[0]
|
||||
if hasattr(first, "request_output") and isinstance(first.request_output, list):
|
||||
inner = first.request_output[0]
|
||||
if isinstance(inner, OmniRequestOutput) and hasattr(inner, "images"):
|
||||
frames = inner.images[0] if inner.images else None
|
||||
else:
|
||||
frames = inner
|
||||
elif hasattr(first, "images") and first.images:
|
||||
frames = first.images
|
||||
else:
|
||||
raise ValueError("Could not extract video frames from output.")
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video = frames.detach().cpu()
|
||||
if video.dim() == 5:
|
||||
video = video[0].permute(1, 2, 3, 0) if video.shape[1] in (3, 4) else video[0]
|
||||
elif video.dim() == 4 and video.shape[0] in (3, 4):
|
||||
video = video.permute(1, 2, 3, 0)
|
||||
if video.is_floating_point():
|
||||
video = video.clamp(-1, 1) * 0.5 + 0.5
|
||||
frames_array = video.float().numpy()
|
||||
else:
|
||||
frames_array = np.asarray(frames)
|
||||
if frames_array.ndim == 5:
|
||||
frames_array = frames_array[0]
|
||||
|
||||
return frames_array, elapsed, peak_mem
|
||||
|
||||
|
||||
def _unload_omni(omni):
|
||||
"""Delete Omni instance and free GPU memory."""
|
||||
del omni
|
||||
gc.collect()
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.empty_cache()
|
||||
torch.cuda.synchronize()
|
||||
|
||||
|
||||
def run_benchmark(args):
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
is_video = args.task == "t2v"
|
||||
prompts = args.prompts
|
||||
seed = args.seed
|
||||
|
||||
# Determine configs to benchmark
|
||||
configs = [] # list of (label, quantization_method)
|
||||
for method in args.quantization:
|
||||
configs.append((method, method))
|
||||
|
||||
# --- Baseline run ---
|
||||
print("\n" + "=" * 60)
|
||||
print("Running BF16 baseline...")
|
||||
print("=" * 60)
|
||||
bl_kwargs = _build_omni_kwargs(args, quantization=None)
|
||||
omni_bl = Omni(**bl_kwargs)
|
||||
|
||||
baseline_outputs = {} # prompt -> (output, time, mem)
|
||||
for prompt in prompts:
|
||||
print(f" Generating: {prompt[:60]}...")
|
||||
if is_video:
|
||||
out, t, mem = _generate_video(omni_bl, args, prompt, seed)
|
||||
else:
|
||||
out, t, mem = _generate_image(omni_bl, args, prompt, seed)
|
||||
baseline_outputs[prompt] = (out, t, mem)
|
||||
|
||||
bl_avg_time = np.mean([v[1] for v in baseline_outputs.values()])
|
||||
bl_mem = baseline_outputs[prompts[0]][2] # use first prompt's memory
|
||||
_unload_omni(omni_bl)
|
||||
|
||||
# Save baseline outputs
|
||||
bl_dir = output_dir / "baseline"
|
||||
bl_dir.mkdir(parents=True, exist_ok=True)
|
||||
for i, prompt in enumerate(prompts):
|
||||
out = baseline_outputs[prompt][0]
|
||||
if is_video:
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
frames_list = list(out) if isinstance(out, np.ndarray) and out.ndim == 4 else out
|
||||
export_to_video(frames_list, str(bl_dir / f"prompt_{i}.mp4"), fps=args.fps)
|
||||
except ImportError:
|
||||
np.save(bl_dir / f"prompt_{i}.npy", out)
|
||||
else:
|
||||
out.save(bl_dir / f"prompt_{i}.png")
|
||||
|
||||
# --- Quantized runs ---
|
||||
all_results = [] # list of dicts
|
||||
|
||||
for config_label, quant_method in configs:
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"Running: {config_label}...")
|
||||
print("=" * 60)
|
||||
|
||||
qt_kwargs = _build_omni_kwargs(args, quantization=quant_method)
|
||||
omni_qt = Omni(**qt_kwargs)
|
||||
|
||||
qt_outputs = {}
|
||||
for prompt in prompts:
|
||||
print(f" Generating: {prompt[:60]}...")
|
||||
if is_video:
|
||||
out, t, mem = _generate_video(omni_qt, args, prompt, seed)
|
||||
else:
|
||||
out, t, mem = _generate_image(omni_qt, args, prompt, seed)
|
||||
qt_outputs[prompt] = (out, t, mem)
|
||||
|
||||
qt_avg_time = np.mean([v[1] for v in qt_outputs.values()])
|
||||
qt_mem = qt_outputs[prompts[0]][2]
|
||||
_unload_omni(omni_qt)
|
||||
|
||||
# Save quantized outputs
|
||||
qt_dir = output_dir / config_label.replace(" ", "_")
|
||||
qt_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Compute LPIPS per prompt
|
||||
per_prompt = []
|
||||
for i, prompt in enumerate(prompts):
|
||||
bl_out = baseline_outputs[prompt][0]
|
||||
qt_out = qt_outputs[prompt][0]
|
||||
if is_video:
|
||||
lpips_score = compute_lpips_video(bl_out, qt_out, net=args.lpips_net)
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
frames_list = list(qt_out) if isinstance(qt_out, np.ndarray) and qt_out.ndim == 4 else qt_out
|
||||
export_to_video(frames_list, str(qt_dir / f"prompt_{i}.mp4"), fps=args.fps)
|
||||
except ImportError:
|
||||
np.save(qt_dir / f"prompt_{i}.npy", qt_out)
|
||||
else:
|
||||
lpips_score = compute_lpips_images([bl_out], [qt_out], net=args.lpips_net)[0]
|
||||
qt_out.save(qt_dir / f"prompt_{i}.png")
|
||||
per_prompt.append({"prompt": prompt, "lpips": lpips_score})
|
||||
|
||||
mean_lpips = np.mean([p["lpips"] for p in per_prompt])
|
||||
speedup = bl_avg_time / qt_avg_time if qt_avg_time > 0 else float("inf")
|
||||
mem_reduction = (bl_mem - qt_mem) / bl_mem * 100
|
||||
|
||||
all_results.append(
|
||||
{
|
||||
"config": config_label,
|
||||
"avg_time": qt_avg_time,
|
||||
"speedup": speedup,
|
||||
"memory_gib": qt_mem,
|
||||
"mem_reduction_pct": mem_reduction,
|
||||
"mean_lpips": mean_lpips,
|
||||
"per_prompt": per_prompt,
|
||||
}
|
||||
)
|
||||
|
||||
# --- Print results ---
|
||||
print("\n\n")
|
||||
print("=" * 80)
|
||||
print("RESULTS")
|
||||
print("=" * 80)
|
||||
|
||||
# Summary table
|
||||
lines = []
|
||||
lines.append(f"## Quantization Quality Benchmark — {args.model.split('/')[-1]}")
|
||||
lines.append(
|
||||
f"Setup: {args.height}x{args.width}, {args.num_inference_steps} steps, "
|
||||
f"seed={args.seed}, LPIPS ({args.lpips_net})"
|
||||
)
|
||||
if is_video:
|
||||
lines.append(f"Video: {args.num_frames} frames")
|
||||
lines.append("")
|
||||
lines.append("### Summary")
|
||||
lines.append("")
|
||||
lines.append("| Config | Avg Time | Speedup | Memory (GiB) | Mem Reduction | Mean LPIPS |")
|
||||
lines.append("|--------|----------|---------|--------------|---------------|------------|")
|
||||
lines.append(f"| BF16 baseline | {bl_avg_time:.2f}s | 1.00x | {bl_mem:.2f} | — | (ref) |")
|
||||
for r in all_results:
|
||||
lines.append(
|
||||
f"| {r['config']} | {r['avg_time']:.2f}s | {r['speedup']:.2f}x "
|
||||
f"| {r['memory_gib']:.2f} | {r['mem_reduction_pct']:.0f}% "
|
||||
f"| {r['mean_lpips']:.4f} |"
|
||||
)
|
||||
lines.append("")
|
||||
lines.append("> LPIPS < 0.01 = imperceptible, > 0.1 = clearly noticeable.")
|
||||
lines.append("")
|
||||
|
||||
# Per-prompt table
|
||||
if len(prompts) > 1:
|
||||
lines.append("### Per-Prompt LPIPS")
|
||||
lines.append("")
|
||||
header = "| Prompt |"
|
||||
sep = "|--------|"
|
||||
for r in all_results:
|
||||
header += f" {r['config']} |"
|
||||
sep += "--------|"
|
||||
lines.append(header)
|
||||
lines.append(sep)
|
||||
for i, prompt in enumerate(prompts):
|
||||
short = prompt[:50] + "..." if len(prompt) > 50 else prompt
|
||||
row = f"| {short} |"
|
||||
for r in all_results:
|
||||
row += f" {r['per_prompt'][i]['lpips']:.4f} |"
|
||||
lines.append(row)
|
||||
lines.append("")
|
||||
|
||||
md = "\n".join(lines)
|
||||
print(md)
|
||||
|
||||
# Save markdown
|
||||
results_path = output_dir / "results.md"
|
||||
results_path.write_text(md, encoding="utf-8")
|
||||
print(f"\nResults saved to {results_path}")
|
||||
print(f"Baseline outputs in {bl_dir}")
|
||||
for r in all_results:
|
||||
qt_dir = output_dir / r["config"].replace(" ", "_")
|
||||
print(f"Quantized outputs in {qt_dir}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Benchmark quantization quality loss for diffusion models.",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
)
|
||||
parser.add_argument("--model", required=True, help="Model name or local path.")
|
||||
parser.add_argument(
|
||||
"--task",
|
||||
default="t2i",
|
||||
choices=["t2i", "t2v"],
|
||||
help="Task type: t2i (text-to-image) or t2v (text-to-video).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
nargs="+",
|
||||
required=True,
|
||||
help="One or more quantization methods to benchmark (e.g. fp8 int8 bitsandbytes).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompts",
|
||||
nargs="+",
|
||||
default=["a cup of coffee on the table"],
|
||||
help="One or more prompts to generate.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--height", type=int, default=1024)
|
||||
parser.add_argument("--width", type=int, default=1024)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50)
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of video frames (t2v only).")
|
||||
parser.add_argument("--fps", type=int, default=24, help="Video FPS for saving (t2v only).")
|
||||
parser.add_argument("--guidance-scale", type=float, default=4.0, help="CFG scale (used for video).")
|
||||
parser.add_argument("--output-dir", type=str, default="./quant_bench_output", help="Directory to save outputs.")
|
||||
parser.add_argument(
|
||||
"--lpips-net",
|
||||
type=str,
|
||||
default="alex",
|
||||
choices=["alex", "vgg", "squeeze"],
|
||||
help="LPIPS backbone network.",
|
||||
)
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1)
|
||||
parser.add_argument("--ring-degree", type=int, default=1)
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1)
|
||||
parser.add_argument("--enforce-eager", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_benchmark(args)
|
||||
@@ -56,9 +56,9 @@ What it does:
|
||||
- Runs `examples/offline_inference/qwen3_omni/end2end.py` with `--log-stats`.
|
||||
- Uses `benchmarks/build_dataset/top100.txt` and writes to:
|
||||
- Logs: `benchmarks/qwen3-omni/vllm_omni/logs/`
|
||||
- `omni_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
|
||||
- `omni_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
|
||||
- `omni_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
|
||||
- `omni_llm_pipeline_text.orchestrator.stats.jsonl` — per-stage latency stats.
|
||||
- `omni_llm_pipeline_text.overall.stats.jsonl` — end-to-end latency/TPS.
|
||||
- `omni_llm_pipeline_text.stage{0,1,2}.log` — per-stage detailed logs/errors.
|
||||
- Outputs: `benchmarks/qwen3-omni/vllm_omni/outputs/` — ~100 text and `.wav` files.
|
||||
|
||||
Key checks:
|
||||
|
||||
@@ -26,12 +26,12 @@ else
|
||||
--log-stats \
|
||||
--log-dir $log_dir
|
||||
echo "Logs and outputs are saved in ${log_dir} and ${outputs_dir} respectively:"
|
||||
echo " - omni_pipeline_text run dir/base name"
|
||||
echo " - omni_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
|
||||
echo " - omni_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
|
||||
echo " - omni_pipeline_text.stage0.log per-stage detailed logs"
|
||||
echo " - omni_pipeline_text.stage1.log"
|
||||
echo " - omni_pipeline_text.stage2.log"
|
||||
echo " - omni_llm_pipeline_text run dir/base name"
|
||||
echo " - omni_llm_pipeline_text.orchestrator.stats.jsonl orchestrator-stage latency stats"
|
||||
echo " - omni_llm_pipeline_text.overall.stats.jsonl overall latency/TPS stats"
|
||||
echo " - omni_llm_pipeline_text.stage0.log per-stage detailed logs"
|
||||
echo " - omni_llm_pipeline_text.stage1.log"
|
||||
echo " - omni_llm_pipeline_text.stage2.log"
|
||||
echo "Key checks: overall.stats.jsonl for end-to-end latency/TPS; orchestrator.stats.jsonl for stable per-stage latency; stage*.log for errors or long tails."
|
||||
echo " - outputs/ Generated txt and wav files, there should be 100 text and wav files generated respectively"
|
||||
fi
|
||||
@@ -1,103 +0,0 @@
|
||||
# Qwen3-TTS Benchmark
|
||||
|
||||
Benchmarks for Qwen3-TTS text-to-speech models, comparing vLLM-Omni streaming serving against HuggingFace Transformers offline inference.
|
||||
|
||||
## Prerequisites
|
||||
|
||||
```bash
|
||||
pip install matplotlib aiohttp soundfile numpy tqdm
|
||||
pip install qwen_tts # for HF baseline
|
||||
```
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the full benchmark (vllm-omni + HF baseline) with a single command:
|
||||
|
||||
```bash
|
||||
cd benchmarks/qwen3-tts
|
||||
bash run_benchmark.sh
|
||||
```
|
||||
|
||||
Results (JSON + PNG plots) are saved to `results/`.
|
||||
|
||||
### Common options
|
||||
|
||||
```bash
|
||||
# Only vllm-omni (skip HF baseline)
|
||||
bash run_benchmark.sh --async-only
|
||||
|
||||
# Only HF baseline
|
||||
bash run_benchmark.sh --hf-only
|
||||
|
||||
# Use a different model (e.g. 1.7B)
|
||||
MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
|
||||
# Use 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
|
||||
|
||||
# Custom GPU, prompt count, concurrency levels
|
||||
GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
```
|
||||
|
||||
## Manual Steps
|
||||
|
||||
### 1) Start the vLLM-Omni server
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0 python -m vllm_omni.entrypoints.cli.main serve \
|
||||
"Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--omni --host 127.0.0.1 --port 8000 \
|
||||
--stage-configs-path benchmarks/qwen3-tts/vllm_omni/configs/qwen3_tts_bs1.yaml \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
### 2) Run online serving benchmark
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/vllm_omni/bench_tts_serve.py \
|
||||
--port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--config-name "async_chunk" \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 3) Run HuggingFace baseline
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/transformers/bench_tts_hf.py \
|
||||
--model "Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice" \
|
||||
--num-prompts 50 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
```
|
||||
|
||||
### 4) Generate comparison plots
|
||||
|
||||
```bash
|
||||
python benchmarks/qwen3-tts/plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_hf_transformers_*.json \
|
||||
--labels "vllm-omni" "hf_transformers" \
|
||||
--output results/comparison.png
|
||||
```
|
||||
|
||||
## Stage Configs
|
||||
|
||||
| Config | max_num_seqs | Description |
|
||||
|--------|:------------:|-------------|
|
||||
| `vllm_omni/configs/qwen3_tts_bs1.yaml` | 1 | Single-request processing (lowest latency) |
|
||||
| `vllm_omni/configs/qwen3_tts_bs16.yaml` | 16 | High-throughput concurrent processing |
|
||||
|
||||
All configs use a 2-stage pipeline (Talker -> Code2Wav) with `async_chunk` streaming enabled. The `SharedMemoryConnector` streams codec frames (25-frame chunks with 25-frame context overlap) between stages.
|
||||
|
||||
The model is specified via the CLI `--model` flag (or `MODEL` env var), so the same configs work for both the 0.6B and 1.7B model variants.
|
||||
|
||||
## Metrics
|
||||
|
||||
- **TTFP (Time to First Audio Packet)**: Time from request to first audio chunk (streaming latency)
|
||||
- **E2E (End-to-End Latency)**: Total time from request to complete audio response
|
||||
- **RTF (Real-Time Factor)**: E2E latency / audio duration. RTF < 1.0 means faster-than-real-time synthesis
|
||||
- **Throughput**: Total audio seconds generated per wall-clock second
|
||||
@@ -1,254 +0,0 @@
|
||||
"""Plot Qwen3-TTS benchmark results.
|
||||
|
||||
Generates comparison bar charts similar to the async_chunk design doc:
|
||||
- TTFP (Time-to-First-Packet) across concurrency levels
|
||||
- E2E latency across concurrency levels
|
||||
- RTF (Real-Time Factor) across concurrency levels
|
||||
|
||||
Usage:
|
||||
# Compare two configs (async_chunk vs no_async_chunk):
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json results/bench_no_async_chunk_*.json \
|
||||
--labels "async_chunk" "no_async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
|
||||
# Single config:
|
||||
python plot_results.py \
|
||||
--results results/bench_async_chunk_*.json \
|
||||
--labels "async_chunk" \
|
||||
--output results/qwen3_tts_benchmark.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
|
||||
def load_results(result_files: list[str]) -> list[list[dict]]:
|
||||
"""Load benchmark results from JSON files."""
|
||||
all_results = []
|
||||
for f in result_files:
|
||||
with open(f) as fh:
|
||||
data = json.load(fh)
|
||||
all_results.append(data)
|
||||
return all_results
|
||||
|
||||
|
||||
def plot_comparison(
|
||||
all_results: list[list[dict]],
|
||||
labels: list[str],
|
||||
output_path: str,
|
||||
title_prefix: str = "Qwen3-TTS",
|
||||
):
|
||||
"""Generate comparison bar charts."""
|
||||
n_configs = len(all_results)
|
||||
|
||||
# Collect concurrency levels present in ALL configs (skip missing data)
|
||||
all_concurrencies = [set(r["concurrency"] for r in results) for results in all_results]
|
||||
concurrencies = sorted(set.union(*all_concurrencies))
|
||||
|
||||
# Build data arrays, using None for missing concurrency levels
|
||||
ttfp_data = {label: [] for label in labels}
|
||||
e2e_data = {label: [] for label in labels}
|
||||
rtf_data = {label: [] for label in labels}
|
||||
throughput_data = {label: [] for label in labels}
|
||||
|
||||
for results, label in zip(all_results, labels):
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
for c in concurrencies:
|
||||
r = conc_map.get(c)
|
||||
ttfp_data[label].append(r["mean_ttfp_ms"] if r else None)
|
||||
e2e_data[label].append(r["mean_e2e_ms"] if r else None)
|
||||
rtf_data[label].append(r["mean_rtf"] if r else None)
|
||||
throughput_data[label].append(r["audio_throughput"] if r else None)
|
||||
|
||||
fig, axes = plt.subplots(2, 2, figsize=(14, 10))
|
||||
fig.suptitle(f"{title_prefix} Performance Benchmark", fontsize=16, fontweight="bold")
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.35 if n_configs == 2 else 0.5
|
||||
if n_configs > 1:
|
||||
offsets = np.linspace(-width / 2 * (n_configs - 1), width / 2 * (n_configs - 1), n_configs)
|
||||
else:
|
||||
offsets = [0]
|
||||
|
||||
colors = ["#2196F3", "#FF5722", "#4CAF50", "#FFC107"]
|
||||
|
||||
def plot_metric(ax, data_dict, ylabel, title, fmt=".1f"):
|
||||
bars = []
|
||||
for i, (label, values) in enumerate(data_dict.items()):
|
||||
# Replace None with 0 for plotting, but track which are missing
|
||||
plot_values = [v if v is not None else 0 for v in values]
|
||||
color = colors[i % len(colors)]
|
||||
bar = ax.bar(x + offsets[i], plot_values, width, label=label, color=color, alpha=0.85)
|
||||
bars.append(bar)
|
||||
# Add value labels on bars (skip None/missing data)
|
||||
max_val = max((v for v in values if v is not None), default=1)
|
||||
for rect, val in zip(bar, values):
|
||||
if val is not None and val > 0:
|
||||
ax.text(
|
||||
rect.get_x() + rect.get_width() / 2,
|
||||
rect.get_height() + max_val * 0.02,
|
||||
f"{val:{fmt}}",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=9,
|
||||
fontweight="bold",
|
||||
)
|
||||
ax.set_xlabel("Concurrency", fontsize=12)
|
||||
ax.set_ylabel(ylabel, fontsize=12)
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.legend(fontsize=10)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plot_metric(axes[0, 0], ttfp_data, "TTFP (ms)", "Time to First Audio Packet (TTFP)")
|
||||
plot_metric(axes[0, 1], e2e_data, "E2E Latency (ms)", "End-to-End Latency (E2E)")
|
||||
plot_metric(axes[1, 0], rtf_data, "RTF", "Real-Time Factor (RTF)", fmt=".3f")
|
||||
plot_metric(axes[1, 1], throughput_data, "Audio-sec / Wall-sec", "Audio Throughput", fmt=".2f")
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_single_summary(results: list[dict], label: str, output_path: str):
|
||||
"""Generate a single-config summary with percentile breakdown."""
|
||||
concurrencies = [r["concurrency"] for r in results]
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(16, 5))
|
||||
fig.suptitle(f"Qwen3-TTS Benchmark - {label}", fontsize=15, fontweight="bold")
|
||||
|
||||
# TTFP breakdown
|
||||
ax = axes[0]
|
||||
means = [r["mean_ttfp_ms"] for r in results]
|
||||
medians = [r["median_ttfp_ms"] for r in results]
|
||||
p90s = [r["p90_ttfp_ms"] for r in results]
|
||||
p99s = [r["p99_ttfp_ms"] for r in results]
|
||||
x = np.arange(len(concurrencies))
|
||||
w = 0.2
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("TTFP (ms)")
|
||||
ax.set_title("Time to First Audio Packet")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# E2E breakdown
|
||||
ax = axes[1]
|
||||
means = [r["mean_e2e_ms"] for r in results]
|
||||
medians = [r["median_e2e_ms"] for r in results]
|
||||
p90s = [r["p90_e2e_ms"] for r in results]
|
||||
p99s = [r["p99_e2e_ms"] for r in results]
|
||||
ax.bar(x - 1.5 * w, means, w, label="mean", color="#2196F3")
|
||||
ax.bar(x - 0.5 * w, medians, w, label="median", color="#4CAF50")
|
||||
ax.bar(x + 0.5 * w, p90s, w, label="p90", color="#FF9800")
|
||||
ax.bar(x + 1.5 * w, p99s, w, label="p99", color="#F44336")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("E2E Latency (ms)")
|
||||
ax.set_title("End-to-End Latency")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
# RTF
|
||||
ax = axes[2]
|
||||
means = [r["mean_rtf"] for r in results]
|
||||
medians = [r["median_rtf"] for r in results]
|
||||
ax.bar(x - 0.15, means, 0.3, label="mean", color="#2196F3")
|
||||
ax.bar(x + 0.15, medians, 0.3, label="median", color="#4CAF50")
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_xlabel("Concurrency")
|
||||
ax.set_ylabel("RTF")
|
||||
ax.set_title("Real-Time Factor")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3)
|
||||
|
||||
plt.tight_layout()
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_comparison_table(all_results: list[list[dict]], labels: list[str]):
|
||||
"""Print a markdown-formatted comparison table."""
|
||||
concurrencies = sorted(set(r["concurrency"] for r in all_results[0]))
|
||||
|
||||
print("\n## Benchmark Results\n")
|
||||
header = "| Metric | Concurrency |"
|
||||
sep = "| --- | --- |"
|
||||
for label in labels:
|
||||
header += f" {label} |"
|
||||
sep += " --- |"
|
||||
print(header)
|
||||
print(sep)
|
||||
|
||||
for metric, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput (audio-s/s)", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
row = f"| {metric} | {c} |"
|
||||
for results in all_results:
|
||||
conc_map = {r["concurrency"]: r for r in results}
|
||||
val = conc_map.get(c, {}).get(key, 0)
|
||||
row += f" {val:{fmt}} |"
|
||||
print(row)
|
||||
|
||||
# Improvement calculation (only if 2 configs)
|
||||
if len(all_results) == 2:
|
||||
print(f"\n## Improvement ({labels[0]} vs {labels[1]})\n")
|
||||
print("| Metric | Concurrency | Improvement |")
|
||||
print("| --- | --- | --- |")
|
||||
for metric, key in [("TTFP", "mean_ttfp_ms"), ("E2E", "mean_e2e_ms"), ("RTF", "mean_rtf")]:
|
||||
for c in concurrencies:
|
||||
m0 = {r["concurrency"]: r for r in all_results[0]}
|
||||
m1 = {r["concurrency"]: r for r in all_results[1]}
|
||||
v0 = m0.get(c, {}).get(key, 0)
|
||||
v1 = m1.get(c, {}).get(key, 0)
|
||||
if v1 > 0:
|
||||
pct = (v1 - v0) / v1 * 100
|
||||
print(f"| {metric} | {c} | {pct:+.1f}% |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot Qwen3-TTS benchmark results")
|
||||
parser.add_argument(
|
||||
"--results", type=str, nargs="+", required=True, help="Path(s) to result JSON files (one per config)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--labels", type=str, nargs="+", required=True, help="Labels for each config (must match --results count)"
|
||||
)
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_benchmark.png", help="Output image path")
|
||||
parser.add_argument("--title", type=str, default="Qwen3-TTS", help="Title prefix for the plot")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
assert len(args.results) == len(args.labels), "--results and --labels must have the same count"
|
||||
|
||||
all_results = load_results(args.results)
|
||||
print_comparison_table(all_results, args.labels)
|
||||
|
||||
Path(args.output).parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
if len(all_results) == 1:
|
||||
plot_single_summary(all_results[0], args.labels[0], args.output)
|
||||
else:
|
||||
plot_comparison(all_results, args.labels, args.output, title_prefix=args.title)
|
||||
@@ -1,3 +0,0 @@
|
||||
# Benchmark results are machine-specific - do not commit
|
||||
*
|
||||
!.gitignore
|
||||
@@ -1,280 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS Benchmark Runner
|
||||
#
|
||||
# Compares vllm-omni streaming serving vs HuggingFace transformers offline inference.
|
||||
# Produces JSON results and comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# # Full comparison (vllm-omni + HF):
|
||||
# bash run_benchmark.sh
|
||||
#
|
||||
# # Only vllm-omni async_chunk config:
|
||||
# bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Only HuggingFace baseline:
|
||||
# bash run_benchmark.sh --hf-only
|
||||
#
|
||||
# # vllm-omni only (skip HF):
|
||||
# bash run_benchmark.sh --skip-hf
|
||||
#
|
||||
# # Custom settings:
|
||||
# GPU_DEVICE=1 NUM_PROMPTS=20 CONCURRENCY="1 4" bash run_benchmark.sh
|
||||
#
|
||||
# # Use 1.7B model:
|
||||
# MODEL=Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice bash run_benchmark.sh --async-only
|
||||
#
|
||||
# # Use 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
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index to use (default: 0)
|
||||
# NUM_PROMPTS - Number of prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 4 10")
|
||||
# MODEL - Model name (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# PORT - Server port (default: 8000)
|
||||
# GPU_MEM_TALKER - gpu_memory_utilization for talker stage (default: 0.3)
|
||||
# GPU_MEM_CODE2WAV - gpu_memory_utilization for code2wav stage (default: 0.2)
|
||||
# STAGE_CONFIG - Path to stage config YAML (default: configs/qwen3_tts_bs1.yaml)
|
||||
# TASK_TYPE - Task type: CustomVoice, VoiceDesign, Base (default: CustomVoice)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../.." && pwd)"
|
||||
|
||||
# Defaults
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 4 10}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
PORT="${PORT:-8000}"
|
||||
GPU_MEM_TALKER="${GPU_MEM_TALKER:-0.3}"
|
||||
GPU_MEM_CODE2WAV="${GPU_MEM_CODE2WAV:-0.2}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
STAGE_CONFIG="${STAGE_CONFIG:-vllm_omni/configs/qwen3_tts_bs1.yaml}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
TASK_TYPE="${TASK_TYPE:-CustomVoice}"
|
||||
|
||||
# Parse args
|
||||
RUN_ASYNC=true
|
||||
RUN_HF=true
|
||||
for arg in "$@"; do
|
||||
case "$arg" in
|
||||
--async-only) RUN_HF=false ;;
|
||||
--hf-only) RUN_ASYNC=false ;;
|
||||
--skip-hf) RUN_HF=false ;;
|
||||
esac
|
||||
done
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port: ${PORT}"
|
||||
echo " Stage config: ${STAGE_CONFIG}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo " Task type: ${TASK_TYPE}"
|
||||
echo "============================================================"
|
||||
|
||||
# Prepare stage config with correct GPU device and memory settings
|
||||
prepare_config() {
|
||||
local config_template="$1"
|
||||
local config_name="$2"
|
||||
local output_path="${RESULT_DIR}/${config_name}_stage_config.yaml"
|
||||
|
||||
# Use sed to patch GPU device and memory utilization
|
||||
sed \
|
||||
-e "s/devices: \"0\"/devices: \"${GPU_DEVICE}\"/g" \
|
||||
-e "s/gpu_memory_utilization: 0.3/gpu_memory_utilization: ${GPU_MEM_TALKER}/g" \
|
||||
-e "s/gpu_memory_utilization: 0.2/gpu_memory_utilization: ${GPU_MEM_CODE2WAV}/g" \
|
||||
"${config_template}" > "${output_path}"
|
||||
|
||||
echo "${output_path}"
|
||||
}
|
||||
|
||||
# Start server and wait for it to be ready
|
||||
start_server() {
|
||||
local stage_config="$1"
|
||||
local config_name="$2"
|
||||
local log_file="${RESULT_DIR}/server_${config_name}_${TIMESTAMP}.log"
|
||||
|
||||
echo ""
|
||||
echo "Starting server with config: ${config_name}"
|
||||
echo " Stage config: ${stage_config}"
|
||||
echo " Log file: ${log_file}"
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
CUDA_VISIBLE_DEVICES="${GPU_DEVICE}" \
|
||||
python -m vllm_omni.entrypoints.cli.main serve "${MODEL}" \
|
||||
--omni \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--stage-configs-path "${stage_config}" \
|
||||
--stage-init-timeout 120 \
|
||||
--trust-remote-code \
|
||||
--disable-log-stats \
|
||||
> "${log_file}" 2>&1 &
|
||||
|
||||
SERVER_PID=$!
|
||||
echo " Server PID: ${SERVER_PID}"
|
||||
|
||||
# Wait for server to be ready
|
||||
echo " Waiting for server to be ready..."
|
||||
local max_wait=300
|
||||
local waited=0
|
||||
while [ ${waited} -lt ${max_wait} ]; do
|
||||
if curl -sf "http://127.0.0.1:${PORT}/v1/models" > /dev/null 2>&1; then
|
||||
echo " Server is ready! (waited ${waited}s)"
|
||||
return 0
|
||||
fi
|
||||
# Check if process is still alive
|
||||
if ! kill -0 ${SERVER_PID} 2>/dev/null; then
|
||||
echo " ERROR: Server process died. Check log: ${log_file}"
|
||||
tail -20 "${log_file}"
|
||||
return 1
|
||||
fi
|
||||
sleep 2
|
||||
waited=$((waited + 2))
|
||||
done
|
||||
|
||||
echo " ERROR: Server did not start within ${max_wait}s. Check log: ${log_file}"
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
return 1
|
||||
}
|
||||
|
||||
# Stop the server
|
||||
stop_server() {
|
||||
if [ -n "${SERVER_PID:-}" ]; then
|
||||
echo " Stopping server (PID: ${SERVER_PID})..."
|
||||
kill ${SERVER_PID} 2>/dev/null || true
|
||||
wait ${SERVER_PID} 2>/dev/null || true
|
||||
# Kill any remaining child processes on the port
|
||||
local pids
|
||||
pids=$(lsof -ti:${PORT} 2>/dev/null || true)
|
||||
if [ -n "${pids}" ]; then
|
||||
echo " Cleaning up remaining processes on port ${PORT}..."
|
||||
echo "${pids}" | xargs kill -9 2>/dev/null || true
|
||||
fi
|
||||
echo " Server stopped."
|
||||
SERVER_PID=""
|
||||
fi
|
||||
}
|
||||
|
||||
# Cleanup on exit
|
||||
trap 'stop_server' EXIT
|
||||
|
||||
# Run benchmark for a given config
|
||||
run_bench() {
|
||||
local config_name="$1"
|
||||
local config_template="$2"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: ${config_name}"
|
||||
echo "============================================================"
|
||||
|
||||
local stage_config
|
||||
stage_config=$(prepare_config "${config_template}" "${config_name}")
|
||||
|
||||
start_server "${stage_config}" "${config_name}"
|
||||
|
||||
# Convert concurrency string to args
|
||||
local conc_args=""
|
||||
for c in ${CONCURRENCY}; do
|
||||
conc_args="${conc_args} ${c}"
|
||||
done
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/vllm_omni/bench_tts_serve.py" \
|
||||
--host 127.0.0.1 \
|
||||
--port "${PORT}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${conc_args} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--config-name "${config_name}" \
|
||||
--result-dir "${RESULT_DIR}" \
|
||||
--task-type "${TASK_TYPE}"
|
||||
|
||||
stop_server
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
}
|
||||
|
||||
# Run vllm-omni benchmark
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
run_bench "async_chunk" "${SCRIPT_DIR}/${STAGE_CONFIG}"
|
||||
fi
|
||||
|
||||
# Run HuggingFace baseline benchmark
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmarking: HuggingFace transformers (offline)"
|
||||
echo "============================================================"
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
python "${SCRIPT_DIR}/transformers/bench_tts_hf.py" \
|
||||
--model "${MODEL}" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--gpu-device "${GPU_DEVICE}" \
|
||||
--config-name "hf_transformers" \
|
||||
--result-dir "${RESULT_DIR}" \
|
||||
--task-type "${TASK_TYPE}"
|
||||
|
||||
# Allow GPU memory to settle
|
||||
sleep 5
|
||||
fi
|
||||
|
||||
# Plot results
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Generating plots..."
|
||||
echo "============================================================"
|
||||
|
||||
RESULT_FILES=""
|
||||
LABELS=""
|
||||
|
||||
if [ "${RUN_ASYNC}" = true ]; then
|
||||
ASYNC_FILE=$(ls -t "${RESULT_DIR}"/bench_async_chunk_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${ASYNC_FILE}" ]; then
|
||||
RESULT_FILES="${ASYNC_FILE}"
|
||||
LABELS="async_chunk"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ "${RUN_HF}" = true ]; then
|
||||
HF_FILE=$(ls -t "${RESULT_DIR}"/bench_hf_transformers_*.json 2>/dev/null | head -1)
|
||||
if [ -n "${HF_FILE}" ]; then
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
RESULT_FILES="${RESULT_FILES} ${HF_FILE}"
|
||||
LABELS="${LABELS} hf_transformers"
|
||||
else
|
||||
RESULT_FILES="${HF_FILE}"
|
||||
LABELS="hf_transformers"
|
||||
fi
|
||||
fi
|
||||
fi
|
||||
|
||||
if [ -n "${RESULT_FILES}" ]; then
|
||||
python "${SCRIPT_DIR}/plot_results.py" \
|
||||
--results ${RESULT_FILES} \
|
||||
--labels ${LABELS} \
|
||||
--output "${RESULT_DIR}/qwen3_tts_benchmark_${TIMESTAMP}.png"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
@@ -1,301 +0,0 @@
|
||||
"""Benchmark Qwen3-TTS using HuggingFace transformers (qwen_tts library).
|
||||
|
||||
Measures E2E latency, RTF, and audio duration for offline (non-serving) inference.
|
||||
Results are saved in the same JSON format as bench_tts_serve.py for unified plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_hf.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice \
|
||||
--num-prompts 50 \
|
||||
--num-warmups 3 \
|
||||
--gpu-device 0 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
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:
|
||||
config_name: str = ""
|
||||
concurrency: int = 1 # always 1 for offline
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats - not applicable for HF offline, set to E2E for compatibility
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def 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
|
||||
|
||||
device = f"cuda:{args.gpu_device}"
|
||||
print(f"Loading model: {args.model} on {device}")
|
||||
model = Qwen3TTSModel.from_pretrained(
|
||||
args.model,
|
||||
device_map=device,
|
||||
dtype=torch.bfloat16,
|
||||
)
|
||||
print("Model loaded.")
|
||||
|
||||
# Build prompt list
|
||||
prompts = [PROMPTS[i % len(PROMPTS)] for i in range(args.num_prompts)]
|
||||
|
||||
# Warmup
|
||||
if args.num_warmups > 0:
|
||||
print(f"Warming up with {args.num_warmups} requests...")
|
||||
for i in range(args.num_warmups):
|
||||
p = PROMPTS[i % len(PROMPTS)]
|
||||
wavs, sr = generate_audio(model, p, args)
|
||||
# Sync GPU
|
||||
torch.cuda.synchronize(device)
|
||||
print("Warmup done.")
|
||||
|
||||
# Benchmark
|
||||
print(f"Running {args.num_prompts} requests sequentially...")
|
||||
e2e_times = []
|
||||
rtfs = []
|
||||
audio_durations = []
|
||||
per_request = []
|
||||
failed = 0
|
||||
|
||||
audio_dir = None
|
||||
if args.save_audio:
|
||||
audio_dir = Path(args.result_dir) / "audio_hf"
|
||||
audio_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
total_start = time.perf_counter()
|
||||
|
||||
for i, prompt in enumerate(prompts):
|
||||
try:
|
||||
torch.cuda.synchronize(device)
|
||||
st = time.perf_counter()
|
||||
|
||||
wavs, sr = generate_audio(model, prompt, args)
|
||||
|
||||
torch.cuda.synchronize(device)
|
||||
elapsed = time.perf_counter() - st
|
||||
|
||||
# Compute audio duration
|
||||
audio_samples = wavs[0]
|
||||
if isinstance(audio_samples, torch.Tensor):
|
||||
audio_samples = audio_samples.cpu().numpy()
|
||||
audio_dur = len(audio_samples) / sr
|
||||
|
||||
rtf = elapsed / audio_dur if audio_dur > 0 else 0.0
|
||||
|
||||
e2e_times.append(elapsed)
|
||||
rtfs.append(rtf)
|
||||
audio_durations.append(audio_dur)
|
||||
per_request.append(
|
||||
{
|
||||
"e2e_ms": elapsed * 1000,
|
||||
"ttfp_ms": elapsed * 1000, # no streaming, TTFP = E2E
|
||||
"rtf": rtf,
|
||||
"audio_duration_s": audio_dur,
|
||||
"prompt": prompt,
|
||||
}
|
||||
)
|
||||
|
||||
if audio_dir:
|
||||
sf.write(str(audio_dir / f"output_{i:04d}.wav"), audio_samples, sr)
|
||||
|
||||
if (i + 1) % 10 == 0 or i == 0:
|
||||
print(
|
||||
f" [{i + 1}/{args.num_prompts}] e2e={elapsed * 1000:.0f}ms rtf={rtf:.3f} audio={audio_dur:.2f}s"
|
||||
)
|
||||
|
||||
except Exception as e:
|
||||
print(f" [{i + 1}/{args.num_prompts}] FAILED: {e}")
|
||||
failed += 1
|
||||
|
||||
total_duration = time.perf_counter() - total_start
|
||||
completed = len(e2e_times)
|
||||
|
||||
# Compute stats
|
||||
result = BenchmarkResult(
|
||||
config_name=args.config_name,
|
||||
concurrency=1,
|
||||
num_prompts=args.num_prompts,
|
||||
completed=completed,
|
||||
failed=failed,
|
||||
duration_s=total_duration,
|
||||
)
|
||||
|
||||
if e2e_times:
|
||||
e2e_ms = [t * 1000 for t in e2e_times]
|
||||
|
||||
result.mean_e2e_ms = float(np.mean(e2e_ms))
|
||||
result.median_e2e_ms = float(np.median(e2e_ms))
|
||||
result.std_e2e_ms = float(np.std(e2e_ms))
|
||||
result.p90_e2e_ms = float(np.percentile(e2e_ms, 90))
|
||||
result.p95_e2e_ms = float(np.percentile(e2e_ms, 95))
|
||||
result.p99_e2e_ms = float(np.percentile(e2e_ms, 99))
|
||||
|
||||
# For HF offline, TTFP = E2E (no streaming)
|
||||
result.mean_ttfp_ms = result.mean_e2e_ms
|
||||
result.median_ttfp_ms = result.median_e2e_ms
|
||||
result.std_ttfp_ms = result.std_e2e_ms
|
||||
result.p90_ttfp_ms = result.p90_e2e_ms
|
||||
result.p95_ttfp_ms = result.p95_e2e_ms
|
||||
result.p99_ttfp_ms = result.p99_e2e_ms
|
||||
|
||||
result.mean_rtf = float(np.mean(rtfs))
|
||||
result.median_rtf = float(np.median(rtfs))
|
||||
result.std_rtf = float(np.std(rtfs))
|
||||
result.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
result.mean_audio_duration_s = float(np.mean(audio_durations))
|
||||
result.total_audio_duration_s = float(np.sum(audio_durations))
|
||||
result.audio_throughput = result.total_audio_duration_s / total_duration
|
||||
result.request_throughput = completed / total_duration
|
||||
result.per_request = per_request
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{1:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{total_duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{result.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{result.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{result.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{result.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{result.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{result.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{result.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{result.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{result.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{result.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{result.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{result.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
# Save results (as a list with single concurrency=1 entry, matching serve format)
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump([asdict(result)], f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return result
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS HuggingFace Benchmark")
|
||||
parser.add_argument(
|
||||
"--model", type=str, default="Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice", help="HuggingFace model name or path"
|
||||
)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--gpu-device", type=int, default=0)
|
||||
parser.add_argument("--voice", type=str, default="Vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--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)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
parser.add_argument("--save-audio", action="store_true", help="Save generated audio files")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_benchmark(args)
|
||||
@@ -1,301 +0,0 @@
|
||||
"""Benchmark comparing async_chunk on vs off for Qwen3-TTS.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF across
|
||||
concurrency levels for both async_chunk modes. Saves results as JSON.
|
||||
|
||||
Usage:
|
||||
# Run against a server already serving with a given config:
|
||||
python bench_async_chunk.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--config-name async_chunk_on \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
|
||||
|
||||
@dataclass
|
||||
class RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0
|
||||
e2e: float = 0.0
|
||||
audio_bytes: int = 0
|
||||
audio_duration: float = 0.0
|
||||
rtf: float = 0.0
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0
|
||||
request_throughput: float = 0.0
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
return num_bytes / sample_width / sample_rate
|
||||
|
||||
|
||||
async def send_tts_request(
|
||||
session: aiohttp.ClientSession,
|
||||
api_url: str,
|
||||
prompt: str,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
pbar: tqdm | None = None,
|
||||
) -> RequestResult:
|
||||
payload = {
|
||||
"input": prompt,
|
||||
"voice": voice,
|
||||
"language": language,
|
||||
"stream": stream,
|
||||
"response_format": "pcm",
|
||||
}
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
stream: bool = True,
|
||||
) -> BenchmarkResult:
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60)
|
||||
session = aiohttp.ClientSession(connector=connector, timeout=aiohttp.ClientTimeout(total=600))
|
||||
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = [
|
||||
send_tts_request(session, api_url, PROMPTS[i % len(PROMPTS)], voice, language, stream)
|
||||
for i in range(num_warmups)
|
||||
]
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, voice, language, stream, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful]
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f" Concurrency: {max_concurrency} | Completed: {bench.completed} | Failed: {bench.failed}")
|
||||
print(f" Duration: {duration:.2f}s | Throughput: {bench.request_throughput:.2f} req/s")
|
||||
print(
|
||||
f" TTFP (ms): mean={bench.mean_ttfp_ms:.1f} median={bench.median_ttfp_ms:.1f}"
|
||||
f" p90={bench.p90_ttfp_ms:.1f} p99={bench.p99_ttfp_ms:.1f}"
|
||||
)
|
||||
print(
|
||||
f" E2E (ms): mean={bench.mean_e2e_ms:.1f} median={bench.median_e2e_ms:.1f}"
|
||||
f" p90={bench.p90_e2e_ms:.1f} p99={bench.p99_e2e_ms:.1f}"
|
||||
)
|
||||
print(f" RTF: mean={bench.mean_rtf:.3f} median={bench.median_rtf:.3f}")
|
||||
print(f" Throughput: {bench.audio_throughput:.2f} audio-sec/wall-sec")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
stream=args.stream,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS async_chunk benchmark client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50)
|
||||
parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1, 10])
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--voice", type=str, default="vivian")
|
||||
parser.add_argument("--language", type=str, default="English")
|
||||
parser.add_argument("--stream", action="store_true", default=True)
|
||||
parser.add_argument("--no-stream", dest="stream", action="store_false")
|
||||
parser.add_argument("--config-name", type=str, default="async_chunk_on")
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -1,371 +0,0 @@
|
||||
"""Benchmark client for Qwen3-TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Measures TTFP (Time-to-First-Packet), E2E latency, and RTF (Real-Time Factor)
|
||||
across configurable concurrency levels. Saves results as JSON for plotting.
|
||||
|
||||
Usage:
|
||||
python bench_tts_serve.py \
|
||||
--host 127.0.0.1 --port 8000 \
|
||||
--num-prompts 50 \
|
||||
--max-concurrency 1 4 10 \
|
||||
--result-dir results/
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
from dataclasses import asdict, dataclass, field
|
||||
from datetime import datetime
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
from tqdm.asyncio import tqdm
|
||||
|
||||
PROMPTS = [
|
||||
"Hello, welcome to the voice synthesis benchmark test.",
|
||||
"She said she would be here by noon, but nobody showed up.",
|
||||
"The quick brown fox jumps over the lazy dog near the riverbank.",
|
||||
"I can't believe how beautiful the sunset looks from up here on the mountain.",
|
||||
"Please remember to bring your identification documents to the appointment tomorrow morning.",
|
||||
"Have you ever wondered what it would be like to travel through time and visit ancient civilizations?",
|
||||
"The restaurant on the corner serves the best pasta I have ever tasted in my entire life.",
|
||||
"After the meeting, we should discuss the quarterly results and plan for the next phase.",
|
||||
"Learning a new language takes patience, practice, and a genuine curiosity about other cultures.",
|
||||
"The train leaves at half past seven, so we need to arrive at the station before then.",
|
||||
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
|
||||
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
|
||||
]
|
||||
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 RequestResult:
|
||||
success: bool = False
|
||||
ttfp: float = 0.0 # Time to first audio packet (seconds)
|
||||
e2e: float = 0.0 # End-to-end latency (seconds)
|
||||
audio_bytes: int = 0 # Total audio bytes received
|
||||
audio_duration: float = 0.0 # Audio duration in seconds (estimated from PCM)
|
||||
rtf: float = 0.0 # Real-time factor = e2e / audio_duration
|
||||
prompt: str = ""
|
||||
error: str = ""
|
||||
|
||||
|
||||
@dataclass
|
||||
class BenchmarkResult:
|
||||
config_name: str = ""
|
||||
concurrency: int = 0
|
||||
num_prompts: int = 0
|
||||
completed: int = 0
|
||||
failed: int = 0
|
||||
duration_s: float = 0.0
|
||||
# TTFP stats (ms)
|
||||
mean_ttfp_ms: float = 0.0
|
||||
median_ttfp_ms: float = 0.0
|
||||
std_ttfp_ms: float = 0.0
|
||||
p90_ttfp_ms: float = 0.0
|
||||
p95_ttfp_ms: float = 0.0
|
||||
p99_ttfp_ms: float = 0.0
|
||||
# E2E stats (ms)
|
||||
mean_e2e_ms: float = 0.0
|
||||
median_e2e_ms: float = 0.0
|
||||
std_e2e_ms: float = 0.0
|
||||
p90_e2e_ms: float = 0.0
|
||||
p95_e2e_ms: float = 0.0
|
||||
p99_e2e_ms: float = 0.0
|
||||
# RTF stats
|
||||
mean_rtf: float = 0.0
|
||||
median_rtf: float = 0.0
|
||||
std_rtf: float = 0.0
|
||||
p99_rtf: float = 0.0
|
||||
# Audio stats
|
||||
mean_audio_duration_s: float = 0.0
|
||||
total_audio_duration_s: float = 0.0
|
||||
audio_throughput: float = 0.0 # audio_duration / wall_time
|
||||
request_throughput: float = 0.0 # requests / second
|
||||
# Per-request details
|
||||
per_request: list = field(default_factory=list)
|
||||
|
||||
|
||||
def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width: int = 2) -> float:
|
||||
"""Convert raw PCM byte count to duration in seconds."""
|
||||
num_samples = num_bytes / sample_width
|
||||
return num_samples / sample_rate
|
||||
|
||||
|
||||
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)
|
||||
|
||||
result = RequestResult(prompt=prompt)
|
||||
st = time.perf_counter()
|
||||
|
||||
try:
|
||||
async with session.post(api_url, json=payload) as response:
|
||||
if response.status != 200:
|
||||
result.error = f"HTTP {response.status}: {await response.text()}"
|
||||
result.success = False
|
||||
return result
|
||||
|
||||
first_chunk = True
|
||||
total_bytes = 0
|
||||
|
||||
async for chunk in response.content.iter_any():
|
||||
if first_chunk and len(chunk) > 0:
|
||||
result.ttfp = time.perf_counter() - st
|
||||
first_chunk = False
|
||||
total_bytes += len(chunk)
|
||||
|
||||
result.e2e = time.perf_counter() - st
|
||||
result.audio_bytes = total_bytes
|
||||
result.audio_duration = pcm_bytes_to_duration(total_bytes)
|
||||
|
||||
if result.audio_duration > 0:
|
||||
result.rtf = result.e2e / result.audio_duration
|
||||
result.success = True
|
||||
|
||||
except Exception as e:
|
||||
result.error = str(e)
|
||||
result.success = False
|
||||
result.e2e = time.perf_counter() - st
|
||||
|
||||
if pbar:
|
||||
pbar.update(1)
|
||||
return result
|
||||
|
||||
|
||||
async def run_benchmark(
|
||||
host: str,
|
||||
port: int,
|
||||
num_prompts: int,
|
||||
max_concurrency: int,
|
||||
num_warmups: int = 3,
|
||||
task_type: str = "CustomVoice",
|
||||
voice: str = "vivian",
|
||||
language: str = "English",
|
||||
) -> BenchmarkResult:
|
||||
"""Run benchmark at a given concurrency level."""
|
||||
api_url = f"http://{host}:{port}/v1/audio/speech"
|
||||
|
||||
connector = aiohttp.TCPConnector(
|
||||
limit=max_concurrency,
|
||||
limit_per_host=max_concurrency,
|
||||
keepalive_timeout=60,
|
||||
)
|
||||
session = aiohttp.ClientSession(
|
||||
connector=connector,
|
||||
timeout=aiohttp.ClientTimeout(total=600),
|
||||
)
|
||||
|
||||
# Warmup
|
||||
if num_warmups > 0:
|
||||
print(f" Warming up with {num_warmups} requests...")
|
||||
warmup_tasks = []
|
||||
for i in range(num_warmups):
|
||||
prompt = PROMPTS[i % len(PROMPTS)]
|
||||
warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
|
||||
await asyncio.gather(*warmup_tasks)
|
||||
print(" Warmup done.")
|
||||
|
||||
# Build request list
|
||||
request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)]
|
||||
|
||||
# Run benchmark
|
||||
print(f" Running {num_prompts} requests with concurrency={max_concurrency}...")
|
||||
semaphore = asyncio.Semaphore(max_concurrency)
|
||||
pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}")
|
||||
|
||||
async def limited_request(prompt):
|
||||
async with semaphore:
|
||||
return await send_tts_request(session, api_url, prompt, task_type, voice, language, pbar)
|
||||
|
||||
start_time = time.perf_counter()
|
||||
tasks = [asyncio.create_task(limited_request(p)) for p in request_prompts]
|
||||
results: list[RequestResult] = await asyncio.gather(*tasks)
|
||||
duration = time.perf_counter() - start_time
|
||||
pbar.close()
|
||||
|
||||
await session.close()
|
||||
|
||||
# Compute stats
|
||||
successful = [r for r in results if r.success]
|
||||
failed = [r for r in results if not r.success]
|
||||
|
||||
bench = BenchmarkResult(
|
||||
concurrency=max_concurrency,
|
||||
num_prompts=num_prompts,
|
||||
completed=len(successful),
|
||||
failed=len(failed),
|
||||
duration_s=duration,
|
||||
)
|
||||
|
||||
if successful:
|
||||
ttfps = [r.ttfp * 1000 for r in successful] # convert to ms
|
||||
e2es = [r.e2e * 1000 for r in successful]
|
||||
rtfs = [r.rtf for r in successful]
|
||||
audio_durs = [r.audio_duration for r in successful]
|
||||
|
||||
bench.mean_ttfp_ms = float(np.mean(ttfps))
|
||||
bench.median_ttfp_ms = float(np.median(ttfps))
|
||||
bench.std_ttfp_ms = float(np.std(ttfps))
|
||||
bench.p90_ttfp_ms = float(np.percentile(ttfps, 90))
|
||||
bench.p95_ttfp_ms = float(np.percentile(ttfps, 95))
|
||||
bench.p99_ttfp_ms = float(np.percentile(ttfps, 99))
|
||||
|
||||
bench.mean_e2e_ms = float(np.mean(e2es))
|
||||
bench.median_e2e_ms = float(np.median(e2es))
|
||||
bench.std_e2e_ms = float(np.std(e2es))
|
||||
bench.p90_e2e_ms = float(np.percentile(e2es, 90))
|
||||
bench.p95_e2e_ms = float(np.percentile(e2es, 95))
|
||||
bench.p99_e2e_ms = float(np.percentile(e2es, 99))
|
||||
|
||||
bench.mean_rtf = float(np.mean(rtfs))
|
||||
bench.median_rtf = float(np.median(rtfs))
|
||||
bench.std_rtf = float(np.std(rtfs))
|
||||
bench.p99_rtf = float(np.percentile(rtfs, 99))
|
||||
|
||||
bench.mean_audio_duration_s = float(np.mean(audio_durs))
|
||||
bench.total_audio_duration_s = float(np.sum(audio_durs))
|
||||
bench.audio_throughput = bench.total_audio_duration_s / duration
|
||||
bench.request_throughput = len(successful) / duration
|
||||
|
||||
bench.per_request = [
|
||||
{
|
||||
"ttfp_ms": r.ttfp * 1000,
|
||||
"e2e_ms": r.e2e * 1000,
|
||||
"rtf": r.rtf,
|
||||
"audio_duration_s": r.audio_duration,
|
||||
"prompt": r.prompt,
|
||||
}
|
||||
for r in successful
|
||||
]
|
||||
|
||||
# Print summary in standardized performance template
|
||||
W = 50
|
||||
print("")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Serving Benchmark Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Successful requests:':<40}{bench.completed:<10}")
|
||||
print(f"{'Failed requests:':<40}{bench.failed:<10}")
|
||||
print(f"{'Maximum request concurrency:':<40}{max_concurrency:<10}")
|
||||
print(f"{'Benchmark duration (s):':<40}{duration:<10.2f}")
|
||||
print(f"{'Request throughput (req/s):':<40}{bench.request_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'End-to-end Latency':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean E2EL (ms):':<40}{bench.mean_e2e_ms:<10.2f}")
|
||||
print(f"{'Median E2EL (ms):':<40}{bench.median_e2e_ms:<10.2f}")
|
||||
print(f"{'P99 E2EL (ms):':<40}{bench.p99_e2e_ms:<10.2f}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Audio Result':^{W}}")
|
||||
print(f"{'=' * W}")
|
||||
print(f"{'Total audio duration generated (s):':<40}{bench.total_audio_duration_s:<10.2f}")
|
||||
print(f"{'Audio throughput (audio duration/s):':<40}{bench.audio_throughput:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Time to First Packet':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_TTFP (ms):':<40}{bench.mean_ttfp_ms:<10.2f}")
|
||||
print(f"{'Median AUDIO_TTFP (ms):':<40}{bench.median_ttfp_ms:<10.2f}")
|
||||
print(f"{'P99 AUDIO_TTFP (ms):':<40}{bench.p99_ttfp_ms:<10.2f}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Real Time Factor':^{W}}")
|
||||
print(f"{'-' * W}")
|
||||
print(f"{'Mean AUDIO_RTF:':<40}{bench.mean_rtf:<10.3f}")
|
||||
print(f"{'Median AUDIO_RTF:':<40}{bench.median_rtf:<10.3f}")
|
||||
print(f"{'P99 AUDIO_RTF:':<40}{bench.p99_rtf:<10.3f}")
|
||||
print(f"{'=' * W}")
|
||||
print("")
|
||||
|
||||
if failed:
|
||||
for r in failed[:3]:
|
||||
print(f" [ERROR] {r.error[:200]}")
|
||||
|
||||
return bench
|
||||
|
||||
|
||||
async def main(args):
|
||||
all_results = []
|
||||
|
||||
for concurrency in args.max_concurrency:
|
||||
result = await run_benchmark(
|
||||
host=args.host,
|
||||
port=args.port,
|
||||
num_prompts=args.num_prompts,
|
||||
max_concurrency=concurrency,
|
||||
num_warmups=args.num_warmups,
|
||||
task_type=args.task_type,
|
||||
voice=args.voice,
|
||||
language=args.language,
|
||||
)
|
||||
result.config_name = args.config_name
|
||||
all_results.append(asdict(result))
|
||||
|
||||
# Save results
|
||||
result_dir = Path(args.result_dir)
|
||||
result_dir.mkdir(parents=True, exist_ok=True)
|
||||
timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
|
||||
result_file = result_dir / f"bench_{args.config_name}_{timestamp}.json"
|
||||
|
||||
with open(result_file, "w") as f:
|
||||
json.dump(all_results, f, indent=2)
|
||||
print(f"Results saved to {result_file}")
|
||||
|
||||
return all_results
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Qwen3-TTS Benchmark Client")
|
||||
parser.add_argument("--host", type=str, default="127.0.0.1")
|
||||
parser.add_argument("--port", type=int, default=8000)
|
||||
parser.add_argument("--num-prompts", type=int, default=50, help="Number of prompts per concurrency level")
|
||||
parser.add_argument( # noqa: E501
|
||||
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
|
||||
)
|
||||
parser.add_argument("--num-warmups", type=int, default=3)
|
||||
parser.add_argument("--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(
|
||||
"--config-name", type=str, default="async_chunk", help="Label for this config (used in filenames)"
|
||||
)
|
||||
parser.add_argument("--result-dir", type=str, default="results")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
asyncio.run(main(args))
|
||||
@@ -1,93 +0,0 @@
|
||||
# Qwen3-TTS batch_size=1 config (streaming with async_chunk)
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -1,94 +0,0 @@
|
||||
# Qwen3-TTS max_num_seqs=16 config (streaming with async_chunk)
|
||||
# High-throughput concurrent request processing
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 4096
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 16384
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 16
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -1,94 +0,0 @@
|
||||
# Qwen3-TTS batch_size=4 config (streaming with async_chunk)
|
||||
# Enables concurrent request processing
|
||||
# 2-stage pipeline: Talker -> Code2Wav
|
||||
async_chunk: true
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_comprehension: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 4
|
||||
model_stage: qwen3_tts
|
||||
model_arch: Qwen3TTSTalkerForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
enforce_eager: false
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: latent
|
||||
gpu_memory_utilization: 0.3
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 512
|
||||
max_model_len: 4096
|
||||
custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_tts.talker2code2wav_async_chunk
|
||||
output_connectors:
|
||||
to_stage_1: connector_of_shared_memory
|
||||
default_sampling_params:
|
||||
temperature: 0.9
|
||||
top_k: 50
|
||||
max_tokens: 4096
|
||||
seed: 42
|
||||
detokenize: false
|
||||
repetition_penalty: 1.05
|
||||
stop_token_ids: [2150]
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 4
|
||||
model_stage: code2wav
|
||||
model_arch: Qwen3TTSCode2Wav
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
async_scheduling: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: audio
|
||||
gpu_memory_utilization: 0.2
|
||||
distributed_executor_backend: "mp"
|
||||
max_num_batched_tokens: 8192
|
||||
max_model_len: 32768
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
input_connectors:
|
||||
from_stage_0: connector_of_shared_memory
|
||||
tts_args:
|
||||
max_instructions_length: 500
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 65536
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.0
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 4
|
||||
|
||||
connectors:
|
||||
connector_of_shared_memory:
|
||||
name: SharedMemoryConnector
|
||||
extra:
|
||||
shm_threshold_bytes: 65536
|
||||
codec_streaming: true
|
||||
connector_get_sleep_s: 0.01
|
||||
connector_get_max_wait_first_chunk: 3000
|
||||
connector_get_max_wait: 300
|
||||
codec_chunk_frames: 25
|
||||
codec_left_context_frames: 25
|
||||
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -1,249 +0,0 @@
|
||||
"""Plot TTFP comparison: async_chunk off vs on.
|
||||
|
||||
Generates a bar chart with improvement arrows, matching the Qwen3-Omni
|
||||
async_chunk benchmark figure style.
|
||||
|
||||
Usage:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--output results/qwen3_tts_async_chunk_ttfp.png
|
||||
|
||||
# Also supports E2E and RTF metrics:
|
||||
python plot_async_chunk.py \
|
||||
--off results/bench_async_chunk_off_*.json \
|
||||
--on results/bench_async_chunk_on_*.json \
|
||||
--metric e2e \
|
||||
--output results/qwen3_tts_async_chunk_e2e.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
|
||||
METRIC_CONFIG = {
|
||||
"ttfp": {
|
||||
"key": "mean_ttfp_ms",
|
||||
"ylabel": "TTFP (s)",
|
||||
"title": "TTFP (Time to First Audio Packet) - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"e2e": {
|
||||
"key": "mean_e2e_ms",
|
||||
"ylabel": "E2E (s)",
|
||||
"title": "E2E Latency - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": True,
|
||||
},
|
||||
"rtf": {
|
||||
"key": "mean_rtf",
|
||||
"ylabel": "RTF",
|
||||
"title": "Real-Time Factor - Qwen3-TTS, by concurrency",
|
||||
"to_seconds": False,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def load_results(path: str) -> list[dict]:
|
||||
with open(path) as f:
|
||||
return json.load(f)
|
||||
|
||||
|
||||
def plot_ttfp_comparison(
|
||||
off_results: list[dict],
|
||||
on_results: list[dict],
|
||||
metric: str,
|
||||
output_path: str,
|
||||
title_override: str | None = None,
|
||||
):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
fig, ax = plt.subplots(figsize=(8, 6))
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB", edgecolor="none")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
# Draw improvement arrows and labels
|
||||
for i in range(len(concurrencies)):
|
||||
v_off = off_vals[i]
|
||||
v_on = on_vals[i]
|
||||
if v_on > 0:
|
||||
improvement = v_off / v_on
|
||||
else:
|
||||
improvement = float("inf")
|
||||
|
||||
# Arrow from top of off-bar to top of on-bar
|
||||
arrow_start_x = x[i] - width / 2
|
||||
arrow_start_y = v_off * 0.95
|
||||
arrow_end_x = x[i] + width / 2
|
||||
arrow_end_y = v_on * 1.05
|
||||
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(arrow_end_x, arrow_end_y),
|
||||
xytext=(arrow_start_x, arrow_start_y),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
|
||||
# Improvement label
|
||||
label_x = (arrow_start_x + arrow_end_x) / 2
|
||||
label_y = arrow_start_y + (v_off - v_on) * 0.15
|
||||
ax.text(
|
||||
label_x,
|
||||
label_y,
|
||||
f"{improvement:.1f}x improvement",
|
||||
ha="center",
|
||||
va="bottom",
|
||||
fontsize=10,
|
||||
color="red",
|
||||
fontweight="bold",
|
||||
)
|
||||
|
||||
title = title_override or cfg["title"]
|
||||
ax.set_title(title, fontsize=13, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=12)
|
||||
ax.set_xlabel("Max concurrency", fontsize=12)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
ax.set_yscale("log")
|
||||
ax.legend(loc="upper left", fontsize=11)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def plot_all_metrics(off_results: list[dict], on_results: list[dict], output_path: str):
|
||||
"""Generate a 1x3 subplot with TTFP, E2E, and RTF comparisons."""
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
fig, axes = plt.subplots(1, 3, figsize=(18, 6))
|
||||
fig.suptitle("Qwen3-TTS: async_chunk on vs off", fontsize=15, fontweight="bold")
|
||||
|
||||
for ax, metric in zip(axes, ["ttfp", "e2e", "rtf"]):
|
||||
cfg = METRIC_CONFIG[metric]
|
||||
key = cfg["key"]
|
||||
to_seconds = cfg["to_seconds"]
|
||||
|
||||
off_vals = []
|
||||
on_vals = []
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c][key]
|
||||
v_on = on_map[c][key]
|
||||
if to_seconds:
|
||||
v_off /= 1000.0
|
||||
v_on /= 1000.0
|
||||
off_vals.append(v_off)
|
||||
on_vals.append(v_on)
|
||||
|
||||
x = np.arange(len(concurrencies))
|
||||
width = 0.3
|
||||
ax.bar(x - width / 2, off_vals, width, label="async_chunk off", color="#87CEEB")
|
||||
ax.bar(x + width / 2, on_vals, width, label="async_chunk on", color="#FFF8DC", edgecolor="#DDD8B8")
|
||||
|
||||
for i in range(len(concurrencies)):
|
||||
if on_vals[i] > 0:
|
||||
improvement = off_vals[i] / on_vals[i]
|
||||
ax.annotate(
|
||||
"",
|
||||
xy=(x[i] + width / 2, on_vals[i] * 1.05),
|
||||
xytext=(x[i] - width / 2, off_vals[i] * 0.95),
|
||||
arrowprops=dict(arrowstyle="->", color="red", lw=1.5),
|
||||
)
|
||||
label_y = off_vals[i] * 0.85
|
||||
ax.text(x[i], label_y, f"{improvement:.1f}x", ha="center", fontsize=10, color="red", fontweight="bold")
|
||||
|
||||
ax.set_title(cfg["title"].split(" - ")[0], fontsize=12, fontweight="bold")
|
||||
ax.set_ylabel(cfg["ylabel"], fontsize=11)
|
||||
ax.set_xlabel("Max concurrency", fontsize=11)
|
||||
ax.set_xticks(x)
|
||||
ax.set_xticklabels([str(c) for c in concurrencies])
|
||||
if metric != "rtf":
|
||||
ax.set_yscale("log")
|
||||
ax.legend(fontsize=9)
|
||||
ax.grid(axis="y", alpha=0.3, linestyle="--")
|
||||
ax.set_axisbelow(True)
|
||||
|
||||
plt.tight_layout()
|
||||
Path(output_path).parent.mkdir(parents=True, exist_ok=True)
|
||||
plt.savefig(output_path, dpi=150, bbox_inches="tight")
|
||||
print(f"Plot saved to {output_path}")
|
||||
plt.close()
|
||||
|
||||
|
||||
def print_table(off_results: list[dict], on_results: list[dict]):
|
||||
off_map = {r["concurrency"]: r for r in off_results}
|
||||
on_map = {r["concurrency"]: r for r in on_results}
|
||||
concurrencies = sorted(set(off_map.keys()) & set(on_map.keys()))
|
||||
|
||||
print("\n## Benchmark Results: async_chunk off vs on\n")
|
||||
print("| Metric | Concurrency | async_chunk off | async_chunk on | Improvement |")
|
||||
print("| --- | --- | --- | --- | --- |")
|
||||
|
||||
for name, key, fmt in [
|
||||
("TTFP (ms)", "mean_ttfp_ms", ".1f"),
|
||||
("E2E (ms)", "mean_e2e_ms", ".1f"),
|
||||
("RTF", "mean_rtf", ".3f"),
|
||||
("Throughput", "audio_throughput", ".2f"),
|
||||
]:
|
||||
for c in concurrencies:
|
||||
v_off = off_map[c].get(key, 0)
|
||||
v_on = on_map[c].get(key, 0)
|
||||
if v_on > 0 and key != "audio_throughput":
|
||||
ratio = f"{v_off / v_on:.1f}x"
|
||||
elif v_off > 0 and key == "audio_throughput":
|
||||
ratio = f"{v_on / v_off:.1f}x"
|
||||
else:
|
||||
ratio = "N/A"
|
||||
print(f"| {name} | {c} | {v_off:{fmt}} | {v_on:{fmt}} | {ratio} |")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Plot async_chunk comparison for Qwen3-TTS")
|
||||
parser.add_argument("--off", type=str, required=True, help="JSON results for async_chunk off")
|
||||
parser.add_argument("--on", type=str, required=True, help="JSON results for async_chunk on")
|
||||
parser.add_argument("--metric", type=str, default="ttfp", choices=["ttfp", "e2e", "rtf", "all"])
|
||||
parser.add_argument("--output", type=str, default="results/qwen3_tts_async_chunk.png")
|
||||
parser.add_argument("--title", type=str, default=None, help="Custom title override")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
off_results = load_results(args.off)
|
||||
on_results = load_results(args.on)
|
||||
|
||||
print_table(off_results, on_results)
|
||||
|
||||
if args.metric == "all":
|
||||
plot_all_metrics(off_results, on_results, args.output)
|
||||
else:
|
||||
plot_ttfp_comparison(off_results, on_results, args.metric, args.output, args.title)
|
||||
@@ -1,167 +0,0 @@
|
||||
#!/bin/bash
|
||||
# Qwen3-TTS async_chunk on vs off Benchmark
|
||||
#
|
||||
# Starts two servers (async_chunk on and off), benchmarks both,
|
||||
# and generates comparison plots.
|
||||
#
|
||||
# Usage:
|
||||
# bash run_async_chunk_benchmark.sh
|
||||
#
|
||||
# Environment variables:
|
||||
# GPU_DEVICE - GPU index (default: 0)
|
||||
# MODEL - Model path (default: Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice)
|
||||
# NUM_PROMPTS - Prompts per concurrency level (default: 50)
|
||||
# CONCURRENCY - Space-separated concurrency levels (default: "1 10")
|
||||
# PORT_ON - Port for async_chunk on server (default: 8000)
|
||||
# PORT_OFF - Port for async_chunk off server (default: 8001)
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
PROJECT_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
cd "$PROJECT_ROOT"
|
||||
|
||||
GPU_DEVICE="${GPU_DEVICE:-0}"
|
||||
MODEL="${MODEL:-Qwen/Qwen3-TTS-12Hz-0.6B-CustomVoice}"
|
||||
NUM_PROMPTS="${NUM_PROMPTS:-50}"
|
||||
CONCURRENCY="${CONCURRENCY:-1 10}"
|
||||
NUM_WARMUPS="${NUM_WARMUPS:-3}"
|
||||
PORT_ON="${PORT_ON:-8000}"
|
||||
PORT_OFF="${PORT_OFF:-8001}"
|
||||
RESULT_DIR="${SCRIPT_DIR}/results"
|
||||
TIMESTAMP="$(date +%Y%m%d_%H%M%S)"
|
||||
|
||||
STAGE_CONFIG_ON="vllm_omni/model_executor/stage_configs/qwen3_tts.yaml"
|
||||
STAGE_CONFIG_OFF="vllm_omni/model_executor/stage_configs/qwen3_tts_no_async_chunk.yaml"
|
||||
|
||||
mkdir -p "${RESULT_DIR}"
|
||||
|
||||
echo "============================================================"
|
||||
echo " Qwen3-TTS async_chunk Benchmark"
|
||||
echo "============================================================"
|
||||
echo " GPU: ${GPU_DEVICE}"
|
||||
echo " Model: ${MODEL}"
|
||||
echo " Prompts: ${NUM_PROMPTS}"
|
||||
echo " Concurrency: ${CONCURRENCY}"
|
||||
echo " Port (on/off): ${PORT_ON} / ${PORT_OFF}"
|
||||
echo " Results: ${RESULT_DIR}"
|
||||
echo "============================================================"
|
||||
|
||||
cleanup() {
|
||||
echo "Cleaning up servers..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup EXIT
|
||||
|
||||
wait_for_server() {
|
||||
local port=$1
|
||||
local name=$2
|
||||
local max_wait=300
|
||||
local elapsed=0
|
||||
echo "Waiting for ${name} server on port ${port}..."
|
||||
while ! curl -s "http://localhost:${port}/health" >/dev/null 2>&1; do
|
||||
sleep 5
|
||||
elapsed=$((elapsed + 5))
|
||||
if [ $elapsed -ge $max_wait ]; then
|
||||
echo "ERROR: ${name} server failed to start within ${max_wait}s"
|
||||
exit 1
|
||||
fi
|
||||
done
|
||||
echo "${name} server ready (${elapsed}s)"
|
||||
}
|
||||
|
||||
# ---- Phase 1: Start async_chunk ON server ----
|
||||
echo ""
|
||||
echo "[Phase 1] Starting async_chunk ON server on port ${PORT_ON}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_ON}" \
|
||||
--host 0.0.0.0 --port "${PORT_ON}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_on_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_ON=$!
|
||||
|
||||
wait_for_server "${PORT_ON}" "async_chunk_on"
|
||||
|
||||
echo "[Phase 1] Benchmarking async_chunk ON..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_ON}" \
|
||||
--config-name "async_chunk_on" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 1] Stopping async_chunk ON server..."
|
||||
kill "$PID_ON" 2>/dev/null || true
|
||||
wait "$PID_ON" 2>/dev/null || true
|
||||
sleep 5
|
||||
|
||||
# ---- Phase 2: Start async_chunk OFF server ----
|
||||
echo ""
|
||||
echo "[Phase 2] Starting async_chunk OFF server on port ${PORT_OFF}..."
|
||||
CUDA_VISIBLE_DEVICES=${GPU_DEVICE} vllm-omni serve "${MODEL}" \
|
||||
--stage-configs-path "${STAGE_CONFIG_OFF}" \
|
||||
--host 0.0.0.0 --port "${PORT_OFF}" \
|
||||
--trust-remote-code --enforce-eager --omni \
|
||||
> "${RESULT_DIR}/server_off_${TIMESTAMP}.log" 2>&1 &
|
||||
PID_OFF=$!
|
||||
|
||||
wait_for_server "${PORT_OFF}" "async_chunk_off"
|
||||
|
||||
echo "[Phase 2] Benchmarking async_chunk OFF (non-streaming)..."
|
||||
# shellcheck disable=SC2086
|
||||
python "${SCRIPT_DIR}/bench_async_chunk.py" \
|
||||
--host 127.0.0.1 --port "${PORT_OFF}" \
|
||||
--config-name "async_chunk_off" \
|
||||
--num-prompts "${NUM_PROMPTS}" \
|
||||
--max-concurrency ${CONCURRENCY} \
|
||||
--num-warmups "${NUM_WARMUPS}" \
|
||||
--no-stream \
|
||||
--result-dir "${RESULT_DIR}"
|
||||
|
||||
echo "[Phase 2] Stopping async_chunk OFF server..."
|
||||
kill "$PID_OFF" 2>/dev/null || true
|
||||
wait "$PID_OFF" 2>/dev/null || true
|
||||
|
||||
# ---- Phase 3: Plot results ----
|
||||
echo ""
|
||||
echo "[Phase 3] Generating plots..."
|
||||
|
||||
# Find the latest result files
|
||||
RESULT_ON=$(ls -t "${RESULT_DIR}"/bench_async_chunk_on_*.json 2>/dev/null | head -1)
|
||||
RESULT_OFF=$(ls -t "${RESULT_DIR}"/bench_async_chunk_off_*.json 2>/dev/null | head -1)
|
||||
|
||||
if [ -z "$RESULT_ON" ] || [ -z "$RESULT_OFF" ]; then
|
||||
echo "ERROR: Could not find result files. Check logs in ${RESULT_DIR}/"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo " ON results: ${RESULT_ON}"
|
||||
echo " OFF results: ${RESULT_OFF}"
|
||||
|
||||
# TTFP comparison (main figure)
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric ttfp \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
|
||||
# All metrics comparison
|
||||
python "${SCRIPT_DIR}/plot_async_chunk.py" \
|
||||
--off "${RESULT_OFF}" \
|
||||
--on "${RESULT_ON}" \
|
||||
--metric all \
|
||||
--output "${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
|
||||
echo ""
|
||||
echo "============================================================"
|
||||
echo " Benchmark complete!"
|
||||
echo " Results: ${RESULT_DIR}/"
|
||||
echo " Plots:"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_ttfp.png"
|
||||
echo " - ${RESULT_DIR}/qwen3_tts_async_chunk_all.png"
|
||||
echo "============================================================"
|
||||
@@ -1,17 +1,20 @@
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.19.0
|
||||
ARG VLLM_BASE_TAG=v0.16.0
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||
ARG APP_DIR=/workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
# Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
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
|
||||
|
||||
|
||||
@@ -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 []
|
||||
+6
-19
@@ -1,30 +1,17 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.18.0rc1
|
||||
ARG VLLM_ASCEND_TAG=v0.14.0rc1
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
# RUN git fetch origin --tags && git checkout v0.18.0
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm-ascend
|
||||
# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
|
||||
# # Install vllm-ascend
|
||||
# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
# source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
# python3 -m pip cache purge
|
||||
WORKDIR /vllm-workspace/vllm-ascend
|
||||
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
|
||||
RUN pip install -v -e .
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
|
||||
RUN pip install -v -e .
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
|
||||
@@ -1,30 +1,17 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
|
||||
ARG VLLM_ASCEND_TAG=v0.14.0rc1-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
# RUN git fetch origin --tags && git checkout v0.18.0
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm-ascend
|
||||
# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
|
||||
# # Install vllm-ascend
|
||||
# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
|
||||
# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
# source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
|
||||
# python3 -m pip cache purge
|
||||
WORKDIR /vllm-workspace/vllm-ascend
|
||||
RUN git checkout e2175d9c7e62b437391dfee996b1375674ba7c18
|
||||
RUN pip install -v -e .
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
|
||||
COPY . .
|
||||
|
||||
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
|
||||
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
|
||||
source /usr/local/Ascend/nnal/atb/set_env.sh && \
|
||||
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
|
||||
python3 -m pip install -v -e /vllm-workspace/vllm-omni/ --no-build-isolation
|
||||
RUN pip install -v -e .
|
||||
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.16.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
@@ -10,7 +10,7 @@ FROM ${BASE_IMAGE} AS base
|
||||
# we should swap over to use stable release ASAP.
|
||||
# We should avoid relying on custom commits.
|
||||
ARG USE_NIGHTLY_BUILD=0
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=89138b21cc246ae944c741d5c399c148e2b770ab
|
||||
ARG VLLM_VERSION_OR_COMMIT_HASH=2d5be1dd5ce2e44dfea53ea03ff61143da5137eb
|
||||
ARG ARG_PYTORCH_ROCM_ARCH
|
||||
ENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}
|
||||
|
||||
@@ -19,7 +19,7 @@ WORKDIR ${COMMON_WORKDIR}
|
||||
|
||||
# Step 1: Setup - Install system dependencies
|
||||
RUN apt-get update && \
|
||||
apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get install -y ffmpeg git sox libsox-fmt-all jq && \
|
||||
apt-get clean && \
|
||||
rm -rf /var/lib/apt/lists/*
|
||||
|
||||
|
||||
+4
-17
@@ -9,12 +9,12 @@ ARG PYTHON_VERSION=3.12
|
||||
ARG PIP_EXTRA_INDEX_URL="https://download.pytorch.org/whl/xpu"
|
||||
|
||||
RUN wget -O- https://apt.repos.intel.com/intel-gpg-keys/GPG-PUB-KEY-INTEL-SW-PRODUCTS.PUB | gpg --dearmor | tee /usr/share/keyrings/oneapi-archive-keyring.gpg > /dev/null && \
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list
|
||||
echo "deb [signed-by=/usr/share/keyrings/oneapi-archive-keyring.gpg] https://apt.repos.intel.com/oneapi all main" | tee /etc/apt/sources.list.d/oneAPI.list && \
|
||||
add-apt-repository -y ppa:kobuk-team/intel-graphics
|
||||
|
||||
RUN apt clean && apt-get update -y && \
|
||||
apt-get install -y --no-install-recommends --fix-missing \
|
||||
curl \
|
||||
espeak-ng \
|
||||
ffmpeg \
|
||||
git \
|
||||
libsndfile1 \
|
||||
@@ -31,22 +31,9 @@ RUN apt clean && apt-get update -y && \
|
||||
python3-pip
|
||||
|
||||
RUN apt update && apt upgrade -y && \
|
||||
apt install -y libze1 libze-dev libze-intel-gpu1 intel-opencl-icd libze-intel-gpu-raytracing intel-ocloc && \
|
||||
apt install -y intel-oneapi-compiler-dpcpp-cpp-2025.3
|
||||
|
||||
# Install UMD with fixed version
|
||||
RUN mkdir neo && \
|
||||
cd neo && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-core-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/intel-graphics-compiler/releases/download/v2.24.8/intel-igc-opencl-2_2.24.8+20344_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-ocloc_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/intel-opencl-icd_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libigdgmm12_22.8.2_amd64.deb && \
|
||||
wget https://github.com/intel/compute-runtime/releases/download/25.48.36300.8/libze-intel-gpu1_25.48.36300.8-0_amd64.deb && \
|
||||
wget https://github.com/oneapi-src/level-zero/releases/download/v1.26.0/level-zero_1.26.0+u24.04_amd64.deb && \
|
||||
dpkg -i *.deb && \
|
||||
cd .. && \
|
||||
rm -rf neo
|
||||
|
||||
ENV PATH="/root/.local/bin:$PATH"
|
||||
ENV VIRTUAL_ENV="/opt/venv"
|
||||
ENV UV_PYTHON_INSTALL_DIR=/opt/uv/python
|
||||
@@ -76,7 +63,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.16.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
|
||||
+5
-34
@@ -6,7 +6,6 @@ nav:
|
||||
- getting_started/installation/*
|
||||
- Serving:
|
||||
- OpenAI-Compatible API:
|
||||
- Diffusion Chat API: serving/diffusion_chat_api.md
|
||||
- Image Generation: serving/image_generation_api.md
|
||||
- Image Edit: serving/image_edit_api.md
|
||||
- Text to Speech: serving/speech_api.md
|
||||
@@ -14,14 +13,9 @@ nav:
|
||||
- examples/README.md
|
||||
- Offline Inference:
|
||||
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
|
||||
- CosyVoice3: user_guide/examples/offline_inference/cosyvoice3.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/offline_inference/fish_speech.md
|
||||
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
|
||||
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
|
||||
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
|
||||
- MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md
|
||||
- MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md
|
||||
- Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md
|
||||
- Qwen3-Omni: user_guide/examples/offline_inference/qwen3_omni.md
|
||||
@@ -29,12 +23,8 @@ nav:
|
||||
- Text-To-Audio: user_guide/examples/offline_inference/text_to_audio.md
|
||||
- Text-To-Image: user_guide/examples/offline_inference/text_to_image.md
|
||||
- Text-To-Video: user_guide/examples/offline_inference/text_to_video.md
|
||||
- Voxtral TTS Offline Inference: user_guide/examples/offline_inference/voxtral_tts.md
|
||||
- X-To-Video-Audio: user_guide/examples/offline_inference/x_to_video_audio.md
|
||||
- Online Serving:
|
||||
- BAGEL-7B-MoT: user_guide/examples/online_serving/bagel.md
|
||||
- vLLM-Omni Helm Chart: user_guide/examples/online_serving/chart-helm.md
|
||||
- Fish Speech S2 Pro: user_guide/examples/online_serving/fish_speech.md
|
||||
- GLM-Image Online Serving: user_guide/examples/online_serving/glm_image.md
|
||||
- Image-To-Image: user_guide/examples/online_serving/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/online_serving/image_to_video.md
|
||||
@@ -54,28 +44,18 @@ nav:
|
||||
- Features:
|
||||
- Sleep Mode: features/sleep_mode.md
|
||||
- Diffusion Features:
|
||||
- Overview: user_guide/diffusion_features.md
|
||||
- Feature Compatibility: user_guide/feature_compatibility.md
|
||||
- Cache Acceleration:
|
||||
- TeaCache: user_guide/diffusion/cache_acceleration/teacache.md
|
||||
- Cache-DiT: user_guide/diffusion/cache_acceleration/cache_dit.md
|
||||
- Acceleration Overview: user_guide/diffusion_acceleration.md
|
||||
- TeaCache: user_guide/diffusion/teacache.md
|
||||
- Cache-DiT: user_guide/diffusion/cache_dit_acceleration.md
|
||||
- Quantization:
|
||||
- Overview: user_guide/diffusion/quantization/overview.md
|
||||
- FP8: user_guide/diffusion/quantization/fp8.md
|
||||
- Int8: user_guide/diffusion/quantization/int8.md
|
||||
- GGUF: user_guide/diffusion/quantization/gguf.md
|
||||
- Parallelism:
|
||||
- Overview: user_guide/diffusion/parallelism/overview.md
|
||||
- CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md
|
||||
- Expert Parallel: user_guide/diffusion/parallelism/expert_parallel.md
|
||||
- Hybrid Sharded Data Parallel: user_guide/diffusion/parallelism/hsdp.md
|
||||
- Sequence Parallel: user_guide/diffusion/parallelism/sequence_parallel.md
|
||||
- Tensor Parallel: user_guide/diffusion/parallelism/tensor_parallel.md
|
||||
- VAE Patch Parallel: user_guide/diffusion/parallelism/vae_patch_parallel.md
|
||||
- Parallelism Acceleration: user_guide/diffusion/parallelism_acceleration.md
|
||||
- CPU Offloading: user_guide/diffusion/cpu_offload_diffusion.md
|
||||
- LoRA: user_guide/diffusion/lora.md
|
||||
- Hybrid Sharded Data Parallel: design/feature/hsdp.md
|
||||
- Custom Pipeline: features/custom_pipeline.md
|
||||
- Step Execution: user_guide/diffusion/step_execution.md
|
||||
- ComfyUI: features/comfyui.md
|
||||
- Developer Guide:
|
||||
- General:
|
||||
@@ -86,9 +66,7 @@ nav:
|
||||
- Model Implementation:
|
||||
- contributing/model/README.md
|
||||
- contributing/model/adding_omni_model.md
|
||||
- contributing/model/adding_tts_model.md
|
||||
- contributing/model/adding_diffusion_model.md
|
||||
- contributing/model/adding_quantization_model.md
|
||||
- CI: contributing/ci
|
||||
- Design Documents:
|
||||
- design/index.md
|
||||
@@ -98,28 +76,21 @@ 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
|
||||
- design/feature/hsdp.md
|
||||
- design/feature/cache_dit.md
|
||||
- design/feature/teacache.md
|
||||
- design/feature/async_chunk_design.md
|
||||
- design/feature/vae_parallel.md
|
||||
- design/feature/diffusion_step_execution.md
|
||||
- Module Design:
|
||||
- design/module/ar_module.md
|
||||
- design/module/dit_module.md
|
||||
- design/module/entrypoint_module.md
|
||||
- design/module/async_omni_architecture.md
|
||||
- Docs Guide: contributing/DOCS_GUIDE.md
|
||||
- API Reference:
|
||||
- api/README.md
|
||||
- api/vllm_omni
|
||||
- CLI Reference: cli
|
||||
- Community:
|
||||
- Governance: community/governance.md
|
||||
- community/*
|
||||
- Slack: https://slack.vllm.ai
|
||||
- Blog: https://blog.vllm.ai
|
||||
|
||||
+1
-1
@@ -60,5 +60,5 @@ vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, in
|
||||
|
||||
For more information, checkout the following:
|
||||
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1XJWgv79lORl8rbaVvp2d5Sqs6ZEBgAgj/edit?slide=id.p1#slide=id.p1)
|
||||
- [vllm-omni architecture design and recent roadmaps](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true)
|
||||
- [vllm-omni announcement blogpost](https://blog.vllm.ai/2025/11/30/vllm-omni.html)
|
||||
+10
-42
@@ -5,21 +5,25 @@
|
||||
Main entry points for vLLM-Omni inference and serving.
|
||||
|
||||
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
|
||||
- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
|
||||
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
|
||||
- [vllm_omni.entrypoints.async_omni_llm.AsyncOmniLLM][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][]
|
||||
- [vllm_omni.entrypoints.cli.serve.OmniServeCommand][]
|
||||
- [vllm_omni.entrypoints.client_request_state.ClientRequestState][]
|
||||
- [vllm_omni.entrypoints.omni.Omni][]
|
||||
- [vllm_omni.entrypoints.omni_base.OmniBase][]
|
||||
- [vllm_omni.entrypoints.pd_utils.PDDisaggregationMixin][]
|
||||
- [vllm_omni.entrypoints.omni.OmniBase][]
|
||||
- [vllm_omni.entrypoints.omni_diffusion.OmniDiffusion][]
|
||||
- [vllm_omni.entrypoints.omni_llm.OmniLLM][]
|
||||
- [vllm_omni.entrypoints.omni_stage.OmniStage][]
|
||||
- [vllm_omni.entrypoints.stage_utils.OmniStageTaskType][]
|
||||
- [vllm_omni.entrypoints.zmq_utils.ZmqQueue][]
|
||||
|
||||
## Inputs
|
||||
|
||||
Input data structures for multi-modal inputs.
|
||||
|
||||
- [vllm_omni.inputs.data.OmniCustomPrompt][]
|
||||
- [vllm_omni.inputs.data.OmniDiffusionSamplingParams][]
|
||||
- [vllm_omni.inputs.data.OmniEmbedsPrompt][]
|
||||
- [vllm_omni.inputs.data.OmniTextPrompt][]
|
||||
@@ -44,19 +48,11 @@ Engine classes for offline and online inference.
|
||||
- [vllm_omni.engine.OmniEngineCoreOutputs][]
|
||||
- [vllm_omni.engine.OmniEngineCoreRequest][]
|
||||
- [vllm_omni.engine.PromptEmbedsPayload][]
|
||||
- [vllm_omni.engine.arg_utils.AsyncOmniEngineArgs][]
|
||||
- [vllm_omni.engine.arg_utils.OmniEngineArgs][]
|
||||
- [vllm_omni.engine.async_omni_engine.AsyncOmniEngine][]
|
||||
- [vllm_omni.engine.mm_outputs.MultimodalCompletionOutput][]
|
||||
- [vllm_omni.engine.mm_outputs.MultimodalPayload][]
|
||||
- [vllm_omni.engine.orchestrator.Orchestrator][]
|
||||
- [vllm_omni.engine.orchestrator.OrchestratorRequestState][]
|
||||
- [vllm_omni.engine.output_modality.OutputModality][]
|
||||
- [vllm_omni.engine.output_modality.TensorAccumulationStrategy][]
|
||||
- [vllm_omni.engine.input_processor.OmniInputProcessor][]
|
||||
- [vllm_omni.engine.output_processor.MultimodalOutputProcessor][]
|
||||
- [vllm_omni.engine.output_processor.OmniRequestState][]
|
||||
- [vllm_omni.engine.stage_engine_core_client.StageEngineCoreClient][]
|
||||
- [vllm_omni.engine.stage_init_utils.StageMetadata][]
|
||||
- [vllm_omni.engine.stage_init_utils.StartedLlmStage][]
|
||||
|
||||
## Core
|
||||
|
||||
@@ -68,19 +64,6 @@ Core scheduling and caching components.
|
||||
- [vllm_omni.core.sched.output.OmniCachedRequestData][]
|
||||
- [vllm_omni.core.sched.output.OmniNewRequestData][]
|
||||
- [vllm_omni.core.sched.output.OmniSchedulerOutput][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.BASECFM][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalConditionalCFM][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.CausalMaskedDiffWithDiT][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.cfm.ConditionalCFM][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1d][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConv1dUpsample][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalConvRNNF0Predictor][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.CausalHiFTGenerator][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.HiFTGenerator][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SineGen2][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.Snake][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.code2wav_core.hifigan.SourceModuleHnNSF][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedGroupResidualVectorQuantization][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.DistributedResidualVectorQuantization][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.vq.core_vq.EuclideanCodebook][]
|
||||
@@ -92,17 +75,9 @@ Core scheduling and caching components.
|
||||
Configuration classes.
|
||||
|
||||
- [vllm_omni.config.model.OmniModelConfig][]
|
||||
- [vllm_omni.config.stage_config.ModelPipeline][]
|
||||
- [vllm_omni.config.stage_config.StageConfig][]
|
||||
- [vllm_omni.config.stage_config.StageConfigFactory][]
|
||||
- [vllm_omni.config.stage_config.StageType][]
|
||||
- [vllm_omni.diffusion.cache.teacache.config.TeaCacheConfig][]
|
||||
- [vllm_omni.distributed.omni_connectors.utils.config.ConnectorSpec][]
|
||||
- [vllm_omni.distributed.omni_connectors.utils.config.OmniTransferConfig][]
|
||||
- [vllm_omni.model_executor.models.cosyvoice3.config.CosyVoice3Config][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechConfig][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechFastARConfig][]
|
||||
- [vllm_omni.model_executor.models.fish_speech.configuration_fish_speech.FishSpeechSlowARConfig][]
|
||||
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioConfig][]
|
||||
- [vllm_omni.model_executor.models.mimo_audio.config_mimo_audio.MiMoAudioTokenizerConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts.Qwen3TTSConfig][]
|
||||
@@ -116,10 +91,6 @@ Configuration classes.
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1DecoderDiTConfig][]
|
||||
- [vllm_omni.model_executor.models.qwen3_tts.tokenizer_25hz.configuration_qwen3_tts_tokenizer_v1.Qwen3TTSTokenizerV1EncoderConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Config][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLTextConfig][]
|
||||
- [vllm_omni.transformers_utils.configs.mammoth_moda2.Mammothmoda2Qwen2_5_VLVisionConfig][]
|
||||
|
||||
## Workers
|
||||
|
||||
@@ -130,8 +101,6 @@ Worker classes and model runners for distributed inference.
|
||||
- [vllm_omni.diffusion.worker.diffusion_worker.DiffusionWorker][]
|
||||
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerProc][]
|
||||
- [vllm_omni.diffusion.worker.diffusion_worker.WorkerWrapperBase][]
|
||||
- [vllm_omni.diffusion.worker.utils.DiffusionRequestState][]
|
||||
- [vllm_omni.diffusion.worker.utils.RunnerOutput][]
|
||||
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.ExecuteModelState][]
|
||||
- [vllm_omni.platforms.npu.worker.npu_ar_model_runner.NPUARModelRunner][]
|
||||
- [vllm_omni.platforms.npu.worker.npu_ar_worker.NPUARWorker][]
|
||||
@@ -149,6 +118,5 @@ Worker classes and model runners for distributed inference.
|
||||
- [vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner][]
|
||||
- [vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker][]
|
||||
- [vllm_omni.worker.gpu_memory_utils.parse_cuda_visible_devices][]
|
||||
- [vllm_omni.worker.gpu_model_runner.CUDAGraphWrapper][]
|
||||
- [vllm_omni.worker.gpu_model_runner.OmniGPUModelRunner][]
|
||||
- [vllm_omni.worker.mixins.OmniWorkerMixin][]
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 96 KiB After Width: | Height: | Size: 127 KiB |
@@ -1,54 +0,0 @@
|
||||
# Governance
|
||||
|
||||
vLLM-Omni's governance is inspired by the [vLLM governance process](https://docs.vllm.ai/en/latest/governance/process/). We share the same commitment to open source community and meritocratic norms.
|
||||
|
||||
## Values
|
||||
|
||||
vLLM-Omni aims to be the fastest and easiest-to-use omni-modality inference and serving engine. Our values are aligned with [vLLM's values](https://docs.vllm.ai/en/latest/governance/process/#values):
|
||||
|
||||
### Design Values
|
||||
|
||||
1. **Top performance**: System performance is our top priority. We monitor overheads, optimize kernels, and publish benchmarks. We never leave performance on the table.
|
||||
2. **Ease of use**: vLLM-Omni must be simple to install, configure, and operate. We provide clear documentation, fast startup, clean logs, helpful error messages, and monitoring guides. Many users fork our code or study it deeply, so we keep it readable and modular.
|
||||
3. **Wide coverage**: vLLM-Omni supports frontier models and high-performance accelerators. We make it easy to add new models and hardware. vLLM-Omni + PyTorch form a simple interface that avoids complexity.
|
||||
4. **Production ready**: vLLM-Omni runs 24/7 in production. It must be easy to operate and monitor for health issues.
|
||||
5. **Extensibility**: vLLM-Omni serves as fundamental omni-modality infrastructure. Our codebase cannot cover every use case, so we design for easy forking and customization.
|
||||
|
||||
### Collaboration Values
|
||||
|
||||
1. **Tightly Knit and Fast-Moving**: Our maintainer team is aligned on vision, philosophy, and roadmap. We work closely to unblock each other and move quickly.
|
||||
2. **Individual Merit**: No one buys their way into governance. Committer status belongs to individuals, not companies. We reward contribution, maintenance, and project stewardship.
|
||||
|
||||
## Project Maintainers
|
||||
|
||||
### Lead Maintainers
|
||||
|
||||
Lead maintainers are responsible for the overall direction and strategy of the project:
|
||||
|
||||
- [@Gaohan123](https://github.com/Gaohan123)
|
||||
- [@hsliuustc0106](https://github.com/hsliuustc0106)
|
||||
- [@ywang96](https://github.com/ywang96)
|
||||
|
||||
### Active Committers
|
||||
|
||||
Committers have write access and merge rights. They typically have deep expertise in specific areas of this project and shepherd the community contributions:
|
||||
|
||||
- [@david6666666](https://github.com/david6666666): Quantization and Community Relationship
|
||||
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
|
||||
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
|
||||
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
|
||||
- [@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
|
||||
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
|
||||
|
||||
## Meetings
|
||||
|
||||
Committers hold **bi-weekly meetings** to discuss future directions and collaborations of the project.
|
||||
|
||||
## Committer Nomination Process
|
||||
|
||||
Every month, any active committer can nominate new committer(s) to the project. Up to **two new committers** will be admitted per month based on the quality and impact of their contributions.
|
||||
@@ -2,7 +2,7 @@
|
||||
|
||||
This section lists the most common options for running vLLM-Omni.
|
||||
|
||||
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.16.0/configuration/index.html)
|
||||
For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html)
|
||||
|
||||
Currently, the main options are maintained by stage configs for each model.
|
||||
|
||||
@@ -14,10 +14,8 @@ For introduction, please check [Introduction for stage config](./stage_configs.m
|
||||
|
||||
- **[GPU Memory Calculation and Configuration](./gpu_memory_utilization.md)** - Guide on how to calculate memory requirements and set up `gpu_memory_utilization` for optimal performance
|
||||
|
||||
## Multi-Stage Recipes
|
||||
|
||||
- **[Prefill-Decode Disaggregation](./pd_disaggregation.md)** - How to derive a PD-aware Qwen3-Omni stage config from the default config without introducing another bundled YAML
|
||||
|
||||
## Optimization Features
|
||||
|
||||
- **[Diffusion Features Overview](../user_guide/diffusion_features.md)** - Complete overview of all diffusion model features and supported models
|
||||
- **[TeaCache Configuration](../user_guide/diffusion/teacache.md)** - Enable TeaCache adaptive caching for DiT models to achieve 1.5x-2.0x speedup with minimal quality loss
|
||||
- **[Cache-DiT Configuration](../user_guide/diffusion/cache_dit_acceleration.md)** - Enable Cache-DiT as cache acceleration backends for DiT models
|
||||
- **[Parallelism Configuration](../user_guide/diffusion/parallelism_acceleration.md)** - Enable parallelism (e.g., sequence parallelism) for for DiT models
|
||||
@@ -158,7 +158,7 @@ The stage initialized but ran out of memory during processing.
|
||||
|
||||
**Solutions:**
|
||||
1. Reduce `max_num_batched_tokens`
|
||||
2. Reduce `max_num_seqs` in engine_args
|
||||
2. Reduce `max_batch_size` in runtime config
|
||||
3. Lower `gpu_memory_utilization` slightly
|
||||
4. Enable quantization if supported
|
||||
|
||||
|
||||
@@ -1,177 +0,0 @@
|
||||
# Prefill-Decode (PD) Disaggregation
|
||||
|
||||
PD disaggregation splits the Qwen3-Omni thinker into separate prefill and decode
|
||||
stages so prompt processing and token generation can run on different workers.
|
||||
|
||||
This is documented as a stage-config recipe instead of a bundled YAML because the
|
||||
deployment-specific values usually change per environment:
|
||||
|
||||
- GPU placement
|
||||
- `tensor_parallel_size`
|
||||
- connector backend and connector ports
|
||||
- connector IPs or bootstrap addresses
|
||||
|
||||
Start from the [default Qwen3-Omni stage config](gh-file:vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml)
|
||||
and copy it to your own file, for example `qwen3_omni_pd.yaml`. Then apply the
|
||||
changes below.
|
||||
|
||||
## Requirements
|
||||
|
||||
- 3+ GPUs for a basic layout: prefill, decode, and talker+code2wav
|
||||
- A KV connector supported by vLLM, such as `MooncakeConnector`
|
||||
- Matching `tensor_parallel_size` on the prefill and decode thinker stages
|
||||
|
||||
## 1. Split the thinker into prefill and decode stages
|
||||
|
||||
Replace the original thinker stage with two stages:
|
||||
|
||||
```yaml
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
stage_type: llm
|
||||
is_prefill_only: true
|
||||
runtime:
|
||||
devices: "0"
|
||||
engine_args:
|
||||
max_num_seqs: 16
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_producer"
|
||||
kv_rank: 0
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25201
|
||||
final_output: false
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
stage_type: llm
|
||||
is_decode_only: true
|
||||
runtime:
|
||||
devices: "1"
|
||||
engine_args:
|
||||
max_num_seqs: 64
|
||||
model_stage: thinker
|
||||
model_arch: Qwen3OmniMoeForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
gpu_memory_utilization: 0.9
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
engine_output_type: latent
|
||||
distributed_executor_backend: "mp"
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 32768
|
||||
hf_config_name: thinker_config
|
||||
tensor_parallel_size: 1
|
||||
kv_transfer_config:
|
||||
kv_connector: "MooncakeConnector"
|
||||
kv_role: "kv_consumer"
|
||||
kv_rank: 1
|
||||
kv_parallel_size: 2
|
||||
kv_connector_extra_config:
|
||||
mooncake_bootstrap_port: 25202
|
||||
engine_input_source: [0]
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
is_comprehension: true
|
||||
default_sampling_params:
|
||||
temperature: 0.4
|
||||
top_p: 0.9
|
||||
top_k: 1
|
||||
max_tokens: 2048
|
||||
seed: 42
|
||||
detokenize: True
|
||||
repetition_penalty: 1.05
|
||||
```
|
||||
|
||||
Notes:
|
||||
|
||||
- `is_prefill_only: true` marks the thinker stage that only saves KV.
|
||||
- `is_decode_only: true` marks the thinker stage that resumes from remote KV.
|
||||
- `kv_transfer_config` is required on both stages.
|
||||
- The orchestrator forces the prefill stage to run with `max_tokens=1`, so the
|
||||
prefill side only processes the prompt and exports KV.
|
||||
|
||||
## 2. Shift the downstream stages by one index
|
||||
|
||||
After inserting the extra thinker stage, renumber the remaining stages:
|
||||
|
||||
```yaml
|
||||
- stage_id: 2
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_input_source: [1]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker
|
||||
|
||||
- stage_id: 3
|
||||
runtime:
|
||||
devices: "2"
|
||||
engine_args:
|
||||
max_num_seqs: 1
|
||||
engine_input_source: [2]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav
|
||||
```
|
||||
|
||||
Compared with the default Qwen3-Omni config:
|
||||
|
||||
- the talker becomes stage `2` instead of stage `1`
|
||||
- the code2wav stage becomes stage `3` instead of stage `2`
|
||||
- the talker now reads from decode stage `1`
|
||||
|
||||
## 3. Add runtime edges for the four-stage pipeline
|
||||
|
||||
```yaml
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
- from: 1
|
||||
to: 2
|
||||
window_size: -1
|
||||
- from: 2
|
||||
to: 3
|
||||
window_size: -1
|
||||
```
|
||||
|
||||
## 4. Launch with your custom config
|
||||
|
||||
```bash
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 \
|
||||
--stage-configs-path /path/to/qwen3_omni_pd.yaml
|
||||
```
|
||||
|
||||
## Operational Notes
|
||||
|
||||
- `MooncakeConnector` does not support heterogeneous TP sizes across the PD
|
||||
pair. Keep prefill and decode at the same `tensor_parallel_size`.
|
||||
- If the thinker requires TP=2, both thinker stages must use TP=2 and be given
|
||||
separate GPU sets, for example `"0,1"` for prefill and `"2,3"` for decode.
|
||||
- Choose connector ports and addresses that match your deployment. The values
|
||||
shown above are examples only.
|
||||
@@ -18,7 +18,7 @@ If users want to modify some part of it. The custom stage_configs file can be in
|
||||
For offline (Assume necessary dependencies have ben imported):
|
||||
```python
|
||||
model_name = "Qwen/Qwen2.5-Omni-7B"
|
||||
omni = Omni(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
omni_llm = OmniLLM(model=model_name, stage_configs_path="/path/to/custom_stage_configs.yaml")
|
||||
```
|
||||
|
||||
For online serving:
|
||||
@@ -30,15 +30,15 @@ vllm serve Qwen/Qwen2.5-Omni-7B --omni --port 8091 --stage-configs-path /path/to
|
||||
|
||||
Below is a specific example of stage_configs.yaml in Qwen2.5-omni.
|
||||
```python
|
||||
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
stage_args:
|
||||
- stage_id: 0 # mark the unique id for each stage
|
||||
runtime: # The disaggregated configuration
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1 # the batch_size for offline inference
|
||||
engine_args: # Engine arguments for a certain engine
|
||||
model_stage: thinker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration # The model implementation registered in model_executor/models/registry.py
|
||||
worker_type: ar # The specific worker used
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler # The specific scehduler used
|
||||
@@ -62,9 +62,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "1"
|
||||
max_batch_size: 3
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
max_num_seqs: 3
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
@@ -88,9 +88,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
@@ -135,14 +135,6 @@ Each stage in the `stage_args` list contains the following configuration options
|
||||
|
||||
A unique identifier for each stage in the multi-stage pipeline. Stages are numbered sequentially starting from 0, and this ID is used to reference stages in inter-stage dependencies (e.g., `engine_input_source`).
|
||||
|
||||
### `prompt_expand_func` (Optional)
|
||||
|
||||
A custom Python function hook for the LLM stage (Stage 0) that expands a single incoming prompt object into multiple prompts. This is primarily used for multi-modal Classifier-Free Guidance (CFG), where it generates the necessary companion requests (like a negative text prompt) and tags them with internal roles (e.g., `cfg_text`). This ensures the upstream LLM generates the needed contextual hidden states for both the conditional and unconditional generations simultaneously.
|
||||
|
||||
### `cfg_kv_collect_func` (Optional)
|
||||
|
||||
A custom Python function hook for downstream diffusion stages (Stage 1+) to collect, map, and process the KV caches transferred from the companion requests fired by `prompt_expand_func`. It aggregates the hidden condition states cleanly (e.g., binding them as `cfg_text_past_key_values` and `cfg_text_kv_metadata`), allowing the diffusion runtime to perform CFG smoothly without redundantly evaluating text paths on the DiT workers.
|
||||
|
||||
### `runtime`
|
||||
|
||||
Configuration for disaggregated execution of the stage, controlling how the stage is deployed and executed.
|
||||
@@ -159,9 +151,9 @@ Visible devices for this stage, specified as a string. This controls which GPU d
|
||||
|
||||
Default: `"0"`
|
||||
|
||||
#### `engine_args.max_num_seqs`
|
||||
#### `runtime.max_batch_size`
|
||||
|
||||
The maximum number of sequences for concurrent processing in this stage. For LLM stages, this controls the vLLM scheduler's maximum concurrent sequences. For all stage types, this also controls how many tasks can be batched together in the task processing loop.
|
||||
The maximum batch size for offline inference in this stage. This limits how many sequences can be processed together in a single batch during offline inference operations.
|
||||
|
||||
Default: `1`
|
||||
|
||||
|
||||
@@ -1,12 +1,12 @@
|
||||
# stage config for running qwen2.5-omni with AsyncOmniEngine + Orchestrator runtime.
|
||||
# stage config for running qwen2.5-omni with architecture of OmniLLM.
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
process: true # Run this stage in a separate process
|
||||
devices: "0" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device)
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
@@ -30,9 +30,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "1"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: talker
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: ar
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
@@ -56,9 +56,9 @@ stage_args:
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0" # Example: use a different GPU than the previous stage; use "0" if single GPU
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: code2wav
|
||||
max_num_seqs: 1
|
||||
model_arch: Qwen2_5OmniForConditionalGeneration
|
||||
worker_type: generation
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
|
||||
@@ -71,7 +71,14 @@ For additional features and advanced configurations, refer to the:
|
||||
### Testing
|
||||
|
||||
vLLM-Omni uses `pytest` to test the codebase.
|
||||
Please refer to the [test instructions](./ci/test_guide.md) for detailed testing information.
|
||||
|
||||
```bash
|
||||
# Run all tests
|
||||
pytest tests/
|
||||
|
||||
# Run tests for a single test file with detailed output
|
||||
pytest -s -v tests/test_omni_llm.py
|
||||
```
|
||||
|
||||
!!! warning
|
||||
Currently, not all unit tests pass when run on CPU platforms. If you don't have access to a GPU platform to run unit tests locally, rely on the continuous integration system to run the tests for now.
|
||||
@@ -107,7 +114,7 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
- `[CI/Build]` for build or continuous integration improvements.
|
||||
- `[Doc]` for documentation fixes and improvements.
|
||||
- `[Model]` for adding a new model or improving an existing model. Model name should appear in the title.
|
||||
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `Omni`/`AsyncOmni`, etc.)
|
||||
- `[Frontend]` For changes on the vLLM-Omni frontend (e.g., OpenAI API server, `OmniLLM` class, etc.)
|
||||
- `[Kernel]` for changes affecting CUDA kernels or other compute kernels.
|
||||
- `[Core]` for changes in the core vLLM-Omni logic (e.g., `OmniProcessor`, `OmniARScheduler`, etc.)
|
||||
- `[Hardware][Vendor]` for hardware-specific changes. Vendor name should appear in the prefix, such as [Ascend] for Ascend NPUs.
|
||||
@@ -116,9 +123,6 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
!!! note
|
||||
If the PR spans more than one category, please include all relevant prefixes.
|
||||
|
||||
### Local Test
|
||||
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
|
||||
|
||||
### Code Quality
|
||||
|
||||
The PR needs to meet the following code quality standards:
|
||||
|
||||
@@ -1,6 +0,0 @@
|
||||
nav:
|
||||
- CI_5levels.md
|
||||
- failures.md
|
||||
- test_guide.md
|
||||
- test_markers.md
|
||||
- test_style.md
|
||||
@@ -67,8 +67,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
|
||||
<td><30min</td>
|
||||
<td>
|
||||
/tests/e2e/online_serving/test_{model_name}.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}.py
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}_expansion.py
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
|
||||
@@ -86,7 +86,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
|
||||
<strong>Performance:</strong><br>
|
||||
/tests/dfx/perf/tests/test.json<br>
|
||||
/tests/e2e/perf/nightly.json<br>
|
||||
<strong>Doc Test:</strong><br>
|
||||
tests/example/online_serving/test_{model_name}.py<br>
|
||||
tests/example/offline_inference/test_{model_name}.py
|
||||
@@ -104,7 +104,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td> Depends on reality </td>
|
||||
<td>
|
||||
<strong>Stability:</strong><br>
|
||||
/tests/dfx/stability/tests/test.json<br>
|
||||
tests/e2e/stability/weekly.json<br>
|
||||
<strong>Reliability:</strong><br>
|
||||
tests/e2e/reliability/test_{model_name}.py
|
||||
</td>
|
||||
@@ -168,6 +168,12 @@ vllm_omni/ tests/
|
||||
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
|
||||
│
|
||||
├── entrypoints/ → ├── entrypoints/
|
||||
│ ├── omni.py │ ├── test_omni.py ⬜ (E2E covered by e2e/offline, e2e/online)
|
||||
│ ├── omni_llm.py │ ├── test_omni_llm.py ✅
|
||||
│ ├── omni_stage.py │ ├── test_omni_stage.py ⬜ (partial in test_omni_stage_diffusion_config.py)
|
||||
│ ├── omni_diffusion.py │ ├── test_omni_diffusion.py ✅
|
||||
│ ├── async_omni.py │ ├── test_async_omni.py ✅ actually in e2e/online_serving/test_async_omni.py
|
||||
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
|
||||
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
|
||||
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
|
||||
│ │ └── ... │ │ └── test_*.py ⬜
|
||||
@@ -224,11 +230,9 @@ vllm_omni/ tests/
|
||||
│
|
||||
└── e2e (tests side only) → ├── e2e/
|
||||
├── online_serving/ ✅ non-empty
|
||||
│ ├── test_qwen2_5_omni.py
|
||||
│ ├── test_async_omni.py
|
||||
│ ├── test_qwen3_omni.py
|
||||
│ ├── test_qwen3_omni_expansion.py
|
||||
│ ├── test_mimo_audio.py
|
||||
│ ├── test_image_gen_edit.py
|
||||
│ ├── test_images_generations_lora.py
|
||||
│ └── stage_configs/
|
||||
@@ -271,7 +275,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 .
|
||||
|
||||
|
||||
|
||||
@@ -504,9 +508,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
**2.4.5 Request Execution**
|
||||
|
||||
```python
|
||||
openai_client.send_omni_request(request_config, request_num=1) # for omni-understanding models
|
||||
# or
|
||||
openai_client.send_diffusion_request(request_config, request_num=1) # for diffusion models
|
||||
openai_client.send_request(request_config, request_num=1)
|
||||
```
|
||||
|
||||
**Explanation**:
|
||||
@@ -515,7 +517,7 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Single Request**: The comment clearly states this is a single-request completion test. For concurrent testing, it can be extended to multiple requests using request_num = n.
|
||||
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
|
||||
**Implicit Validation**: The send_request method internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model performs deep validation.
|
||||
|
||||
- ***Run Command***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
|
||||
@@ -530,13 +532,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
|
||||
### 3.2 Testing Content and Scope
|
||||
|
||||
- ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling.
|
||||
- ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
|
||||
- ***Performance Testing***: Uses the `/tests/e2e/perf/nightly.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization.
|
||||
- ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description.
|
||||
|
||||
### 3.3 Test Directory and Execution Files
|
||||
|
||||
- ***Functional Testing***: Same directories as L3.
|
||||
- ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json`
|
||||
- ***Performance Test Configuration***: `/tests/e2e/perf/nightly.json`
|
||||
- ***Documentation Example Tests***:
|
||||
- - `tests/example/online_serving/test_{model_name}.py`
|
||||
- `tests/example/offline_inference/test_{model_name}.py`
|
||||
@@ -546,52 +548,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e
|
||||
- ***Trigger Timing***: **`Nightly`**, automatically executed every night.
|
||||
- ***Execution Environment***: ***GPU*** server clusters to meet the resource demands of performance testing.
|
||||
- ***Script Example***:
|
||||
|
||||
??? example "Test Examples: Documentation Example Tests"
|
||||
|
||||
--8<-- "docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md"
|
||||
|
||||
??? example "Test Examples: Performance Tests"
|
||||
|
||||
--8<-- "docs/contributing/ci/test_examples/l4_performance_tests.inc.md"
|
||||
|
||||
??? example "Test Examples: Functionality Tests"
|
||||
|
||||
--8<-- "docs/contributing/ci/test_examples/l4_functionality_tests.inc.md"
|
||||
|
||||
- ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
|
||||
|
||||
## Chapter 4: L5 Level Testing - Stability and Reliability Testing
|
||||
|
||||
### 4.1 Testing Purpose
|
||||
|
||||
L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
|
||||
|
||||
|
||||
|
||||
### 4.2 Testing Content and Scope
|
||||
|
||||
- ***Long-term Stability (Stability) Testing***: Uses the `tests/dfx/stability/tests/test.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test.json`
|
||||
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
|
||||
|
||||
### 4.4 Execution Method and Example
|
||||
|
||||
- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
|
||||
- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
|
||||
- ***Script Example***:
|
||||
<details>
|
||||
<summary> Test Examples</summary>
|
||||
When you want to add L4-level performance test cases, you can refer to the following format for case addition in tests/perf/tests/test.json:
|
||||
|
||||
When you want to add L5-level stability test cases, you can refer to the following format for case addition in `tests/dfx/stability/tests/test.json`:
|
||||
|
||||
```json
|
||||
```JSON
|
||||
{
|
||||
"test_name": "test_qwen3_omni_stability",
|
||||
"test_name": "test_qwen3_omni",
|
||||
"server_params": {
|
||||
"model": "Qwen/Qwen3-Omni-30B-A3B-Instruct",
|
||||
"stage_config_name": "qwen3_omni.yaml"
|
||||
@@ -599,15 +562,17 @@ When you want to add L5-level stability test cases, you can refer to the followi
|
||||
"benchmark_params": [
|
||||
{
|
||||
"dataset_name": "random",
|
||||
"backend": "openai-chat-omni",
|
||||
"endpoint": "/v1/chat/completions",
|
||||
"duration_sec": 43200,
|
||||
"request_rate": 0.5,
|
||||
"num_prompts_per_batch": 20,
|
||||
"num_prompts": [10, 20],
|
||||
"request_rate": [0.5, 1],
|
||||
"random_input_len": 2500,
|
||||
"random_output_len": 900,
|
||||
"ignore_eos": true,
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration"
|
||||
"percentile-metrics": "ttft,tpot,itl,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"mean_ttft_ms": 100000,
|
||||
"mean_audio_ttfp_ms": 100000,
|
||||
"mean_audio_rtf": 100000
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
@@ -617,20 +582,20 @@ When you want to add L5-level stability test cases, you can refer to the followi
|
||||
|
||||
***Overview***
|
||||
|
||||
| Field | Required | Description |
|
||||
| ---------------- | -------- | --------------------------------------------------------------------------- |
|
||||
| test_name | Yes | Unique identifier for the stability test case |
|
||||
| server_params | Yes | Server-side configuration parameters (model, stage configuration, etc.) |
|
||||
| benchmark_params | Yes | Stability benchmark running parameters (supports multiple configurations) |
|
||||
| Field | Required | Description |
|
||||
| ---------------- | -------- | --------------------------------------------------------------- |
|
||||
| test_name | Yes | Unique identifier for the test case |
|
||||
| server_params | Yes | Server-side configuration parameters |
|
||||
| benchmark_params | Yes | Benchmark running parameters (supports multiple configurations) |
|
||||
|
||||
#### server_params Configuration
|
||||
|
||||
##### Basic Parameters
|
||||
|
||||
| Parameter | Required | Example | Description |
|
||||
| ----------------- | -------- | ---------------------------------- | ----------------------------------- |
|
||||
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
|
||||
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
|
||||
| Parameter | Required | Example | Description |
|
||||
| ----------------- | -------- | ---------------------------------- | ----------------------------- |
|
||||
| model | Yes | "Qwen/Qwen3-Omni-30B-A3B-Instruct" | Model name or path |
|
||||
| stage_config_name | Yes | "qwen3_omni.yaml" | Stage configuration file name |
|
||||
|
||||
##### Dynamic Configuration (update/delete)
|
||||
|
||||
@@ -642,23 +607,70 @@ Supports incremental modifications based on the basic configuration:
|
||||
| delete | Delete specified configuration items |
|
||||
|
||||
***Example***:
|
||||
You can refer to Test Examples in Chapter 3.4
|
||||
|
||||
```
|
||||
"update": {
|
||||
"async_chunk": true, // Enable asynchronous chunk processing
|
||||
"stage_args": {
|
||||
"0": {
|
||||
"engine_args.custom_process_next_stage_input_func": "vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker_async_chunk"
|
||||
}
|
||||
}
|
||||
},
|
||||
"delete": {
|
||||
"stage_args": {
|
||||
"2": ["custom_process_input_func"] // Delete this configuration for stage 2
|
||||
}
|
||||
}
|
||||
```
|
||||
#### benchmark_params Configuration
|
||||
|
||||
For stability testing, the key parameters are:
|
||||
You can add any benchmark running parameters you need here. For all optional parameters, refer to the [benchmark documentation](https://github.com/vllm-project/vllm-omni/blob/main/docs/cli/bench/serve.md). General modifications are as follows:
|
||||
|
||||
- **duration_sec**: Total duration (in seconds) during which benchmark traffic is sent. The stability benchmark will keep sending batches until this duration is reached.
|
||||
- **request_rate** / **max_concurrency**: Exactly one of them must be specified. They control how the traffic is generated for each batch:
|
||||
- `request_rate`: Number of requests per second. The benchmark will send `num_prompts_per_batch` requests at the given rate.
|
||||
- `max_concurrency`: Maximum number of concurrent requests. When this is used, `request_rate` is set to `inf` internally.
|
||||
- **num_prompts_per_batch**: Number of prompts sent in each batch. Multiple batches will be executed sequentially within `duration_sec`.
|
||||
|
||||
All other optional parameters follow the same rules as the in Chapter 3.4.
|
||||
1. Change the ---xxx-xx-xx running parameters to xxx_xx_xx format and fill them as keys in the JSON file.
|
||||
2. For boolean variables in the running parameters, modify them to forms such as ignore_eos: true/false and fill them into the JSON file.
|
||||
3. Add the baseline parameter to specify the required validation values, ensuring the validation metric names match those in the result.json generated by the benchmark.
|
||||
4. The qps and concurrency modes are mutually exclusive. For detailed explanations, see the table below:
|
||||
|
||||
| Parameter | Type | Required | Example/Values | Description |
|
||||
| --------------- | ----------- | -------- | --------------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
|
||||
| num_prompts | int / array | Yes | 10,[10, 20, 30] | Number of requests. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of qps or max_concurrency, e.g., [10,10,10]. If an array is used, its length must match the number of qps or max_concurrency. |
|
||||
| request_rate | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||
| max_concurrency | int / array | No | 1, [1, 2, 3] | Queries per second. Supports single values or arrays. If a single value is used, it will be automatically expanded to match the number of num_prompts, e.g., [1,1,1]. If an array is used, its length must match the number of num_prompts. |
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_{model_name}.py`
|
||||
- - ***Run Command***: (Specific commands would depend on the performance testing tool and configuration defined in `nightly.json`).
|
||||
|
||||
## Chapter 4: L5 Level Testing - Stability and Reliability Testing
|
||||
|
||||
### 4.1 Testing Purpose
|
||||
|
||||
L5 level testing focuses on the performance of model services under ***long-running*** and ***abnormal fault*** scenarios. It aims to uncover deep-seated issues that only manifest under sustained pressure or extreme conditions, such as memory leaks, resource contention, gradual performance degradation, and lack of fault tolerance mechanisms. This is the final, yet crucial, line of defense for ensuring service high availability and production environment robustness.
|
||||
|
||||
|
||||
|
||||
### 4.2 Testing Content and Scope
|
||||
|
||||
- ***Long-term Stability (Stability) Testing***: Uses the `tests/e2e/stability/weekly.json` configuration to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses `tests/e2e/reliability/test_{model_name}.py` to actively simulate various fault and abnormal scenarios, such as: dependent service interruption, abnormal input data, network flicker, hardware resource preemption, etc., to verify the system's fault tolerance, self-healing, and graceful degradation capabilities.
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***Stability Test Configuration***: `tests/e2e/stability/weekly.json`
|
||||
- ***Reliability Test Suite***: `tests/e2e/reliability/test_{model_name}.py`
|
||||
|
||||
### 4.4 Execution Method and Example
|
||||
|
||||
- ***Trigger Timing***: **`Weekly`** (weekly) or **`Days before Release`** (several days before a major release). Due to long execution times, the frequency is lower.
|
||||
- ***Execution Environment***: ***GPU*** servers, requiring a stable and exclusive testing environment.
|
||||
- ***Script Example***:
|
||||
<details>
|
||||
<summary> Test Examples</summary>
|
||||
```python
|
||||
# WIP
|
||||
```
|
||||
</details>
|
||||
|
||||
- - ***Stability***: (Execution would be driven by the configuration in `weekly.json`).
|
||||
- ***Reliability***: `pytest -s -v tests/e2e/reliability/test_{model_name}.py`
|
||||
|
||||
## Summary
|
||||
|
||||
@@ -1,106 +1,10 @@
|
||||
# CI Failures
|
||||
|
||||
## Overview of CI Checks
|
||||
What should I do when a CI job fails on my PR, but I don't think my PR caused the failure?
|
||||
|
||||
When you open a PR against vLLM-Omni, several CI checks run automatically:
|
||||
|
||||
| Check | Platform | What it does |
|
||||
| ----- | -------- | ------------ |
|
||||
| **pre-commit** | GitHub Actions | Runs linting (Ruff), formatting, spell-checking (typos), and YAML validation. |
|
||||
| **Build Wheel** | GitHub Actions | Builds Python wheels for Python 3.11 and 3.12 on Ubuntu. Skipped for docs-only or Markdown-only changes (controlled by `paths-ignore` in the workflow). |
|
||||
| **DCO** | GitHub | Verifies every commit has a `Signed-off-by` line. |
|
||||
| **docs/readthedocs.org:vllm-omni** | Read the Docs | Builds the MkDocs documentation site. |
|
||||
| **buildkite/vllm-omni** | Buildkite | Runs GPU-based tests on NVIDIA CUDA hardware (L4, H100). |
|
||||
| **buildkite/vllm-omni-amd** | Buildkite | Runs GPU-based tests on AMD ROCm hardware (MI325). |
|
||||
| **buildkite/vllm-omni-intel** | Buildkite | Runs GPU-based tests on Intel XPU hardware (Intel Arc BMG). |
|
||||
## Common Case 1:
|
||||
|
||||
## Step 1: Identify the Failing Check
|
||||
🚨 Error: The command was interrupted by a signal: signal: terminated
|
||||
|
||||
Click the **Details** link next to the failing check on your PR to open the build log. The most common failures fall into these categories:
|
||||
|
||||
### pre-commit failures
|
||||
|
||||
These are typically formatting or linting issues introduced by your PR. Fix them locally:
|
||||
|
||||
```bash
|
||||
uv pip install pre-commit
|
||||
pre-commit run --all-files
|
||||
```
|
||||
|
||||
Then commit the fixes and push.
|
||||
|
||||
### DCO failures
|
||||
|
||||
Every commit must include a `Signed-off-by` line. If you forgot, amend your commits:
|
||||
|
||||
```bash
|
||||
git commit --amend -s
|
||||
git push --force-with-lease
|
||||
```
|
||||
|
||||
For multiple commits, use an interactive rebase to add the sign-off to each one.
|
||||
|
||||
### Read the Docs failures
|
||||
|
||||
The documentation build uses MkDocs with `fail_on_warning: true`, so even a minor warning (not just errors) will cause the build to fail. To reproduce locally:
|
||||
|
||||
```bash
|
||||
uv pip install -e ".[docs]"
|
||||
mkdocs build --strict
|
||||
```
|
||||
|
||||
Common causes include broken cross-references, invalid admonition syntax, or missing files referenced by `--8<--` includes.
|
||||
|
||||
### Buildkite failures
|
||||
|
||||
Buildkite runs GPU tests in Docker containers. These are the most complex checks and can fail for reasons unrelated to your PR (infrastructure issues, flaky tests, etc.). See the sections below for how to investigate.
|
||||
|
||||
## Step 2: Check if the Failure Is a Known Issue
|
||||
|
||||
Before spending time debugging, check whether the failure already exists on the `main` branch:
|
||||
|
||||
1. **Look at the Buildkite build log** — the test name and error message are usually enough to identify the issue.
|
||||
2. **Check recent CI runs on `main`** — if the same test is failing there, the failure is not caused by your PR.
|
||||
3. **Search existing issues** — look for open issues in the [vllm-omni issue tracker](https://github.com/vllm-project/vllm-omni/issues) with the test name or error message.
|
||||
|
||||
If the failure is already tracked, leave a comment on your PR noting that the failure is pre-existing and link the issue.
|
||||
|
||||
## Step 3: Investigate the Failure
|
||||
|
||||
If the failure appears to be new, investigate whether your changes caused it.
|
||||
|
||||
### Reading Buildkite Logs
|
||||
|
||||
1. Click **Details** next to the Buildkite check on your PR.
|
||||
2. Find the failing step in the pipeline (e.g., "Diffusion Model Test", "Simple Unit Test").
|
||||
3. Expand the step to see the full test output with the traceback.
|
||||
|
||||
### Running Tests Locally
|
||||
|
||||
For instructions on running tests locally (including specific test files, functions, and markers), see the [Running Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/test_guide/#running-tests) section in the Test Guide.
|
||||
|
||||
## Step 4: Raise an Issue or Fix It
|
||||
|
||||
### If the failure is pre-existing (not caused by your PR)
|
||||
|
||||
1. **Raise a new issue** if one doesn't already exist, using the title format:
|
||||
`[CI Failure]: [job-name] - [test-path]`
|
||||
2. Include the error message, relevant log excerpts, and the commit hash where the failure occurs (e.g., "Still failing on main as of commit `abc1234`").
|
||||
3. Leave a comment on your PR linking to the issue and noting that the failure is unrelated to your changes.
|
||||
|
||||
### If the failure is caused by your PR
|
||||
|
||||
1. Fix the issue in your branch and push the update.
|
||||
2. If the fix is non-trivial, consider adding a test to prevent regression.
|
||||
|
||||
## Common Failure Patterns
|
||||
|
||||
| Symptom | Likely Cause | Fix |
|
||||
| ------- | ------------ | --- |
|
||||
| `ruff` or formatting errors | Code style violation | Run `pre-commit run --all-files` |
|
||||
| `Signed-off-by` missing | DCO check | Amend commits with `git commit --amend -s` |
|
||||
| MkDocs build warning | Broken docs reference | Run `mkdocs build --strict` locally |
|
||||
| `OOM` or `CUDA out of memory` | Test exceeds GPU memory | Check if your changes increased memory usage; use `--vae_use_slicing` / `--vae_use_tiling` for diffusion tests |
|
||||
| Import errors | Missing or changed dependency | Check `pyproject.toml` and make sure dependencies are correct |
|
||||
| Timeout (step exceeded N minutes) | Test is too slow or hangs | Profile the test; check for infinite loops or deadlocks |
|
||||
| `Agent lost` in Buildkite | Infrastructure issue (not your fault) | Re-trigger the build; comment on your PR |
|
||||
Reason: The test is terminated due to exceed the time limits. Sometimes rebuild the test will pass.
|
||||
@@ -1,49 +0,0 @@
|
||||
**Preferred Test Strategy**
|
||||
|
||||
Use one of the following patterns depending on page type:
|
||||
|
||||
- **Dynamic code-block extraction (preferred for offline docs)**
|
||||
- Extract Python/Bash code blocks from markdown AST analyzer, then execute them directly in tests.
|
||||
- Benefit: test logic stays automatically aligned with docs.
|
||||
- Basic idea: Use `ReadmeSnippet.extract_readme_snippets` to extract a list of code blocks as a global variable in file,
|
||||
use this list as `pytest.mark.parametrize` parameters, and pass each snippet item to `example_runner.run` inside the parametrized test.
|
||||
Additionally pass an `output_subfolder` argument for the 2nd-level output folder explained in **Output Directory Structure** below.
|
||||
If any extra environment variable is need for a test (e.g., the example script reads it), `example_runner.run` also accepts a 3rd `env` parameter.
|
||||
- See [tests/examples/offline_inference/test_text_to_image.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/examples/offline_inference/test_text_to_image.py) for reference implementation.
|
||||
|
||||
- **Explicit copied scripts (used by online docs for now until further update)**
|
||||
- For online serving pages, it is acceptable to copy code from docs into dedicated test functions, because only client-side, request-sending scripts are tested.
|
||||
- Benefit: dynamic extraction is overly complex: need to tell server-launch and client-request scripts.
|
||||
- Requirement: copied test code must be kept in sync with doc updates.
|
||||
|
||||
**Test Case Naming Convention**
|
||||
|
||||
- Dynamic code extraction (auto-generated internally):
|
||||
- `test_{single_function_name_matching_file_name}[h2_heading_00X]`
|
||||
- Example: `test_text_to_image[basic_usage_001]`
|
||||
- Explicit copied scripts:
|
||||
- `test_{h2_heading_00X}[{dummy_param_id_for_omni_server}]`
|
||||
- Example: `test_api_calls_001[omni_server0]`
|
||||
|
||||
**Runtime Configuration**
|
||||
|
||||
In the example code tests, do **not** reduce `num_inference_steps` just to speed up the tests unless there is a strong CI reliability reason to do otherwise.
|
||||
|
||||
**Skipping Rules**
|
||||
|
||||
You may skip examples falling in the following categories using `pytest.mark.skip` or `pytest.skip`:
|
||||
|
||||
- Gradio UI scripts
|
||||
- Scenarios that significantly overlap with existing tests and add little new coverage.
|
||||
|
||||
**Output Directory Structure**
|
||||
|
||||
Use a three-layer output structure to store output artifacts:
|
||||
|
||||
1. Root output directory
|
||||
- Auto-detected from `OUTPUT_DIR` env var or auto-generated under `/tmp`.
|
||||
2. Doc-page directory
|
||||
- Define and use a clear page-level folder name in each `test_*.py` yourself (abbreviations are acceptable, e.g., `example_offline_t2i`).
|
||||
3. Test-case directory
|
||||
- Must match the case identifier (e.g., `basic_usage_001`).
|
||||
- Auto-generated for dynamic extracted tests.
|
||||
@@ -1,46 +0,0 @@
|
||||
**Scope**
|
||||
|
||||
For diffusion models, the L4 functionality test covers all or common *diffusion features* that are supported by this model, including several [parallelism acceleration methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism_acceleration/), [CPU offloading](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cpu_offload_diffusion/), [TeaCache](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/teacache/) and [Cache-DiT](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/cache_dit_acceleration/) cache backends, [quantization methods](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/quantization/overview/).
|
||||
|
||||
**Test Case Design**
|
||||
|
||||
For a *high priority* model (currently listed in [issue #1832](https://github.com/vllm-project/vllm-omni/issues/1832)), we use several test cases, each with multiple features turned on, so that each supported feature is tested in at least one test case. This is to relieve the GPU workload on the CI machine. The suggested test case combination is as follows:
|
||||
|
||||
- If the model can fit into 4 L4 GPU (with quantization and tensor parallel always on) (20GB RAM each)
|
||||
- (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
|
||||
- (4 GPUs) CacheDiT + Ulysses=2 + TP=2 + VAE=2 + FP8
|
||||
- (4 GPUs) CacheDiT + Ring=2 + HSDP=2 + VAE=2 + GGUF
|
||||
- (4 GPUs) TeaCache + CFG=2 + TP=2 + VAE=2 + FP8
|
||||
- Otherwise, consider 2 H100 GPU environment (80GB RAM each) with the following tests
|
||||
- (1 GPU) TeaCache + Layerwise CPU offloading + GGUF
|
||||
- (2 GPUs) CacheDiT + Ulysses=2 + FP8
|
||||
- (2 GPUs) CacheDiT + Ring=2 + GGUF
|
||||
- (2 GPUs) TeaCache + CFG=2 + FP8
|
||||
- (2 GPUs) CacheDiT + TP=2 + VAE=2 + FP8
|
||||
- (2 GPUs) CacheDiT + HSDP=2 + VAE=2 + GGUF
|
||||
- If 2 H100 GPU cannot handle the model either (e.g., HunyuanImage 3.0)
|
||||
- Still design tests and feature combinations that can best fit real-world scenario.
|
||||
- Do not include it in CI (or exclude it from the CI's filtering criteria). Instead, relevant PR authors are suggested to run these tests locally.
|
||||
- Fallback plan
|
||||
- If the model does not support layerwise CPU offloading, replace the corresponding test case with module-wise offloading
|
||||
- If the model only supports specific or no caching feature, use this option or remove this option in all test cases.
|
||||
- If the model only supports specific or no quantization feature, use this option or remove this option in all test cases.
|
||||
- If the model does not support certain other features, remove this option from that test case. If, consequently, the coverage of this modified test case completely overlaps with others, remove this test case.
|
||||
|
||||
For a *normal priority* model, further reduce the number of test cases.
|
||||
|
||||
- Only write one or two test cases for the most common feature combinations.
|
||||
- The author can explore themselves to see which feature combination balances output quality and performance. Alternatively, the author can refer to any example code in the PR that adds the model, or the example code in the PR that adds a feature (if the code involves this model of interest).
|
||||
|
||||
Currently all the features are available in online serving mode. Hence, only need to add `tests/e2e/online_serving/test_{model}_expansion.py`.
|
||||
|
||||
**Code Style**
|
||||
|
||||
- Validation: test that the multimodal output files of your model have the correct shapes. `OpenAIClientHandler.send_diffusion_request` should have taken care of this.
|
||||
- Test marks: always add `advanced_model` and `diffusion`. Add GPU-related marks if needed. Ref: [Markers for Tests](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_markers/).
|
||||
- To maximize code reuse, you may refer to
|
||||
- `tests/conftest.py` for `omni_server` (running server in subprocess) and `openai_client` fixtures (sending requests and validating output), `generate_synthetic_image` and `assert_XXX_valid` helper.
|
||||
- `tests/utils.py` for `@hardware_test(...)` and `hardware_marks`.
|
||||
- [Parametrizing tests (pytest doc)](https://docs.pytest.org/en/stable/example/parametrize.html) to reuse test function implementation for different cases.
|
||||
- Doc: add a concise docstring for each test function.
|
||||
- Reference L4 test implementation: [tests/e2e/online_serving/test_qwen_image_edit_expansion.py](https://github.com/vllm-project/vllm-omni/blob/main/tests/e2e/online_serving/test_qwen_image_edit_expansion.py).
|
||||
Loaded 100 of 945 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user