Compare commits

...
Author SHA1 Message Date
Joshna-Medisetty 191b9a8dbf Voxtral TTS: drop hardcoded CUDA in audio tokenizer; add XPU stage config (#2428)
Signed-off-by: Joshna Medisetty <joshna.medisetty@intel.com>
Signed-off-by: Joshna-Medisetty <joshna.medisetty@intel.com>
2026-04-04 07:04:04 +00:00
Zhou Taichang 2804a85e58 [Rebase] Rebase to vllm v0.19.0 (#2475) 2026-04-04 15:03:32 +08:00
Yueqian Lin 4c031580cf [CosyVoice3] Add online serving support, fix stage config, and add CI tests (#2431)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-04-04 14:55:15 +13:00
+11 f50c5a413f [Qwen3TTS] [TTS] [Feat] Refactor voice cache manager (#2108)
Signed-off-by: JuanPZuluaga <juanz9312@gmal.com>
Signed-off-by: yiliu30 <yi4.liu@intel.com>
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
Signed-off-by: Binh Tang <tangbinhna@gmail.com>
Signed-off-by: Binh Tang <binht@netflix.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
Signed-off-by: vraiti <vraiti@redhat.com>
Signed-off-by: Songrui625 <songrui625@gmail.com>
Signed-off-by: Lidang Jiang <lidangjiang@gmail.com>
Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>
Signed-off-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: JuanPZuluaga <juanz9312@gmal.com>
Co-authored-by: Yi Liu <yi4.liu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: Binh Tang <tangbinhna@gmail.com>
Co-authored-by: Binh Tang <binht@netflix.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: rein yang <73573651+R2-Y@users.noreply.github.com>
Co-authored-by: zhumingjue138 <zhumingjue@huawei.com>
Co-authored-by: ChenWenjing <54166744+Shirley125@users.noreply.github.com>
Co-authored-by: vraiti <vraiti@redhat.com>
Co-authored-by: 汪志鹏 <wangzhipeng628@gmail.com>
Co-authored-by: Sy03 <1370724210@qq.com>
Co-authored-by: chickeyton <ngton2014@gmail.com>
Co-authored-by: Jerry Song <46962917+Songrui625@users.noreply.github.com>
Co-authored-by: Lidang Jiang <119769478+Lidang-Jiang@users.noreply.github.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
Co-authored-by: Alex Brooks <albrooks@redhat.com>
Co-authored-by: linyueqian <linyueqian@outlook.com>
Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com>
2026-04-03 16:51:58 -04:00
Yueqian Lin 0e83ebe1d4 [OmniVoice] Add two-stage TTS serving support (#2463)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-04-03 20:41:20 +00:00
Alex Brooks 10db95f9a9 [Bugfix] Fix Flux2 Dev Guidance (#2433)
Signed-off-by: Alex Brooks <albrooks@redhat.com>
2026-04-04 00:05:19 +08:00
Lidang JiangandClaude Opus 4.6 515d15ef87 [Cleanup] Replace bare print() with logger and use specific exception types (#2228)
Signed-off-by: Lidang Jiang <lidangjiang@gmail.com>
Signed-off-by: Lidang-Jiang <lidangjiang@gmail.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
2026-04-03 23:59:07 +08:00
Jerry Song cd71567b06 [Model] Add two stages inference for model LTX-2 distilled. (#2260)
Signed-off-by: Songrui625 <songrui625@gmail.com>
2026-04-03 23:17:49 +08:00
Canlin Guo 6dc61c9a20 [Perf] Skip Wan2.2 cross attn Ulysses SP (#2459)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-04-03 16:23:36 +08:00
chickeyton 563f73b78a Refactor StageDiffusionClient and StageEngineCoreClient (#2006) 2026-04-03 08:05:19 +00:00
Sy03 7fb86d51cd [BugFix][FishSpeech] Fix structured voice clone prefill conditioning (#2446) 2026-04-03 06:56:33 +00:00
汪志鹏 fa275fd1be [BAGEL] [Feature]: Add thinking mode in Bagel multi-stage serving (#2447) 2026-04-03 11:47:57 +08:00
Canlin Guo 6be5d05a7c [Docs] Update docs to use vllm-ascend v0.18.0rc1 (#2453)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-04-03 09:54:47 +08:00
vraitiandCanlin Guo 6211413677 Update MRoPE config fallback logic (#2278)
Signed-off-by: vraiti <vraiti@redhat.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
2026-04-02 16:28:53 +00:00
ChenWenjingandHongsheng Liu 728cf6d023 [Feature] add session based audio streaming input (#2208)
Signed-off-by: CHEN <116010019@link.cuhk.edu.cn>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-04-02 22:51:41 +08:00
zhumingjue138 50bb47a629 [Test] Skip tests/e2e/online_serving/test_zimage_expansion.py due to issue #2435 (#2454) 2026-04-02 20:32:57 +08:00
rein yang ca02351a1e [skip ci][Bugfix] clean useless log (#2450)
Signed-off-by: Rein Yang <ruiruyang2@gmail.com>
2026-04-02 18:07:22 +08:00
458f402323 Add VACE support for WAN 2.1 conditional video generation (#1885)
Signed-off-by: Binh Tang <tangbinhna@gmail.com>
Signed-off-by: Binh Tang <binht@netflix.com>
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
Co-authored-by: Binh Tang <binht@netflix.com>
Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-04-02 15:42:50 +08:00
Canlin Guo e2892ef6c1 [Perf] Optimize Wan2.2 rotary embedding (#2393)
Signed-off-by: gcanlin <canlinguosdu@gmail.com>
2026-04-02 14:28:44 +08:00
Yi LiuandHongsheng Liu c1d2dcc3ac [AutoRound] Add offline quantized W4A16 model support (#1777)
Signed-off-by: yiliu30 <yi4.liu@intel.com>
Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
2026-04-02 13:57:53 +08:00
Alicia 900f6aa837 [CI] [skip ci] Rename & reset timout mins for nightly L4 tests. (#2251)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-04-02 11:45:00 +08:00
Jason d3daafbe41 [Feat] Add benchmarks for Qwen3-TTS Base/VoiceDesign Model (#2411)
Signed-off-by: Jiahui Sun <jhsun2020@gmail.com>
2026-04-02 03:44:04 +00:00
Didan Deng ebc9a8d875 [skip ci][Doc] Small fix of Doc (#2400)
Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com>
2026-04-02 10:08:58 +08:00
wangyu 9c2a576301 [Test] Add voice or language test case for Qwen3-omni and Qwen-tts (#1844)
Signed-off-by: yenuo26 <410167048@qq.com>
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
2026-04-02 09:29:22 +08:00
Yuanheng Zhao 9595be59ae [skip ci][Doc] Update RFC template doc (#2141)
Signed-off-by: Yuanheng Zhao <jonathan.zhaoyh@gmail.com>
2026-04-02 09:28:48 +08:00
NATURE bbae904f1d [Bugfix] Fix delayed decoding bug for Bagel AR/DIT workflow (L3 test_bagel_img2img error) (#2422)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-04-01 15:01:39 +00:00
Lancer 70a62651b9 [Feat] support HSDP for Qwen-image series, Z-Image, GLM-Image (#2029)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-04-01 20:01:09 +08:00
TJian 4e4bbc42a6 [CI] Tune GPU resources for test (#2401)
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
2026-04-01 19:37:34 +08:00
Roger Wang 6ef0e907af Add new committers to governance page (#2419) 2026-04-01 04:35:48 -07:00
R0CKSTAR 3def008b32 [Feat] Add MUSA platform support for Moore Threads GPUs (#2337)
Signed-off-by: Xiaodong Ye <yeahdongcn@gmail.com>
2026-04-01 17:55:00 +08:00
汪志鹏 bf5bd0a4c0 [BugFix]: Fix bagel single-stage img2img fallback to text2img bug (#2397)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-04-01 02:29:02 -07:00
Wu JIAZHENandjader 3fd4a4dc27 [Feat] Support step-boundary abort in diffusion (#1769)
Signed-off-by: jader <yjader@foxmail.com>
Signed-off-by: asukaqaq-s <1311722138@qq.com>
Co-authored-by: jader <yjader@foxmail.com>
2026-04-01 17:25:02 +08:00
NATURE d40840b114 [Perf] Bagel KV-ready early forwarding and time step consistency for /v1/chat/completions (#2398)
Signed-off-by: natureofnature <wzliu@connect.hku.hk>
2026-04-01 08:50:45 +00:00
zdobaanddaixinning 08cb436d4e Fix: ensure input tensor is contiguous in GroupCoordinator.all_gather (#2367)
Signed-off-by: daixinning <daixinning@163.com>
Co-authored-by: daixinning <daixinning@163.com>
2026-04-01 16:41:25 +08:00
Ding Zuhao c3376a466b [Feature][HunyuanImage3.0] Add cfgP to HunyuanImage3.0 (#1751)
Signed-off-by: Ding Zuhao <e1583181@u.nus.edu>
2026-04-01 16:33:02 +08:00
Alicia 080593c364 [CI] [skip ci]Nightly Report Optim (#2406)
Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-04-01 14:53:17 +08:00
WeiQing Chenanddavid6666666 183775ec36 [Docs] Update WeChat QR code for community support (#2402)
Signed-off-by: david6666666 <david6666666@users.noreply.github.com>
Co-authored-by: david6666666 <david6666666@users.noreply.github.com>
2026-04-01 12:19:47 +08:00
Samit 7b965a7b6c [Feat][Benchmark] Add synchronous video generation endpoint POST /v1/videos/sync for benchmark test (#2049)
Signed-off-by: samithuang <285365963@qq.com>
2026-04-01 10:47:33 +08:00
fattysand 7274e15840 [BugFix] qwen3_tts chunk boundary handling logic in initial chunk (IC) (#2378)
Signed-off-by: Fattysand <fattysand@users.noreply.github.com>
2026-04-01 09:40:12 +08:00
Lancer dd0b6fd4e8 [Bugfix] Fix layer-wise offload incompatibility with HSDP (#2021)
Signed-off-by: Lancer <maruixiang6688@gmail.com>
2026-04-01 09:25:34 +08:00
汪志鹏 369f301a1b [BugFix] Add BAGEL single-stage diffusion config and fix multiple <im_start><im_end> bug (#2381)
Signed-off-by: princepride <wangzhipeng628@gmail.com>
2026-04-01 09:19:13 +08:00
Kevin H. LuuandClaude Opus 4.6 f8d0bf5389 [release] Add nightly wheel release index (#2345)
Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-03-31 16:49:48 +00:00
wangyu 642f16949d [Bugfix] Update Whisper model loading to support multi-GPU configurations and optimize CUDA memory management (#2354)
Signed-off-by: wangyu <410167048@qq.com>
2026-03-31 14:42:28 +00:00
Hongsheng Liu a5bf33f25a [skip ci] update release 0.18.0 (#2380)
Signed-off-by: hsliuustc0106 <liuhongsheng4@huawei.com>
2026-03-31 20:45:07 +08:00
noobHappylife 7cb7b3e8ad Fix Qwen3-TTS gradio demo (#2372)
Signed-off-by: noobhappylife <aratar1991@hotmail.com>
2026-03-31 11:42:09 +00:00
Reid de1ac5ee88 [Bugfix] Fix case-sensitivity in Qwen3 TTS speaker name lookup (#2358)
Signed-off-by: reidliu41 <reid201711@gmail.com>
2026-03-31 09:51:35 +00:00
WeiQing Chen e68ca4aa5c [CI]test: add wan22 i2v video similarity e2e (#2262)
Signed-off-by: David Chen <530634352@qq.com>
2026-03-31 16:15:47 +08:00
WeiQing Chen 0cbd1cbeb8 [CI] fix test: use minimum supported layered output count (#2350)
Signed-off-by: David Chen <530634352@qq.com>
2026-03-31 14:59:00 +08:00
wangyuandAlicia c1a978ae8d [Test] Add qwen3-omni tests for audio_in_video and one word prompt (#2097)
Signed-off-by: yenuo26 <410167048@qq.com>
Signed-off-by: wangyu <410167048@qq.com>
Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com>
Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com>
2026-03-31 14:30:12 +08:00
Yueqian Lin 1734c1ec23 [Test] Fix 4 broken Qwen3-TTS async chunk unit tests (#2351)
Signed-off-by: linyueqian <linyueqian@outlook.com>
2026-03-31 00:45:13 -04:00
LLandCanlin Guo 12a8f57934 Fix Qwen3-TTS Base on NPU running failed (#2353)
Signed-off-by: LL <rimchew@foxmail.com>
Signed-off-by: Canlin Guo <canlinguosdu@gmail.com>
Co-authored-by: Canlin Guo <canlinguosdu@gmail.com>
2026-03-31 12:42:21 +08:00
238 changed files with 14769 additions and 2519 deletions

No files matched your search

+20
View File
@@ -0,0 +1,20 @@
steps:
- label: "Build and upload wheel"
key: "build-wheel"
agents:
queue: cpu_queue_release
commands:
- "curl -LsSf https://astral.sh/uv/install.sh | sh"
- 'export PATH="$HOME/.local/bin:$PATH"'
- "uv venv --python=3.12 && source .venv/bin/activate"
- "uv pip install --upgrade build"
- "python3 -m build"
- "bash .buildkite/scripts/upload-nightly-wheels.sh"
- label: "Generate and upload wheel indices"
depends_on: "build-wheel"
allow_dependency_failure: true
agents:
queue: small_cpu_queue_release
commands:
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
+1 -1
View File
@@ -10,7 +10,7 @@ steps:
DOCKER_BUILDKIT: "1"
# Buildkite will automatically replace this with the actual commit hash
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
VLLM_VERSION: "v0.18.0"
VLLM_VERSION: "v0.19.0"
priority: 100
timeout_in_minutes: 60
soft_fail: true
+87
View File
@@ -0,0 +1,87 @@
#!/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 -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 's/import regex as re/import re/g' .buildkite/scripts/generate-nightly-index.py
# Generate indices -- the version is just the commit hash (not omni/{commit})
# because relative paths are computed between the index and wheel directories,
# 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
echo "Re-generating indices for /omni/$pure_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 "$pure_version" \
--wheel-dir "$BUILDKITE_COMMIT" \
--current-objects "$obj_json" \
--output-dir "$INDICES_OUTPUT_DIR" \
--comment "version $pure_version"
aws s3 cp --recursive "$INDICES_OUTPUT_DIR/" "s3://$BUCKET/omni/$pure_version/"
fi
+193
View File
@@ -0,0 +1,193 @@
#!/usr/bin/env python3
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import argparse
import json
import sys
from dataclasses import asdict, dataclass
from datetime import datetime
from pathlib import Path
from typing import Any
from urllib.parse import quote
import regex as re
def normalize_package_name(name: str) -> str:
"""Normalize package name per PEP 503."""
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}")
+33
View File
@@ -0,0 +1,33 @@
#!/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."
+25 -16
View File
@@ -56,8 +56,8 @@ steps:
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_diffusion_cpu_offload.py
- pytest -s -v tests/e2e/offline_inference/test_diffusion_layerwise_offload.py
# 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
agents:
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
plugins:
@@ -111,8 +111,7 @@ steps:
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
commands:
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py
- pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
- pytest -s -v tests/e2e/offline_inference/test_sequence_parallel.py tests/diffusion/distributed/test_ulysses_uaa_perf.py
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
plugins:
@@ -174,7 +173,7 @@ steps:
pytest -s -v tests/engine/test_async_omni_engine_abort.py
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
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
@@ -193,8 +192,7 @@ steps:
commands:
- export VLLM_LOGGING_LEVEL=DEBUG
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py
- pytest -s -v tests/e2e/online_serving/test_qwen2_5_omni.py -m "advanced_model" --run-level "advanced_model"
- pytest -s -v tests/e2e/offline_inference/test_qwen2_5_omni.py 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:
@@ -216,7 +214,7 @@ steps:
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
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"
@@ -239,7 +237,7 @@ steps:
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model" && pytest -s -v tests/e2e/offline_inference/test_qwen3_tts_base.py
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"
'
agents:
queue: "gpu_1_queue"
@@ -260,9 +258,7 @@ 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
- pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "advanced_model" --run-level "advanced_model"
- 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"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -348,8 +344,7 @@ steps:
export VLLM_TEST_CLEAN_GPU_MEMORY=1
export VLLM_IMAGE_FETCH_TIMEOUT=60
pytest -s -v tests/e2e/offline_inference/test_bagel_text2img.py -m "advanced_model" --run-level "advanced_model" -k "shared_memory"
pytest -s -v tests/e2e/offline_inference/test_bagel_img2img.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/online_serving/test_bagel_online.py -m "advanced_model" --run-level "advanced_model"
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"
@@ -393,8 +388,17 @@ steps:
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
'
- label: "CosyVoice3-TTS E2E Test"
timeout_in_minutes: 20
depends_on: upload-merge-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "advanced_model" --run-level "advanced_model"
'
agents:
queue: "mithril-h100-pool"
@@ -414,6 +418,11 @@ 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:
+53 -18
View File
@@ -6,16 +6,10 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- |
set +e
pytest -s -v tests/e2e/online_serving/test_qwen3_omni_expansion.py -m "advanced_model" --run-level "advanced_model"
EXIT1=$$?
pytest -s -v tests/examples/online_serving/test_qwen3_omni.py -m "advanced_model" --run-level "advanced_model"
EXIT2=$$?
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
EXIT3=$$?
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
EXIT4=$$?
exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4))
pytest -s -v \
tests/examples/ \
tests/e2e/online_serving/test_*_expansion.py \
-m "advanced_model and H100 and omni" --run-level "advanced_model"
agents:
queue: "mithril-h100-pool"
plugins:
@@ -57,8 +51,7 @@ steps:
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
- pytest -s -v tests/examples/ -m "advanced_model and L4 and omni" --run-level "advanced_model"
- pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model"
- pytest -s -v tests/examples/ 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:
@@ -114,7 +107,7 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion Model Wan22 completed Test with H100"
- label: ":full_moon: Diffusion Model (Wan2.2) Test with H100"
timeout_in_minutes: 90
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
@@ -155,7 +148,7 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Diffusion Model Test with L4"
- label: ":full_moon: Diffusion Model Test"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
@@ -177,7 +170,7 @@ steps:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: ":full_moon: Documentation Example Code Test with H100"
- label: ":full_moon: Doc Example Code Test with H100"
timeout_in_minutes: 60
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
@@ -219,7 +212,7 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Omni Model Perf Test & Test Case Statistics"
- label: ":full_moon: Omni Model Perf Test & Testcase Statistics with H100"
key: nightly-omni-performance
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
@@ -355,9 +348,51 @@ steps:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: ":full_moon: Qwen-Image Diffusion Perf Test with H100"
- label: ":full_moon: Wan22 I2V Accuracy Test with H100"
key: nightly-wan22-i2v-accuracy
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
- pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$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 Perf Test with H100"
key: nightly-qwen-image-performance
timeout_in_minutes: 300
timeout_in_minutes: 180
depends_on: upload-nightly-pipeline
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"
commands:
+105 -42
View File
@@ -180,7 +180,7 @@ steps:
pytest -s -v tests/engine/test_async_omni_engine_abort.py
'
agents:
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
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
@@ -216,44 +216,42 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
# - label: "Omni Model Test with H100"
# depends_on: upload-ready-pipeline
# commands:
# - |
# timeout 20m bash -c '
# export VLLM_WORKER_MULTIPROC_METHOD=spawn
# export VLLM_TEST_CLEAN_GPU_MEMORY="1"
# # - pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
# pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
# '
# agents:
# queue: "mithril-h100-pool"
# plugins:
# - kubernetes:
# podSpec:
# containers:
# - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
# resources:
# limits:
# nvidia.com/gpu: 2
# volumeMounts:
# - name: devshm
# mountPath: /dev/shm
# - name: hf-cache
# mountPath: /root/.cache/huggingface
# env:
# - name: HF_HOME
# value: /root/.cache/huggingface
# nodeSelector:
# node.kubernetes.io/instance-type: gpu-h100-sxm
# volumes:
# - name: devshm
# emptyDir:
# medium: Memory
# - name: hf-cache
# hostPath:
# path: /mnt/hf-cache
# type: DirectoryOrCreate
- label: "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
@@ -273,7 +271,7 @@ steps:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 2
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
@@ -322,6 +320,29 @@ steps:
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "OmniVoice E2E Test"
timeout_in_minutes: 20
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_omnivoice.py -m "core_model" --run-level "core_model"
'
agents:
queue: "gpu_1_queue"
plugins:
- docker#v5.2.0:
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
always-pull: true
propagate-environment: true
shm-size: "8gb"
environment:
- "HF_HOME=/fsx/hf_cache"
volumes:
- "/fsx/hf_cache:/fsx/hf_cache"
- label: "Voxtral-TTS E2E Test"
timeout_in_minutes: 20
depends_on: upload-ready-pipeline
@@ -330,8 +351,7 @@ steps:
timeout 20m bash -c '
export VLLM_LOGGING_LEVEL=DEBUG
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
@@ -528,3 +548,46 @@ steps:
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
- label: "CosyVoice3-TTS E2E Test"
timeout_in_minutes: 20
depends_on: upload-ready-pipeline
commands:
- |
timeout 20m bash -c '
export VLLM_WORKER_MULTIPROC_METHOD=spawn
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts.py -m "core_model" --run-level "core_model"
'
agents:
queue: "mithril-h100-pool"
plugins:
- kubernetes:
podSpec:
containers:
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
resources:
limits:
nvidia.com/gpu: 1
volumeMounts:
- name: devshm
mountPath: /dev/shm
- name: hf-cache
mountPath: /root/.cache/huggingface
env:
- name: HF_HOME
value: /root/.cache/huggingface
- name: HF_TOKEN
valueFrom:
secretKeyRef:
name: hf-token-secret
key: token
nodeSelector:
node.kubernetes.io/instance-type: gpu-h100-sxm
volumes:
- name: devshm
emptyDir:
medium: Memory
- name: hf-cache
hostPath:
path: /mnt/hf-cache
type: DirectoryOrCreate
+2 -2
View File
@@ -7,7 +7,7 @@ body:
- type: markdown
attributes:
value: >
#### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=label%3ARFC+sort%3Aupdated-desc) for reference.
#### Please take a look at previous [RFCs](https://github.com/vllm-project/vllm-omni/issues?q=in%3Atitle%20RFC%20sort%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/12YxSsVeD1jvL-InClkeAEnZyWFDndz_65JmXvsamuV4/edit?tab=t.0#heading=h.4ef4szrsgspp).
Please provide the detailed design document of the RFC using the [template](https://docs.google.com/document/d/1jcgR3cDaUQH3VczD4ZcKaJAoYWHjCmnYzHkCNyz-9fk/edit?usp=sharing).
validations:
required: true
- type: textarea
+1 -1
View File
@@ -16,7 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
---
*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.
+3
View File
@@ -32,6 +32,9 @@ 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
+10 -2
View File
@@ -23,6 +23,9 @@
# # 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
#
@@ -35,6 +38,7 @@
# 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
@@ -53,6 +57,7 @@ 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
@@ -77,6 +82,7 @@ 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
@@ -195,7 +201,8 @@ run_bench() {
--max-concurrency ${conc_args} \
--num-warmups "${NUM_WARMUPS}" \
--config-name "${config_name}" \
--result-dir "${RESULT_DIR}"
--result-dir "${RESULT_DIR}" \
--task-type "${TASK_TYPE}"
stop_server
@@ -222,7 +229,8 @@ if [ "${RUN_HF}" = true ]; then
--num-warmups "${NUM_WARMUPS}" \
--gpu-device "${GPU_DEVICE}" \
--config-name "hf_transformers" \
--result-dir "${RESULT_DIR}"
--result-dir "${RESULT_DIR}" \
--task-type "${TASK_TYPE}"
# Allow GPU memory to settle
sleep 5
@@ -38,6 +38,10 @@ PROMPTS = [
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
@dataclass
class BenchmarkResult:
@@ -75,6 +79,29 @@ class BenchmarkResult:
per_request: list = field(default_factory=list)
def generate_audio(model, prompt: str, args):
if args.task_type == "Base":
return model.generate_voice_clone(
text=prompt,
language=args.language,
ref_audio=REF_AUDIO,
ref_text=REF_TEXT,
)
if args.task_type == "VoiceDesign":
return model.generate_voice_design(
text=prompt,
language=args.language,
instruct=INSTRUCT,
)
return model.generate_custom_voice(
text=prompt,
language=args.language,
speaker=args.voice,
)
def run_benchmark(args):
from qwen_tts import Qwen3TTSModel
@@ -95,11 +122,7 @@ def run_benchmark(args):
print(f"Warming up with {args.num_warmups} requests...")
for i in range(args.num_warmups):
p = PROMPTS[i % len(PROMPTS)]
wavs, sr = model.generate_custom_voice(
text=p,
language=args.language,
speaker=args.voice,
)
wavs, sr = generate_audio(model, p, args)
# Sync GPU
torch.cuda.synchronize(device)
print("Warmup done.")
@@ -124,11 +147,7 @@ def run_benchmark(args):
torch.cuda.synchronize(device)
st = time.perf_counter()
wavs, sr = model.generate_custom_voice(
text=prompt,
language=args.language,
speaker=args.voice,
)
wavs, sr = generate_audio(model, prompt, args)
torch.cuda.synchronize(device)
elapsed = time.perf_counter() - st
@@ -268,6 +287,7 @@ def parse_args():
parser.add_argument("--gpu-device", type=int, default=0)
parser.add_argument("--voice", type=str, default="Vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
parser.add_argument(
"--config-name", type=str, default="hf_transformers", help="Label for this config (used in filenames)"
)
@@ -37,6 +37,9 @@ PROMPTS = [
"Could you please turn down the music a little bit, I'm trying to concentrate on my work.",
"It was a dark and stormy night when the old lighthouse keeper heard a knock at the door.",
]
REF_AUDIO = "https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen3-TTS-Repo/clone_2.wav"
REF_TEXT = "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
INSTRUCT = "Speak in an incredulous tone, but with a hint of panic beginning to creep into your voice."
@dataclass
@@ -93,22 +96,39 @@ def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = 24000, sample_width
return num_samples / sample_rate
def create_payload(
prompt: str, task_type: str = "CustomVoice", voice: str = "vivian", language: str = "English"
) -> dict:
payload = {
"input": prompt,
"language": language,
"stream": True,
"response_format": "pcm",
"task_type": task_type,
}
if task_type == "Base":
payload["ref_audio"] = REF_AUDIO
payload["ref_text"] = REF_TEXT
elif task_type == "CustomVoice":
payload["voice"] = voice
elif task_type == "VoiceDesign":
payload["instructions"] = INSTRUCT
return payload
async def send_tts_request(
session: aiohttp.ClientSession,
api_url: str,
prompt: str,
task_type: str = "CustomVoice",
voice: str = "vivian",
language: str = "English",
pbar: tqdm | None = None,
) -> RequestResult:
"""Send a streaming TTS request and measure latency metrics."""
payload = {
"input": prompt,
"voice": voice,
"language": language,
"stream": True,
"response_format": "pcm",
}
payload = create_payload(prompt, task_type, voice, language)
result = RequestResult(prompt=prompt)
st = time.perf_counter()
@@ -153,6 +173,7 @@ async def run_benchmark(
num_prompts: int,
max_concurrency: int,
num_warmups: int = 3,
task_type: str = "CustomVoice",
voice: str = "vivian",
language: str = "English",
) -> BenchmarkResult:
@@ -175,7 +196,7 @@ async def run_benchmark(
warmup_tasks = []
for i in range(num_warmups):
prompt = PROMPTS[i % len(PROMPTS)]
warmup_tasks.append(send_tts_request(session, api_url, prompt, voice, language))
warmup_tasks.append(send_tts_request(session, api_url, prompt, task_type, voice, language))
await asyncio.gather(*warmup_tasks)
print(" Warmup done.")
@@ -189,7 +210,7 @@ async def run_benchmark(
async def limited_request(prompt):
async with semaphore:
return await send_tts_request(session, api_url, prompt, voice, language, pbar)
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]
@@ -306,6 +327,7 @@ async def main(args):
num_prompts=args.num_prompts,
max_concurrency=concurrency,
num_warmups=args.num_warmups,
task_type=args.task_type,
voice=args.voice,
language=args.language,
)
@@ -334,6 +356,7 @@ def parse_args():
"--max-concurrency", type=int, nargs="+", default=[1, 4, 10], help="Concurrency levels to test"
)
parser.add_argument("--num-warmups", type=int, default=3)
parser.add_argument("--task-type", type=str, default="CustomVoice", choices=["CustomVoice", "VoiceDesign", "Base"])
parser.add_argument("--voice", type=str, default="vivian")
parser.add_argument("--language", type=str, default="English")
parser.add_argument(
+2 -4
View File
@@ -1,5 +1,5 @@
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
ARG VLLM_BASE_TAG=v0.18.0
ARG VLLM_BASE_TAG=v0.19.0
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
ARG APP_DIR=/workspace/vllm-omni
WORKDIR ${APP_DIR}
@@ -11,9 +11,7 @@ RUN apt-get update && \
apt-get clean && \
rm -rf /var/lib/apt/lists/*
# Install vllm-omni into the same uv-managed Python environment used by the base image.
# Use bash -c so that $(python3 -c ...) is expanded inside the container.
RUN uv pip install --system --no-cache-dir ".[dev]"
RUN uv pip install --system ".[dev]"
RUN ln -sf /usr/bin/python3 /usr/bin/python
+13 -13
View File
@@ -1,20 +1,20 @@
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.17.0rc1
ARG VLLM_ASCEND_TAG=v0.18.0rc1
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm
RUN git fetch origin --tags && git checkout v0.18.0
# WORKDIR /vllm-workspace/vllm
# RUN git fetch origin --tags && git checkout v0.18.0
WORKDIR /vllm-workspace/vllm-ascend
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# Install vllm-ascend
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip cache purge
# WORKDIR /vllm-workspace/vllm-ascend
# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# # Install vllm-ascend
# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
# source /usr/local/Ascend/nnal/atb/set_env.sh && \
# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
# python3 -m pip cache purge
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
+13 -13
View File
@@ -1,20 +1,20 @@
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
ARG VLLM_ASCEND_TAG=v0.17.0rc1-a3
ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
WORKDIR /vllm-workspace/vllm
RUN git fetch origin --tags && git checkout v0.18.0
# WORKDIR /vllm-workspace/vllm
# RUN git fetch origin --tags && git checkout v0.18.0
WORKDIR /vllm-workspace/vllm-ascend
RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# Install vllm-ascend
# Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
source /usr/local/Ascend/nnal/atb/set_env.sh && \
export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
python3 -m pip cache purge
# WORKDIR /vllm-workspace/vllm-ascend
# RUN git fetch origin releases/v0.18.0 && git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
# # Install vllm-ascend
# # Append `libascend_hal.so` path (devlib) to LD_LIBRARY_PATH
# RUN export PIP_EXTRA_INDEX_URL=https://mirrors.huaweicloud.com/ascend/repos/pypi && \
# source /usr/local/Ascend/ascend-toolkit/set_env.sh && \
# source /usr/local/Ascend/nnal/atb/set_env.sh && \
# export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/Ascend/ascend-toolkit/latest/`uname -i`-linux/devlib && \
# python3 -m pip install -v -e /vllm-workspace/vllm-ascend/ --extra-index https://download.pytorch.org/whl/cpu/ && \
# python3 -m pip cache purge
ARG APP_DIR=/vllm-workspace/vllm-omni
WORKDIR ${APP_DIR}
+1 -1
View File
@@ -1,4 +1,4 @@
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0
FROM ${BASE_IMAGE} AS base
# Declare a variable to know if we want to use the nightly build or the stable build.
+1 -1
View File
@@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
ENV UV_INDEX_STRATEGY="unsafe-best-match"
ENV UV_LINK_MODE="copy"
ARG VLLM_VERSION=v0.18.0
ARG VLLM_VERSION=v0.19.0
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
WORKDIR /workspace/vllm
-1
View File
@@ -5,7 +5,6 @@
Main entry points for vLLM-Omni inference and serving.
- [vllm_omni.entrypoints.async_omni.AsyncOmni][]
- [vllm_omni.entrypoints.async_omni_diffusion.AsyncOmniDiffusion][]
- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][]
- [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][]
- [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][]
Binary file not shown.

Before

Width:  |  Height:  |  Size: 99 KiB

After

Width:  |  Height:  |  Size: 98 KiB

+3 -1
View File
@@ -37,10 +37,12 @@ Committers have write access and merge rights. They typically have deep expertis
- [@gcanlin](https://github.com/gcanlin): Hardware plugin and NPU integration
- [@Isotr0py](https://github.com/Isotr0py): Diffusion and Quantization
- [@linyueqian](https://github.com/linyueqian): TTS and Omni Support
- [@lishunyang12](https://github.com/lishunyang12): Quantization and Configuration
- [@princepride](https://github.com/princepride): Diffusion and Omni Support
- [@SamitHuang](https://github.com/SamitHuang): RL and Diffusion
- [@tzhouam](https://github.com/tzhouam): Engine and New Model Support
- [@wtomin](https://github.com/wtomin):
- [@wtomin](https://github.com/wtomin): Diffusion and Parallelism
- [@ZeldaHuang](https://github.com/ZeldaHuang): Omni Support
- [@ZJY0516](https://github.com/ZJY0516): Diffusion and CustomOp
## Meetings
-1
View File
@@ -168,7 +168,6 @@ vllm_omni/ tests/
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
│
├── entrypoints/ → ├── entrypoints/
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
│ │ └── ... │ │ └── test_*.py ⬜
-1
View File
@@ -73,7 +73,6 @@ vllm_omni/ tests/
│ └── arg_utils.py │ └── test_arg_utils.py ⬜
│
├── entrypoints/ → ├── entrypoints/
│ ├── async_omni_diffusion.py │ ├── test_async_omni_diffusion_config.py ✅
│ ├── stage_utils.py │ ├── test_stage_utils.py ✅
│ ├── cli/ │ ├── cli/ (benchmarks/test_serve_cli.py covers CLI serve)
│ │ └── ... │ │ └── test_*.py ⬜
@@ -69,7 +69,7 @@
[5] Orchestrator._orchestration_loop (loop)
-> poll stage output
- llm stage: await get_output_async()
- diffusion stage: get_diffusion_output_async()
- diffusion stage: get_diffusion_output_nowait()
-> (llm stage) output_processors[i].process_outputs(...)
-> _route_output(...)
-> if finished and not final_stage and non-async-chunk:
@@ -112,7 +112,7 @@ sequenceDiagram
ORCH->>S0: add_request_async
loop poll route forward
ORCH->>S0: get_output_async / get_diffusion_output_async
ORCH->>S0: get_output_async / get_diffusion_output_nowait
ORCH->>ORCH: _route_output
alt need forward to next stage
ORCH->>SN: add_request_async
+2 -2
View File
@@ -192,7 +192,7 @@ class _BaseScheduler(SchedulerInterface):
self._waiting = deque()
self._running = []
self._finished_req_ids = set()
self._max_batch_size = 1
self.max_num_running_reqs = 1
```
**Design Features**:
@@ -201,7 +201,7 @@ class _BaseScheduler(SchedulerInterface):
- **Shared cleanup logic**: Request-id registration, finish handling, and state removal are centralized instead of duplicated in each policy.
- **Current constraint**: `_max_batch_size` remains `1` because the current engine path is still synchronous request-mode execution.
- **Current constraint**: `max_num_running_reqs` remains `1` because the current engine path is still synchronous request-mode execution.
#### 2.4 Current `RequestScheduler` Policy
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
```bash
uv pip install vllm==0.18.0 --torch-backend=auto
uv pip install vllm==0.19.0 --torch-backend=auto
```
The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -66,11 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.18.0
git checkout v0.19.0
```
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
```bash
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl
#For CUDA 13.0
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl
```
Install vllm with command below (If you have no existing PyTorch).
```bash
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
vLLM-Omni is built based on vLLM. Please install it with command below.
```bash
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
```
#### Installation of vLLM-Omni
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox
# --8<-- [start:build-wheel-from-source]
#### Installation of vLLM
If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library
If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library
```bash
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
```
The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
#### Installation of vLLM-Omni
Since vllm-omni is rapidly evolving, it's recommended to install it from source
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout v0.18.0
git checkout v0.19.0
python3 -m pip install -r requirements/rocm.txt
python3 setup.py develop
```
@@ -130,7 +130,7 @@ docker run --rm \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
-p 8091:8091 \
vllm/vllm-omni-rocm:v0.18.0 \
vllm/vllm-omni-rocm:v0.19.0 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
@@ -149,7 +149,7 @@ docker run --rm -it \
-v ~/.cache/huggingface:/root/.cache/huggingface \
--env "HF_TOKEN=$HF_TOKEN" \
--entrypoint bash \
vllm/vllm-omni-rocm:v0.18.0
vllm/vllm-omni-rocm:v0.19.0
```
# --8<-- [end:pre-built-images]
@@ -10,10 +10,10 @@ The recommended way to use vLLM-Omni on NPU is through the vllm-ascend pre-built
```bash
# Update the vllm-ascend image
# Atlas A2:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
# Atlas A3:
# export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1-a3
export IMAGE=quay.io/ascend/vllm-ascend:v0.17.0rc1
# export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1-a3
export IMAGE=quay.io/ascend/vllm-ascend:v0.18.0rc1
docker run --rm \
--name vllm-omni-npu \
--shm-size=1g \
@@ -33,17 +33,6 @@ docker run --rm \
-p 8000:8000 \
-it $IMAGE bash
cd /vllm-workspace/vllm
git fetch origin --tags
git checkout v0.18.0
# Because vllm-ascend will release v0.18.0rc1 after vllm-omni 0.16.0,
# we have to pin vllm-ascend at the current commit.
cd /vllm-workspace/vllm-ascend
git pull origin main
git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
pip install -v -e .
# Inside the container, install vLLM-Omni from source
cd /vllm-workspace
git clone -b v0.18.0 https://github.com/vllm-project/vllm-omni.git
@@ -68,15 +57,10 @@ You can also build vLLM-Omni from the latest main branch if you want to use the
```bash
# Pin vLLM version to 0.18.0
cd /vllm-workspace/vllm
git fetch origin --tags
git checkout v0.18.0
git clone -b v0.18.0 https://github.com/vllm-project/vllm.git
VLLM_TARGET_DEVICE=empty pip install -v -e .
# Because vllm-ascend has not yet entered continuous development and has not been officially released, we need to pin it to a specific commit. Please note that this commit may change over time.
cd /vllm-workspace/vllm-ascend
git pull origin main
git fetch origin --tags
git checkout d781902ce9dbda8ab1e11bb0f2f0c1bc508fee7a
git clone -b v0.18.0rc1 https://github.com/vllm-project/vllm-ascend.git
pip install -v -e .
# Install vLLM-Omni from the latest main branch
+2 -2
View File
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
source .venv/bin/activate
# On CUDA
uv pip install vllm==0.18.0 --torch-backend=auto
uv pip install vllm==0.19.0 --torch-backend=auto
# On ROCm
uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700
uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700
git clone https://github.com/vllm-project/vllm-omni.git
cd vllm-omni
+1
View File
@@ -120,6 +120,7 @@ def extract_omni_serve_subparser_init():
"_FlexibleArgumentParser": _FlexibleArgumentParser,
"FlexibleArgumentParser": _FlexibleArgumentParser,
"make_arg_parser": lambda parser: parser, # no-op for doc
"_ensure_vllm_platform": lambda: None, # no-op for doc
"VLLM_SUBCMD_PARSER_EPILOG": "",
"logger": logger,
"DummySubparsers": DummySubparsers,
+3
View File
@@ -31,8 +31,11 @@ th {
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `Wan22VACEPipeline` | Wan2.1-VACE | `Wan-AI/Wan2.1-VACE-1.3B-diffusers`, `Wan-AI/Wan2.1-VACE-14B-diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
+24 -3
View File
@@ -118,6 +118,7 @@ Content-Type: application/json
| `instructions` | string | "" | Voice style/emotion instructions |
| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
**Supported languages:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
@@ -143,9 +144,23 @@ Lists available voices for the loaded model.
```json
{
"voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian"]
"voices": ["aiden", "dylan", "eric", "ono_anna", "ryan", "serena", "sohee", "uncle_fu", "vivian", "custom_voice_1"],
"uploaded_voices": [
{
"name": "custom_voice_1",
"consent": "user_consent_id",
"created_at": 1738660000,
"file_size": 1024000,
"mime_type": "audio/wav",
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
}
]
}
```
`uploaded_voices` is always present (empty list when no custom voices have been uploaded). Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
```
POST /v1/audio/voices
Content-Type: multipart/form-data
@@ -161,6 +176,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
| `consent` | string | Yes | Consent recording ID |
| `name` | string | Yes | Name for the new voice |
| `ref_text` | string | No | Transcript of the audio. When provided, enables in-context voice cloning (higher quality). Without it, only the speaker embedding is extracted. |
| `speaker_description` | string | No | Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata and returned in `GET /v1/audio/voices`. |
**Response Example:**
@@ -172,11 +188,15 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
"consent": "user_consent_id",
"created_at": 1738660000,
"mime_type": "audio/wav",
"file_size": 1024000
"file_size": 1024000,
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
}
}
```
Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
**Usage Example:**
```bash
@@ -184,7 +204,8 @@ curl -X POST http://localhost:8091/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \
-F "consent=user_consent_id" \
-F "name=custom_voice_1" \
-F "ref_text=The exact transcript of the audio sample."
-F "ref_text=The exact transcript of the audio sample." \
-F "speaker_description=warm narrator"
```
## Streaming Text Input (WebSocket)
@@ -164,7 +164,7 @@ cache_config={
**Performance Tips**:
- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., [FLUX.2-klein](https://github.com/wtomin/vllm-omni/blob/main/vllm_omni/diffusion/cache/cache_dit_backend.py#L363)) use a larger value for `Fn_compute_blocks` for a balanced performance.
- Default `Fn_compute_blocks=1` works well for most cases. Some models (e.g., FLUX.2-klein) use a larger value for `Fn_compute_blocks` for a balanced performance.
- Increase `residual_diff_threshold` (e.g., 0.12-0.15) for faster inference with slight quality trade-off, or decrease from default 0.24 for higher quality.
- Default `max_warmup_steps=4` is optimized for few-step models. Increase to 6-8 for more steps if needed.
@@ -0,0 +1,91 @@
# AutoRound Quantization
## Overview
[AutoRound](https://github.com/intel/auto-round) is an advanced quantization toolkit designed for Large Language Models (LLMs), Vision-Language Models (VLMs), and diffusion models. It achieves high accuracy at ultra-low bit widths (2–4 bits) with minimal tuning by leveraging sign-gradient descent, while providing broad hardware compatibility with multi-datatype support.
The quantization config is auto-detected from the checkpoint's `config.json` (`quantization_config.quant_method = "auto-round"`). No extra CLI flags are needed.
### Supported Schemes
| Scheme | Bits | Status |
|--------|------|--------|
| W4A16 | 4 | ✅ Supported |
| W8A16 | 8 | Planned |
W4A16 is the first supported scheme. Additional schemes will be added in future releases.
## Configuration
1. **Python API**: point `model` at a pre-quantized checkpoint. The quantization is detected automatically.
```python
from vllm_omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
omni = Omni(model="vllm-project-org/FLUX.1-dev-AutoRound-w4a16")
outputs = omni.generate(
"A cat sitting on a windowsill",
OmniDiffusionSamplingParams(num_inference_steps=28),
)
outputs[0].save_images("output.png")
```
2. **CLI**: pass the quantized model path directly.
```bash
python examples/offline_inference/text_to_image/text_to_image.py \
--model vllm-project-org/FLUX.1-dev-AutoRound-w4a16 \
--prompt "A cat sitting on a windowsill" \
--num-inference-steps 28 \
--output outputs/flux_w4a16.png
```
No `--quantization` flag is needed — the quantization method is read from the checkpoint.
## How It Works
The checkpoint's `config.json` contains:
```json
{
"quantization_config": {
"quant_method": "auto-round",
"bits": 4,
"group_size": 128,
"sym": true,
"packing_format": "auto_round:auto_gptq",
"block_name_to_quantize": "transformer_blocks,single_transformer_blocks"
}
}
```
At load time:
1. `TransformerConfig.from_dict()` parses the `quantization_config` section and builds a vLLM `INCConfig` via `build_quant_config("auto-round", ...)`.
2. `OmniDiffusionConfig.set_tf_model_config()` propagates the detected config to the engine.
3. The appropriate compute kernel (e.g. GPTQ-Marlin for W4A16) is selected automatically based on the checkpoint's bit-width and packing format.
## Supported Models
| Model | HF Checkpoint | Scheme | Group Size | Backend |
|-------|--------------|--------|------------|---------|
| FLUX.1-dev | `vllm-project-org/FLUX.1-dev-AutoRound-w4a16` | W4A16 | 128 | GPTQ-Marlin |
## Creating a Quantized Checkpoint
Use [AutoRound](https://github.com/intel/auto-round) to quantize a BF16 model. The `--scheme` flag selects the quantization scheme:
```bash
# W4A16 (4-bit weight, 16-bit activation)
auto-round \
--model black-forest-labs/FLUX.1-dev \
--scheme W4A16 \
--batch_size 1 \
--disable_opt_rtn \
--dataset coco2014 \
--iters 0
```
The output directory can be used directly as the `model` argument. See the [AutoRound documentation](https://github.com/intel/auto-round) for all available schemes and options.
@@ -11,6 +11,7 @@ vLLM-Omni provides a unified quantization framework that supports both diffusion
| FP8 | [FP8](fp8.md) | FP8 W8A8, dynamic or static | Z-Image, Qwen-Image, Flux, Bagel | SM 89 (Ada) |
| Int8 | [Int8](int8.md) | Int8 W8A8 | Z-Image, Qwen-Image | SM 89 (Ada) / Ascend NPU |
| GGUF | [GGUF](gguf.md) | GGUF format, dequant+GEMM for N-D tensors | Z-Image, Flux | SM 60 |
| AutoRound | [AutoRound](autoround.md) | W4A16 (pre-quantized) | Flux | SM 80 (Ampere) |
### Multi-stage Omni Models (Pre-quantized Checkpoints)
@@ -102,6 +103,9 @@ config = build_quant_config("fp8")
# Dict with parameters
config = build_quant_config({"method": "fp8", "activation_scheme": "static"})
# AutoRound / INC (auto-detected from checkpoint, or explicit)
config = build_quant_config("auto-round", bits=4, group_size=128)
# Per-component dict
config = build_quant_config({
"transformer": {"method": "fp8"},
+65 -47
View File
@@ -15,6 +15,7 @@ vLLM-Omni supports various advanced features for diffusion models:
- Acceleration: **cache methods**, **parallelism methods**
- Memory optimization: **cpu offloading**, **quantization**
- Extensions: **LoRA inference**
- Execution modes: **step execution**
## Supported Features
@@ -64,6 +65,16 @@ Extension methods add specialized capabilities to diffusion models beyond standa
| **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions |
### Execution Modes
Execution modes control how the diffusion pipeline processes denoise steps.
| Method | Description | Best For |
|--------|-------------|----------|
| **[Step Execution](diffusion/step_execution.md)** | Per-step denoise execution with mid-request abort support | Request cancellation between denoise steps, fine-grained execution control |
**Note:** Step execution is currently supported by QwenImagePipeline only. See [Supported Models](#supported-models) for details.
### Quantization Methods
| Method | Configuration | Description | Best For |
@@ -87,28 +98,28 @@ The following tables show which models support each feature:
### ImageGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ❌ | ✅ | ✅ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ |
| **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **FLUX.2-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **GLM-Image** | ❌ | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **HunyuanImage3** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
| **LongCat-Image** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **LongCat-Image-Edit** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **MammothModa2(T2I)** | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Nextstep_1(T2I)** | ❓ | ❓ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **OmniGen2** | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Ovis-Image** | ❌ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Qwen-Image** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-2512** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **Qwen-Image-Edit** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ✅ | ❌ | ✅ | ✅ | ❌ |
> Notes:
> 1. Nextstep_1(T2I) does not support cache acceleration methods such as TeaCache or Cache-DiT.
@@ -116,19 +127,20 @@ The following tables show which models support each feature:
### VideoGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ |
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
| **Helios** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ |
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
### AudioGen
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|
| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ |
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:|
| **Stable-Audio-Open** | ❌ | ❌ | ❓ | ❓ | ❌ | ❌ | ❌ | ❌ | ✅ | ❌ |
## Feature Compatibility
@@ -139,21 +151,22 @@ The following tables show which models support each feature:
- ❌: No support plan
- ❓: Not verified yet and Not Recommended
| | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference |
|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| **⚡TeaCache** | | | | | | | | | | | | | |
| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | |
| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | |
| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | |
| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | |
| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | |
| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | |
| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | |
| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | |
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | |
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | |
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | |
| | ⚡TeaCache | ⚡Cache-DiT | 🔀Ulysses-SP | 🔀Ring-Attn | 🔀CFG-Parallel | 🔀Tensor Parallel | 🔀HSDP | 🔀Expert Parallel | 💾CPU Offloading (Layerwise) | 💾CPU Offloading (Module-wise) | 💾VAE Patch Parallel | 💾FP8 Quant | 🔧LoRA Inference | 🔄Step Execution |
|---|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|:---:|
| **⚡TeaCache** | | | | | | | | | | | | | | |
| **⚡Cache-DiT** | ❌ | | | | | | | | | | | | | |
| **🔀Ulysses-SP** | ✅ | ✅ | | | | | | | | | | | | |
| **🔀Ring-Attn** | ✅ | ✅ | ✅ | | | | | | | | | | | |
| **🔀CFG-Parallel** | ✅ | ✅ | ✅ | ✅ | | | | | | | | | | |
| **🔀Tensor Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | | | | | | | | | |
| **🔀HSDP** | ❓ | ❓ | ❓ | ❓ | ❓ | ❌ | | | | | | | | |
| **🔀Expert Parallel** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | | | | | | |
| **💾CPU Offloading (Layerwise)** | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | | | | | | |
| **💾CPU Offloading (Module-wise)** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ❌ | | | | | |
| **💾VAE Patch Parallel** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | | | | |
| **💾FP8 Quant** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ✅ | ✅ | | | |
| **🔧LoRA Inference** | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | ❓ | | |
| **🔄Step Execution** | ❌ | ❌ | ✅ | ✅ | ✅ | ✅ | ❓ | ❓ | ✅ | ❓ | ✅ | ✅ | ❌ | |
!!! info
@@ -162,6 +175,7 @@ The following tables show which models support each feature:
3. CPU Offloading (Layerwise) and CPU Offloading (Module-wise) are not compatible.
4. CPU Offloading (Layerwise) supports single-card for now.
5. Using FP8-Quant as an example of qunatization methods.
6. Step Execution is not compatible with cache backends (TeaCache, Cache-DiT) or LoRA.
## Learn More
@@ -185,6 +199,10 @@ The following tables show which models support each feature:
- **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning
**Execution Modes:**
- **[Step Execution Guide](diffusion/step_execution.md)** - Per-step denoise execution with mid-request abort support
**Advanced Topics:**
- **[Feature Compatibility](feature_compatibility.md)** - How to combine multiple features for maximum performance
@@ -37,12 +37,45 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "input_reference=@/path/to/input.png" \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "negative_prompt=low quality, blurry, static" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42" \
-o sync_i2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
@@ -159,7 +159,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"speaker": "vivian",
"voice": "vivian",
"language": "English"
}' --output output.wav
@@ -168,7 +168,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "I am so excited!",
"speaker": "vivian",
"voice": "vivian",
"instructions": "Speak with great enthusiasm"
}' --output excited.wav
@@ -185,7 +185,7 @@ client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
response = client.audio.speech.create(
model="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice",
speaker="vivian",
voice="vivian",
input="Hello, how are you?",
)
@@ -201,7 +201,7 @@ response = httpx.post(
"http://localhost:8091/v1/audio/speech",
json={
"input": "Hello, how are you?",
"speaker": "vivian",
"voice": "vivian",
"language": "English",
},
timeout=300.0,
@@ -237,12 +237,16 @@ List all available voices/speakers from the loaded model, including both built-i
"consent": "user_consent_id",
"created_at": 1738660000,
"file_size": 1024000,
"mime_type": "audio/wav"
"mime_type": "audio/wav",
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
}
]
}
```
Fields `ref_text` and `speaker_description` are omitted per-entry when not provided at upload time.
#### POST /v1/audio/voices
Upload a new voice sample for voice cloning in Base task TTS requests.
@@ -252,6 +256,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
- `consent` (required): Consent recording ID
- `name` (required): Name for the new voice
- `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter"). Stored as metadata.
**Response Example:**
```json
@@ -262,18 +267,23 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
"consent": "user_consent_id",
"created_at": 1738660000,
"mime_type": "audio/wav",
"file_size": 1024000
"file_size": 1024000,
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
}
}
```
Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
**Usage Example:**
```bash
curl -X POST http://localhost:8000/v1/audio/voices \
curl -X POST http://localhost:8091/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \
-F "consent=user_consent_id" \
-F "name=custom_voice_1" \
-F "ref_text=The exact transcript of the audio sample."
-F "ref_text=The exact transcript of the audio sample." \
-F "speaker_description=warm narrator"
```
### Endpoint
@@ -290,7 +300,7 @@ This endpoint follows the [OpenAI Audio Speech API](https://platform.openai.com/
```json
{
"input": "Text to synthesize",
"speaker": "vivian",
"voice": "vivian",
"response_format": "wav",
"task_type": "CustomVoice",
"language": "Auto",
@@ -310,7 +320,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
### Voice and language (summary)
- **Speaker**: Use the `speaker` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
- **Speaker**: Use the `voice` request field to select the speaker (e.g., `vivian`, `ryan`, `aiden`). List available speakers with `GET /v1/audio/voices`.
- **Language**: Use the `language` field for the codec language tag (`Auto`, `Chinese`, `English`, etc.). Default is `Auto` for automatic detection.
- **CustomVoice**: Requires a valid `voice` from the model’s speaker set. **VoiceDesign**: Use `instructions` to describe the voice. **Base**: Use `ref_audio` and `ref_text` for voice cloning.
@@ -322,7 +332,7 @@ Returns binary audio data with appropriate `Content-Type` header (e.g., `audio/w
| ----------------- | ------ | -------------- | ----------------------------------------------------------- |
| `input` | string | **required** | Text to synthesize |
| `model` | string | server's model | Model to use (optional, should match server if specified) |
| `speaker` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
| `voice` | string | "vivian" | Speaker name (e.g., vivian, ryan, aiden) |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `speed` | float | 1.0 | Playback speed (0.25-4.0, not supported with `stream=true`) |
@@ -357,7 +367,7 @@ curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"speaker": "vivian",
"voice": "vivian",
"language": "English",
"stream": true,
"response_format": "pcm"
@@ -37,12 +37,43 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42" \
-o sync_t2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
+47 -6
View File
@@ -2,6 +2,7 @@ import argparse
import os
from vllm_omni.inputs.data import OmniPromptType
from vllm_omni.model_executor.stage_input_processors.bagel import GEN_THINK_SYSTEM_PROMPT
def parse_args():
@@ -65,6 +66,17 @@ def parse_args():
help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
)
parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
parser.add_argument(
"--cfg-interval",
type=float,
nargs=2,
default=None,
help="CFG interval [start, end] (default: pipeline default)",
)
parser.add_argument(
"--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel"
)
parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min")
parser.add_argument(
"--enable-diffusion-pipeline-profiler",
action="store_true",
@@ -76,6 +88,12 @@ def parse_args():
default=None,
help="Quantization method (e.g. 'fp8').",
)
parser.add_argument(
"--think",
action="store_true",
default=False,
help="Enable thinking mode: AR stage decodes <think>...</think> planning tokens before image generation.",
)
args = parser.parse_args()
return args
@@ -101,7 +119,7 @@ def main():
if not prompts:
# Default prompt for text2img test if none provided
prompts = ["<|im_start|>A cute cat<|im_end|>"]
prompts = ["A cute cat"]
print(f"[Info] No prompts provided, using default: {prompts}")
omni_outputs = []
@@ -110,8 +128,12 @@ def main():
from vllm_omni.entrypoints.omni import Omni
omni_kwargs = {}
if args.stage_configs_path:
omni_kwargs["stage_configs_path"] = args.stage_configs_path
stage_configs_path = args.stage_configs_path
if args.think and stage_configs_path is None:
stage_configs_path = "vllm_omni/model_executor/stage_configs/bagel_think.yaml"
print(f"[Info] Think mode enabled, using stage config: {stage_configs_path}")
if stage_configs_path:
omni_kwargs["stage_configs_path"] = stage_configs_path
omni_kwargs.update(
{
@@ -136,7 +158,8 @@ def main():
if not args.image_path or not os.path.exists(args.image_path):
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
loaded_image = Image.open(args.image_path).convert("RGB")
final_prompt_text = f"<|fim_middle|><|im_start|>{p}<|im_end|>"
think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>"
prompt_dict = {
"prompt": final_prompt_text,
"multi_modal_data": {"img2img": loaded_image},
@@ -160,7 +183,8 @@ def main():
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
formatted_prompts.append(prompt_dict)
else:
final_prompt_text = f"<|im_start|>{p}<|im_end|>"
think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>"
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
if args.negative_prompt is not None:
prompt_dict["negative_prompt"] = args.negative_prompt
@@ -168,7 +192,6 @@ def main():
params_list = omni.default_sampling_params_list
if args.modality in ("text2img", "img2img"):
params_list[0].max_tokens = 1 # type: ignore
if len(params_list) > 1:
diffusion_params = params_list[1]
diffusion_params.num_inference_steps = args.steps # type: ignore
@@ -179,6 +202,12 @@ def main():
"cfg_text_scale": args.cfg_text_scale,
"cfg_img_scale": args.cfg_img_scale,
}
if args.cfg_interval is not None:
extra["cfg_interval"] = tuple(args.cfg_interval)
if args.cfg_renorm_type is not None:
extra["cfg_renorm_type"] = args.cfg_renorm_type
if args.cfg_renorm_min is not None:
extra["cfg_renorm_min"] = args.cfg_renorm_min
if args.negative_prompt is not None:
extra["negative_prompt"] = args.negative_prompt
diffusion_params.extra_args = extra # type: ignore
@@ -187,6 +216,17 @@ def main():
img_idx = 0
for req_output in omni_outputs:
if args.think:
text_output = getattr(req_output, "text", None) or getattr(req_output, "outputs", None)
if text_output:
if isinstance(text_output, list) and text_output:
for out in text_output:
txt = getattr(out, "text", str(out))
if txt:
print(f"[Think] {txt}")
elif isinstance(text_output, str):
print(f"[Think] {text_output}")
images = getattr(req_output, "images", None)
if not images:
@@ -195,6 +235,7 @@ def main():
for j, img in enumerate(images):
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
img.save(save_path)
print(f"[Output] Saved image to {save_path}")
img_idx += 1
print(omni_outputs)
@@ -490,10 +490,6 @@ def main():
if frames_np.ndim == 4 and frames_np.shape[-1] == 4:
frames_np = frames_np[..., :3]
frames_np = np.clip(frames_np, 0.0, 1.0)
frames_u8 = (frames_np * 255).round().clip(0, 255).astype("uint8")
video_tensor = torch.from_numpy(frames_u8)
audio_out = None
if audio is not None:
if isinstance(audio, list):
@@ -507,7 +503,7 @@ def main():
audio_out = audio_out.float().cpu()
encode_video(
video_tensor,
frames_np,
fps=fps,
audio=audio_out,
audio_sample_rate=args.audio_sample_rate if audio_out is not None else None,
@@ -0,0 +1,73 @@
# OmniVoice
This directory contains an offline demo for running OmniVoice TTS models with vLLM Omni. It generates speech from text and saves WAV files locally.
## Model Overview
[OmniVoice](https://huggingface.co/k2-fsa/OmniVoice) is a zero-shot multilingual TTS model supporting 600+ languages. It uses a diffusion language model (Qwen3-0.6B backbone) with iterative masked unmasking to generate speech.
Three inference modes are supported:
- **Auto Voice**: Generate speech without any reference — the model picks a voice automatically.
- **Voice Clone**: Clone a voice from a reference audio + transcription.
- **Voice Design**: Control voice style via natural language instruction (e.g., "female, low pitch, british accent").
## Setup
Ensure the model is downloaded:
```bash
huggingface-cli download k2-fsa/OmniVoice
```
> **Note:** Voice cloning requires `transformers>=5.3.0` for `HiggsAudioV2TokenizerModel`. Auto voice and voice design modes work with `transformers>=4.57.0`.
## Quick Start
Auto voice (text only):
```bash
python end2end.py --model k2-fsa/OmniVoice --text "Hello, this is a test."
```
Voice design (with style instruction):
```bash
python end2end.py --model k2-fsa/OmniVoice \
--text "Hello, this is a test." \
--instruct "female, low pitch, british accent"
```
Voice clone (with reference audio):
```bash
python end2end.py --model k2-fsa/OmniVoice \
--text "Hello, this is a test." \
--ref-audio ref.wav \
--ref-text "This is the reference transcription."
```
## Language Support
Specify a language for improved quality:
```bash
python end2end.py --model k2-fsa/OmniVoice \
--text "你好,这是一个测试。" \
--lang zh
```
## Architecture
OmniVoice uses a two-stage pipeline:
- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text.
- **Stage 1 (Decoder)**: HiggsAudioV2 RVQ quantizer + DAC acoustic decoder. Converts tokens to 24kHz waveform.
Both stages use `GPUGenerationWorker` with `OmniGenerationScheduler`.
## Notes
- Output audio is saved to `output.wav` by default. Use `--output` to change the path.
- The model estimates duration from text automatically via `RuleDurationEstimator`.
- Use `--stage-init-timeout` to increase the stage initialization timeout for first-time model downloads.
@@ -0,0 +1,164 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""End-to-end OmniVoice TTS inference via vLLM-Omni.
Supports:
- Auto voice mode: text only → generated speech
- Voice cloning mode: text + reference audio → cloned voice speech
Usage:
# Auto voice
python end2end.py --model k2-fsa/OmniVoice --text "Hello world"
# Voice cloning
python end2end.py --model k2-fsa/OmniVoice --text "Hello" \
--ref-audio ref.wav --ref-text "reference transcription"
"""
import argparse
import os
import numpy as np
import soundfile as sf
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
def run_e2e():
parser = argparse.ArgumentParser(description="OmniVoice E2E TTS inference")
parser.add_argument(
"--model",
type=str,
default="k2-fsa/OmniVoice",
help="Model name or path (HuggingFace or local)",
)
parser.add_argument(
"--stage-config",
type=str,
default="vllm_omni/model_executor/stage_configs/omnivoice.yaml",
)
parser.add_argument(
"--text",
type=str,
default="Hello, this is a test of the OmniVoice text to speech system.",
)
parser.add_argument(
"--ref-audio",
type=str,
default=None,
help="Reference audio for voice cloning (WAV file)",
)
parser.add_argument(
"--ref-text",
type=str,
default=None,
help="Transcription of reference audio",
)
parser.add_argument(
"--lang",
type=str,
default=None,
help="Language code (e.g., 'en', 'zh')",
)
parser.add_argument(
"--instruct",
type=str,
default=None,
help="Voice design instruction (e.g., 'female, low pitch, british accent')",
)
parser.add_argument(
"--output",
type=str,
default="output.wav",
help="Output audio file path",
)
parser.add_argument(
"--stage-init-timeout",
type=int,
default=600,
help="Stage initialization timeout in seconds",
)
args = parser.parse_args()
if not os.path.exists(args.stage_config):
raise FileNotFoundError(f"Stage config not found: {args.stage_config}")
print(f"Initializing OmniVoice with model={args.model}")
omni = Omni(
model=args.model,
stage_configs_path=args.stage_config,
trust_remote_code=True,
log_stats=True,
)
print("Model initialized. Preparing inputs...")
# Build prompt
mm_processor_kwargs = {}
multi_modal_data = {}
if args.ref_audio:
if not os.path.exists(args.ref_audio):
raise FileNotFoundError(f"Reference audio not found: {args.ref_audio}")
import librosa
audio_signal, sr = librosa.load(args.ref_audio, sr=None)
multi_modal_data["audio"] = (audio_signal.astype(np.float32), sr)
mm_processor_kwargs["ref_text"] = args.ref_text or ""
mm_processor_kwargs["sample_rate"] = sr
if args.lang:
mm_processor_kwargs["lang"] = args.lang
if args.instruct:
mm_processor_kwargs["instruct"] = args.instruct
prompts = {"prompt": args.text}
if multi_modal_data:
prompts["multi_modal_data"] = multi_modal_data
if mm_processor_kwargs:
prompts["mm_processor_kwargs"] = mm_processor_kwargs
sampling_params_list = [OmniDiffusionSamplingParams()]
print(f"Generating speech for: {args.text}")
outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
print(f"Received {len(outputs)} outputs.")
for i, output in enumerate(outputs):
try:
ro = output.request_output
if ro is None:
print("No request_output found.")
continue
mm = getattr(ro, "multimodal_output", None)
if not mm and ro.outputs:
mm = getattr(ro.outputs[0], "multimodal_output", None)
if mm:
print(f"Multimodal output keys: {mm.keys()}")
if "audio" in mm:
audio_out = mm["audio"]
sr = mm.get("sr", 24000)
if isinstance(audio_out, np.ndarray):
audio_np = audio_out
else:
audio_np = audio_out.cpu().numpy().squeeze()
out_path = args.output if i == 0 else f"output_{i}.wav"
sf.write(out_path, audio_np, sr)
print(f"Saved audio to {out_path} ({sr}Hz, {len(audio_np) / sr:.2f}s)")
else:
print("No multimodal output found.")
except Exception as e:
print(f"Error inspecting output: {e}")
omni.close()
print("Done.")
if __name__ == "__main__":
run_e2e()
@@ -56,8 +56,13 @@ def parse_args() -> argparse.Namespace:
"Examples: Wan-AI/Wan2.2-T2V-A14B-Diffusers, "
"hunyuanvideo-community/HunyuanVideo-1.5-480p_t2v",
)
parser.add_argument(
"--model-class-name",
default=None,
help="Override model class name (e.g., LTX2TwoStagesVideoPipeline).",
)
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt (Wan2.2 only).")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
parser.add_argument(
@@ -185,6 +190,7 @@ def parse_args() -> argparse.Namespace:
def main():
args = parse_args()
model_class_name = args.model_class_name
preset = _detect_preset(args.model)
for key, default_val in preset.items():
@@ -229,6 +235,7 @@ def main():
enable_cpu_offload=args.enable_cpu_offload,
parallel_config=parallel_config,
enforce_eager=args.enforce_eager,
model_class_name=model_class_name,
cache_backend=args.cache_backend,
cache_config=cache_config,
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
@@ -0,0 +1,88 @@
# VACE Video Generation
[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single model.
| Model | Architecture | Model Weights (bf16) | HuggingFace |
|-------|-------------|----------------------|-------------|
| Wan2.1-VACE (1.3B) | Wan2.1 | ~10 GB | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
| Wan2.1-VACE (14B) | Wan2.1 | ~38 GB | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
## Text-to-Video (T2V)
```bash
python vace_video_generation.py \
--mode t2v \
--prompt "A sleek robot stands in a vast warehouse filled with boxes" \
--height 480 --width 832 --num-frames 81 \
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
--output t2v_output.mp4
```
## Image-to-Video (I2V)
First frame is kept, remaining frames are generated:
```bash
python vace_video_generation.py \
--mode i2v \
--image astronaut.jpg \
--prompt "An astronaut emerging from a cracked egg on the moon" \
--height 480 --width 832 --num-frames 81 \
--output i2v_output.mp4
```
## First-Last-Frame Interpolation (FLF2V)
```bash
python vace_video_generation.py \
--mode flf2v \
--image first_frame.jpg --last-image last_frame.jpg \
--prompt "A bird takes off from a branch and lands on another" \
--height 512 --width 512 --num-frames 81 \
--output flf2v_output.mp4
```
## Inpainting
Center vertical stripe is masked and regenerated:
```bash
python vace_video_generation.py \
--mode inpaint \
--image scene.jpg \
--prompt "Shrek walks out of a building" \
--height 480 --width 832 --num-frames 81 \
--output inpaint_output.mp4
```
## Reference Image-guided (R2V)
```bash
python vace_video_generation.py \
--mode r2v \
--image reference.jpg \
--prompt "Camera slowly zooms out from the character" \
--height 480 --width 832 --num-frames 81 \
--output r2v_output.mp4
```
## Key Arguments
- `--mode`: VACE task mode (`t2v`, `i2v`, `flf2v`, `inpaint`, `r2v`).
- `--model`: Model ID (default: `Wan-AI/Wan2.1-VACE-1.3B-diffusers`).
- `--image`: Input image for I2V, inpainting, and R2V modes.
- `--last-image`: Last frame image for FLF2V mode.
- `--prompt`: Text description of desired video.
- `--height/--width`: Output resolution (default 480x832). Dimensions should be multiples of 16.
- `--num-frames`: Number of frames (default 81).
- `--guidance-scale`: CFG scale (default 5.0).
- `--flow-shift`: Scheduler flow shift (default 5.0).
- `--num-inference-steps`: Number of denoising steps (default 30).
- `--fps`: Frames per second for the saved MP4 (default 16).
- `--output`: Path to save the generated video.
- `--vae-use-tiling`: Enable VAE tiling for memory optimization.
- `--ulysses-degree`: Ulysses sequence parallelism degree for multi-GPU.
- `--cfg-parallel-size`: CFG parallel size for multi-GPU.
- `--tensor-parallel-size`: Tensor parallel size.
> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options.
@@ -0,0 +1,209 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""VACE video generation example.
VACE (Video All-in-one Creation Engine) supports multiple video tasks:
- T2V: Text-to-Video
- I2V: Image-to-Video (first frame conditioning)
- V2LF: Video-to-Last-Frame
- FLF2V: First-Last-Frame interpolation
- Inpainting: Masked region generation
- R2V: Reference image-guided generation
Usage examples:
# T2V (text-to-video)
python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
# I2V (image-to-video, first frame kept)
python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
# FLF2V (first-last frame interpolation)
python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
# R2V (reference image guided)
python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
"""
import argparse
import time
from pathlib import Path
import numpy as np
import PIL.Image
import torch
from vllm_omni.diffusion.data import DiffusionParallelConfig
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="VACE video generation.")
parser.add_argument(
"--model",
default="Wan-AI/Wan2.1-VACE-14B-diffusers",
help="VACE model ID or local path.",
)
parser.add_argument(
"--mode",
default="t2v",
choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
help="Generation mode.",
)
parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
parser.add_argument("--height", type=int, default=480, help="Video height.")
parser.add_argument("--width", type=int, default=832, help="Video width.")
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
return parser.parse_args()
def build_prompts(args):
"""Build prompt dict with multi_modal_data based on mode."""
h, w, nf = args.height, args.width, args.num_frames
gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
mask_black = PIL.Image.new("L", (w, h), 0)
mask_white = PIL.Image.new("L", (w, h), 255)
prompt_data = {
"prompt": args.prompt,
"negative_prompt": args.negative_prompt,
}
if args.mode == "t2v":
return prompt_data
if args.mode == "r2v":
assert args.image, "--image required for R2V mode"
ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
return prompt_data
if args.mode == "i2v":
assert args.image, "--image required for I2V mode"
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {
"video": [img] + [gray] * (nf - 1),
"mask": [mask_black] + [mask_white] * (nf - 1),
}
return prompt_data
if args.mode == "v2lf":
assert args.image, "--image required for V2LF mode"
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {
"video": [gray] * (nf - 1) + [img],
"mask": [mask_white] * (nf - 1) + [mask_black],
}
return prompt_data
if args.mode == "flf2v":
assert args.image and args.last_image, "--image and --last-image required for FLF2V"
first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
prompt_data["multi_modal_data"] = {
"video": [first] + [gray] * (nf - 2) + [last],
"mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
}
return prompt_data
if args.mode == "inpaint":
assert args.image, "--image required for inpaint mode"
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
d = 80
frames, masks = [], []
for _ in range(nf):
base = np.array(img).copy()
mask = PIL.Image.new("L", (w, h), 0)
stripe = PIL.Image.new("L", (2 * d, h), 255)
mask.paste(stripe, (w // 2 - d, 0))
base[np.array(mask) > 128] = 128
frames.append(PIL.Image.fromarray(base))
masks.append(mask)
prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
return prompt_data
raise ValueError(f"Unknown mode: {args.mode}")
def main():
args = parse_args()
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
parallel_config = DiffusionParallelConfig(
ulysses_degree=args.ulysses_degree,
ring_degree=args.ring_degree,
cfg_parallel_size=args.cfg_parallel_size,
)
omni = Omni(
model=args.model,
vae_use_tiling=args.vae_use_tiling,
flow_shift=args.flow_shift,
enforce_eager=args.enforce_eager,
parallel_config=parallel_config,
)
prompt_data = build_prompts(args)
print(f"\n{'=' * 60}")
print(f"VACE {args.mode.upper()} Generation")
print(f" Model: {args.model}")
print(f" Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
print(f"{'=' * 60}\n")
start = time.perf_counter()
outputs = omni.generate(
prompt_data,
OmniDiffusionSamplingParams(
height=args.height,
width=args.width,
num_frames=args.num_frames,
num_inference_steps=args.num_inference_steps,
guidance_scale=args.guidance_scale,
generator=generator,
),
)
elapsed = time.perf_counter() - start
video = outputs[0].images
if isinstance(video, list):
video = video[0]
if isinstance(video, torch.Tensor):
video = video.cpu().numpy()
if video.ndim == 5:
video = video[0]
print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
output_path = Path(args.output)
output_path.parent.mkdir(parents=True, exist_ok=True)
from diffusers.utils import export_to_video
if np.issubdtype(video.dtype, np.integer):
video = video.astype(np.float32) / 255.0
export_to_video(list(video), str(output_path), fps=args.fps)
print(f"Saved to {output_path}")
omni.close()
if __name__ == "__main__":
main()
@@ -125,7 +125,7 @@ def generate_image(
def main():
parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
parser.add_argument("--prompt", "-p", default="<|im_start|>A cute cat<|im_end|>", help="Text prompt")
parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
@@ -34,12 +34,45 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=A bear playing with yarn, smooth motion" \
-F "input_reference=@/path/to/input.png" \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "negative_prompt=low quality, blurry, static" \
-F "num_inference_steps=40" \
-F "guidance_scale=1.0" \
-F "guidance_scale_2=1.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=12.0" \
-F "seed=42" \
-o sync_i2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
+131
View File
@@ -0,0 +1,131 @@
# OmniVoice
## Model Overview
| Model | Description |
|-------|-------------|
| `k2-fsa/OmniVoice` | Zero-shot multilingual TTS (600+ languages) with diffusion language model (Qwen3-0.6B backbone) |
> **Note:** Requires `transformers>=5.3.0` for voice cloning (HiggsAudioV2 tokenizer). Auto voice and voice design work with `transformers>=4.57.0`.
## Launch the Server
```bash
vllm serve k2-fsa/OmniVoice \
--omni \
--port 8091 \
--trust-remote-code
```
Or use the convenience script:
```bash
./run_server.sh
```
## Send TTS Request
### Using curl
```bash
# Basic TTS (auto voice)
curl -X POST http://localhost:8091/v1/audio/speech \
-H "Content-Type: application/json" \
-d '{
"input": "Hello, how are you?",
"voice": "default",
"response_format": "wav"
}' --output output.wav
```
### Using Python
```python
import httpx
response = httpx.post(
"http://localhost:8091/v1/audio/speech",
json={
"input": "Hello, how are you?",
"voice": "default",
"response_format": "wav",
},
timeout=300.0,
)
with open("output.wav", "wb") as f:
f.write(response.content)
```
### Using OpenAI SDK
```python
from openai import OpenAI
client = OpenAI(base_url="http://localhost:8091/v1", api_key="none")
response = client.audio.speech.create(
model="k2-fsa/OmniVoice",
voice="default",
input="Hello, how are you?",
)
response.stream_to_file("output.wav")
```
### Using the CLI Client
```bash
cd examples/online_serving/omnivoice
# Basic TTS
python speech_client.py --text "Hello, how are you?"
# Specify language for improved quality
python speech_client.py --text "Bonjour, comment allez-vous?" --language French
```
The CLI client supports:
- `--api-base`: API base URL (default: `http://localhost:8091`)
- `--model` (or `-m`): Model name (default: `k2-fsa/OmniVoice`)
- `--text`: Text to synthesize (required)
- `--response-format`: Audio format: wav, mp3, flac, pcm, aac, opus (default: wav)
- `--language`: Language hint (default: Auto)
- `--output` (or `-o`): Output file path (default: `omnivoice_output.wav`)
## Inference Modes
OmniVoice supports three inference modes. Currently, **auto voice** is supported through the online Speech API. Voice cloning and voice design are available via offline inference (see `examples/offline_inference/omnivoice/`).
| Mode | Description | Online API | Offline |
|------|-------------|:----------:|:-------:|
| Auto Voice | Generate speech without reference | Yes | Yes |
| Voice Clone | Clone from reference audio + transcript | - | Yes |
| Voice Design | Control style via natural language instruction | - | Yes |
## Architecture
OmniVoice uses a single-stage diffusion pipeline:
- **Stage 0 (Generator)**: Qwen3-0.6B transformer with 32-step iterative masked unmasking and classifier-free guidance. Generates 8-codebook audio tokens from text, then decodes to 24kHz waveform via HiggsAudioV2 RVQ quantizer + DAC acoustic decoder.
## API Parameters
OmniVoice uses the standard `/v1/audio/speech` endpoint. See the [Speech API reference](https://docs.vllm.ai/projects/vllm-omni/en/latest/serving/speech_api/) for full documentation.
Key parameters:
| Parameter | Type | Default | Description |
|-----------|------|---------|-------------|
| `input` | string | **required** | Text to synthesize |
| `voice` | string | "default" | Voice name |
| `response_format` | string | "wav" | Audio format: wav, mp3, flac, pcm, aac, opus |
| `speed` | float | 1.0 | Playback speed (0.25-4.0) |
## Troubleshooting
1. **TTS model did not produce audio output**: Ensure the model is fully downloaded (`huggingface-cli download k2-fsa/OmniVoice`)
2. **Connection refused**: Make sure the server is running on the correct port
3. **Out of memory**: Reduce `--gpu-memory-utilization` (default stage config uses 0.5)
4. **Slow first request**: The model performs warmup on first inference; subsequent requests are faster
+19
View File
@@ -0,0 +1,19 @@
#!/bin/bash
# Launch vLLM-Omni server for OmniVoice TTS
#
# Usage:
# ./run_server.sh
# CUDA_VISIBLE_DEVICES=0 ./run_server.sh
set -e
MODEL="${MODEL:-k2-fsa/OmniVoice}"
PORT="${PORT:-8091}"
echo "Starting OmniVoice server with model: $MODEL"
vllm serve "$MODEL" \
--host 0.0.0.0 \
--port "$PORT" \
--trust-remote-code \
--omni
@@ -0,0 +1,84 @@
"""Client for OmniVoice TTS via /v1/audio/speech endpoint.
Examples:
# Basic TTS (auto voice)
python speech_client.py --text "Hello, how are you?"
# Specify language
python speech_client.py --text "Bonjour, comment allez-vous?" --language French
"""
import argparse
import httpx
DEFAULT_API_BASE = "http://localhost:8091"
DEFAULT_API_KEY = "EMPTY"
def run_tts(args) -> None:
"""Generate speech via /v1/audio/speech API."""
payload = {
"model": args.model,
"input": args.text,
"voice": "default",
"response_format": args.response_format,
}
if args.language:
payload["language"] = args.language
print(f"Model: {args.model}")
print(f"Text: {args.text}")
if args.language:
print(f"Language: {args.language}")
print("Generating audio...")
api_url = f"{args.api_base}/v1/audio/speech"
headers = {
"Content-Type": "application/json",
"Authorization": f"Bearer {args.api_key}",
}
with httpx.Client(timeout=300.0) as client:
response = client.post(api_url, json=payload, headers=headers)
if response.status_code != 200:
print(f"Error: {response.status_code}")
print(response.text)
return
try:
text = response.content.decode("utf-8")
if text.startswith('{"error"'):
print(f"Error: {text}")
return
except UnicodeDecodeError:
pass
output_path = args.output or "omnivoice_output.wav"
with open(output_path, "wb") as f:
f.write(response.content)
print(f"Audio saved to: {output_path}")
def main():
parser = argparse.ArgumentParser(description="OmniVoice TTS client")
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
parser.add_argument("--model", "-m", default="k2-fsa/OmniVoice", help="Model name")
parser.add_argument("--text", required=True, help="Text to synthesize")
parser.add_argument("--language", default=None, help="Language hint (e.g., English, Chinese, French)")
parser.add_argument(
"--response-format",
default="wav",
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
help="Audio format (default: wav)",
)
parser.add_argument("--output", "-o", default=None, help="Output file path")
args = parser.parse_args()
run_tts(args)
if __name__ == "__main__":
main()
@@ -36,6 +36,45 @@ cd examples/online_serving/qwen3_omni
python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --query-type use_image --port 8091 --host "localhost"
```
#### Realtime WebSocket client (`openai_realtime_client.py`)
[`openai_realtime_client.py`](./openai_realtime_client.py) connects to **`ws://<host>:<port>/v1/realtime`**, uploads a local audio file as **PCM16 mono @ 16 kHz** chunks (OpenAI-style `input_audio_buffer.append` / `commit`), and prints **streaming transcription** (`transcription.delta` / `transcription.done`).
**Dependencies:**
```bash
pip install websockets librosa numpy
```
(ffmpeg may be required by `librosa` for some formats; see the FAQ below.)
**From this directory** (`examples/online_serving/qwen3_omni`):
```bash
python openai_realtime_client.py \
--host localhost \
--port 8091 \
--model Qwen/Qwen3-Omni-30B-A3B-Instruct \
--audio_path /path/to/your.wav
```
If `--audio_path` is omitted, the script uses a bundled default clip (`mary_had_lamb` via vLLM assets).
**Arguments:**
| Flag | Default | Description |
|------|---------|-------------|
| `--host` | `localhost` | API server host |
| `--port` | `8000` | API server port (match your `vllm serve` port, e.g. `8091`) |
| `--model` | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | Must match the served model (also sent in `session.update`) |
| `--audio_path` | *(optional)* | Path to input audio; resampled to 16 kHz mono inside the client |
Ensure the vLLM-Omni server is running with realtime support for this endpoint, for example:
```bash
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
```
The Python client supports the following command-line arguments:
- `--query-type` (or `-q`): Query type (default: `use_video`). Options: `text`, `use_audio`, `use_image`, `use_video`
@@ -0,0 +1,146 @@
"""
This script demonstrates how to use the vLLM-Omni Realtime WebSocket API to perform
audio transcription by uploading an audio file.
Before running this script, you must start the vLLM-Omni server with a realtime-capable
model, for example:
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni
Requirements:
- vllm with audio support
- websockets
- librosa
- numpy
The script:
1. Connects to the Realtime WebSocket endpoint
2. Converts an audio file to PCM16 @ 16kHz
3. Sends audio chunks to the server
4. Receives and prints transcription as it streams
"""
import argparse
import asyncio
import base64
import json
import librosa
import numpy as np
import websockets
from vllm.assets.audio import AudioAsset
def audio_to_pcm16_base64(audio_path: str) -> str:
"""
Load an audio file and convert it to base64-encoded PCM16 @ 16kHz.
"""
# Load audio and resample to 16kHz mono
audio, _ = librosa.load(audio_path, sr=16000, mono=True)
# Convert to PCM16
pcm16 = (audio * 32767).astype(np.int16)
# Encode as base64
return base64.b64encode(pcm16.tobytes()).decode("utf-8")
async def realtime_transcribe(audio_path: str, host: str, port: int, model: str):
"""
Connect to the Realtime API and transcribe an audio file.
"""
uri = f"ws://{host}:{port}/v1/realtime"
async with websockets.connect(uri) as ws:
# Wait for session.created
response = json.loads(await ws.recv())
if response["type"] == "session.created":
print(f"Session created: {response['id']}")
else:
print(f"Unexpected response: {response}")
return
# Validate model
await ws.send(json.dumps({"type": "session.update", "model": model}))
# Signal ready to start
await ws.send(json.dumps({"type": "input_audio_buffer.commit"}))
# Convert audio file to base64 PCM16
print(f"Loading audio from: {audio_path}")
audio_base64 = audio_to_pcm16_base64(audio_path)
# Send audio in chunks (4KB of raw audio = ~8KB base64)
chunk_size = 4096
audio_bytes = base64.b64decode(audio_base64)
total_chunks = (len(audio_bytes) + chunk_size - 1) // chunk_size
print(f"Sending {total_chunks} audio chunks...")
for i in range(0, len(audio_bytes), chunk_size):
chunk = audio_bytes[i : i + chunk_size]
await ws.send(
json.dumps(
{
"type": "input_audio_buffer.append",
"audio": base64.b64encode(chunk).decode("utf-8"),
}
)
)
# Signal all audio is sent
await ws.send(json.dumps({"type": "input_audio_buffer.commit", "final": True}))
print("Audio sent. Waiting for transcription...\n")
# Receive transcription
print("Transcription: ", end="", flush=True)
while True:
response = json.loads(await ws.recv())
if response["type"] == "transcription.delta":
print(response["delta"], end="", flush=True)
elif response["type"] == "transcription.done":
print(f"\n\nFinal transcription: {response['text']}")
if response.get("usage"):
print(f"Usage: {response['usage']}")
break
elif response["type"] == "error":
print(f"\nError: {response['error']}")
break
def main(args):
if args.audio_path:
audio_path = args.audio_path
else:
# Use default audio asset
audio_path = str(AudioAsset("mary_had_lamb").get_local_path())
print(f"No audio path provided, using default: {audio_path}")
asyncio.run(realtime_transcribe(audio_path, args.host, args.port, args.model))
if __name__ == "__main__":
parser = argparse.ArgumentParser(description="Realtime WebSocket Transcription Client")
parser.add_argument(
"--model",
type=str,
default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
help="Model that is served and should be pinged.",
)
parser.add_argument(
"--audio_path",
type=str,
default=None,
help="Path to the audio file to transcribe.",
)
parser.add_argument(
"--host",
type=str,
default="localhost",
help="vLLM-Omni server host (default: localhost)",
)
parser.add_argument(
"--port",
type=int,
default=8000,
help="vLLM-Omni server port (default: 8000)",
)
args = parser.parse_args()
main(args)
+8 -2
View File
@@ -233,6 +233,7 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
- `consent` (required): Consent recording ID
- `name` (required): Name for the new voice
- `ref_text` (optional): Transcript of the audio. Enables in-context voice cloning (higher quality).
- `speaker_description` (optional): Free-form description of the voice (e.g. "warm narrator", "energetic presenter").
**Response Example:**
```json
@@ -243,18 +244,23 @@ Upload a new voice sample for voice cloning in Base task TTS requests.
"consent": "user_consent_id",
"created_at": 1738660000,
"mime_type": "audio/wav",
"file_size": 1024000
"file_size": 1024000,
"ref_text": "The exact transcript of the audio sample.",
"speaker_description": "warm narrator"
}
}
```
Fields `ref_text` and `speaker_description` are omitted when not provided at upload time.
**Usage Example:**
```bash
curl -X POST http://localhost:8000/v1/audio/voices \
-F "audio_sample=@/path/to/voice_sample.wav" \
-F "consent=user_consent_id" \
-F "name=custom_voice_1" \
-F "ref_text=The exact transcript of the audio sample."
-F "ref_text=The exact transcript of the audio sample." \
-F "speaker_description=warm narrator"
```
### Endpoint
@@ -125,7 +125,7 @@ PLAYER_HTML = """
def _build_player_js(sample_rate: int) -> str:
"""Build the JavaScript that powers the AudioWorklet player."""
return f"""
() => {{
<script>
const SR = {sample_rate};
const WC = {json.dumps(WORKLET_JS)};
let ctx = null, node = null, abort = null, gen = false, st = {{}};
@@ -284,7 +284,7 @@ def _build_player_js(sample_rate: int) -> str:
}}
}}
}};
}}
</script>
"""
@@ -435,10 +435,7 @@ def create_app(api_base: str):
)
with gr.Blocks(
css=css,
title="Qwen3-TTS Demo",
js=_build_player_js(PCM_SAMPLE_RATE),
theme=theme,
) as demo:
gr.HTML(f"""
<div style="display:flex; align-items:center; gap:16px; margin-bottom:8px;">
@@ -773,7 +770,14 @@ def create_app(api_base: str):
demo.queue()
return gr.mount_gradio_app(fastapi_app, demo, path="/")
return gr.mount_gradio_app(
fastapi_app,
demo,
path="/",
css=css,
theme=theme,
head=_build_player_js(PCM_SAMPLE_RATE),
)
def parse_args():
@@ -103,7 +103,7 @@ def build_payload(
if task_type == "CustomVoice":
if voice:
payload["speaker"] = voice
payload["voice"] = voice
if instructions and instructions.strip():
payload["instructions"] = instructions.strip()
@@ -34,12 +34,43 @@ artifact, poll the job status and then download the completed file from the
content endpoint.
The main endpoints are:
- `POST /v1/videos`: create a video generation job
- `POST /v1/videos`: create a video generation job (async)
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
- `GET /v1/videos`: list stored video jobs
- `GET /v1/videos/{video_id}/content`: download the generated video file
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
## Sync API (Benchmark / Testing)
`POST /v1/videos/sync` is a synchronous alternative that blocks until generation
completes and returns the raw video bytes (`video/mp4`) directly in the response
body. It is designed for benchmark and testing scenarios where one-shot
request/response latency measurement is needed.
The sync endpoint accepts the same form parameters as `POST /v1/videos`. It does
not create any stored job record — the response is purely the generated video
file. Metadata is returned via response headers:
- `X-Request-Id`: unique identifier for this generation request
- `X-Model`: model name used for generation
- `X-Inference-Time-S`: wall-clock inference time in seconds
```bash
curl -X POST http://localhost:8091/v1/videos/sync \
-F "prompt=Two anthropomorphic cats in comfy boxing gear and bright gloves fight intensely on a spotlighted stage." \
-F "size=832x480" \
-F "num_frames=33" \
-F "fps=16" \
-F "num_inference_steps=40" \
-F "guidance_scale=4.0" \
-F "guidance_scale_2=4.0" \
-F "boundary_ratio=0.875" \
-F "flow_shift=5.0" \
-F "seed=42" \
-o sync_t2v_output.mp4
```
## Storage
Generated video files are stored on local disk by the async video API.
+5
View File
@@ -177,20 +177,24 @@ markers = [
"rocm: Tests that run on AMD/ROCm (auto-added)",
"xpu: Tests that run on XPU (auto-added)",
"npu: Tests that run on NPU/Ascend (auto-added)",
"musa: Tests that run on MUSA/Moore Threads (auto-added)",
# specified computation resources marks (auto-added)
"H100: Tests that require H100 GPU",
"L4: Tests that require L4 GPU",
"MI325: Tests that require MI325 GPU (AMD/ROCm)",
"S5000: Tests that require S5000 GPU (Moore Threads/MUSA)",
"A2: Tests that require A2 NPU",
"A3: Tests that require A3 NPU",
"distributed_cuda: Tests that require multi cards on CUDA platform",
"distributed_rocm: Tests that require multi cards on ROCm platform",
"distributed_xpu: Tests that require multi cards on XPU platform",
"distributed_npu: Tests that require multi cards on NPU platform",
"distributed_musa: Tests that require multi cards on MUSA platform",
"skipif_cuda: Skip if the num of CUDA cards is less than the required",
"skipif_rocm: Skip if the num of ROCm cards is less than the required",
"skipif_xpu: Skip if the num of XPU cards is less than the required",
"skipif_npu: Skip if the num of NPU cards is less than the required",
"skipif_musa: Skip if the num of MUSA cards is less than the required",
# more detailed markers
"slow: Slow tests (may skip in quick CI)",
"benchmark: Benchmark tests",
@@ -219,3 +223,4 @@ extend-ignore-identifiers-re = [
ue = "ue"
semantics = "semantics"
fullset = "fullset"
Vai = "Vai"
+4
View File
@@ -0,0 +1,4 @@
-r common.txt
# MUSA platform dependencies
torchada>=0.1.46
onnxruntime>=1.23.2
+15 -4
View File
@@ -3,7 +3,7 @@ Setup script for vLLM-Omni with hardware-dependent installation.
This setup.py implements platform-aware dependency routing so users can run
`pip install vllm-omni` and automatically receive the correct platform-specific
dependencies (CUDA/ROCm/CPU/XPU/NPU) without requiring extras like `[cuda]`.
dependencies (CUDA/ROCm/CPU/XPU/NPU/MUSA) without requiring extras like `[cuda]`.
"""
import os
@@ -46,16 +46,16 @@ def detect_target_device() -> str:
Priority order:
1. VLLM_OMNI_TARGET_DEVICE environment variable (highest priority)
2. Torch backend detection (cuda, rocm, npu, xpu)
2. Torch backend detection (cuda, rocm, npu, xpu, musa)
3. CPU fallback (default)
Returns:
str: Device name ('cuda', 'rocm', 'npu', 'xpu', or 'cpu')
str: Device name ('cuda', 'rocm', 'npu', 'xpu', 'musa', or 'cpu')
"""
# Priority 1: Explicit override via environment variable
target_device = os.environ.get("VLLM_OMNI_TARGET_DEVICE")
if target_device:
valid_devices = ["cuda", "rocm", "npu", "xpu", "cpu"]
valid_devices = ["cuda", "rocm", "npu", "xpu", "musa", "cpu"]
if target_device.lower() in valid_devices:
print(f"Using target device from VLLM_OMNI_TARGET_DEVICE: {target_device.lower()}")
return target_device.lower()
@@ -97,6 +97,15 @@ def detect_target_device() -> str:
except Exception:
pass
# Check for MUSA (Moore Threads)
if hasattr(torch, "musa"):
try:
if torch.musa.is_available():
print("Detected MUSA backend from torch")
return "musa"
except Exception:
pass
print("No GPU backend detected in torch, defaulting to CPU")
return "cpu"
@@ -152,6 +161,8 @@ def get_vllm_omni_version() -> str:
version += f"{sep}npu"
elif device == "xpu":
version += f"{sep}xpu"
elif device == "musa":
version += f"{sep}musa"
elif device == "cpu":
version += f"{sep}cpu"
else:
+246 -81
View File
@@ -48,6 +48,7 @@ from vllm.utils.network_utils import get_open_port
from vllm_omni.entrypoints.omni import Omni
from vllm_omni.inputs.data import OmniSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
logger = init_logger(__name__)
@@ -688,8 +689,112 @@ def generate_synthetic_audio(
return result
def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_file: bool = False) -> dict[str, Any]:
"""Generate synthetic video with bouncing balls and return base64 string."""
def _mux_mp4_bytes_with_synthetic_audio(
video_mp4_bytes: bytes,
*,
num_frames: int,
fps: float = 30.0,
sample_rate: int = 48000,
) -> bytes:
"""
Mux a video-only MP4 with mono TTS audio from :func:`generate_synthetic_audio` (AAC).
Audio length is at least the video duration in whole seconds (rounded up); ffmpeg
``-shortest`` trims to the video when the WAV is longer.
Uses ffmpeg from ``imageio_ffmpeg`` when available, else ``ffmpeg`` on PATH.
If TTS or mux fails, returns ``video_mp4_bytes`` unchanged.
Mux subprocess does **not** use ``capture_output=True``: ffmpeg can block writing
to a full stderr pipe while :func:`subprocess.run` waits for exit (classic deadlock).
"""
duration_sec = num_frames / fps if fps > 0 else 0.0
# generate_synthetic_audio(duration=int) uses at least 1s of buffer internally
duration_int = max(1, int(math.ceil(duration_sec)))
try:
audio_result = generate_synthetic_audio(
duration=duration_int,
num_channels=1,
sample_rate=sample_rate,
save_to_file=False,
)
audio_pcm = audio_result["np_array"]
except Exception as e:
logger.warning("Synthetic video: generate_synthetic_audio failed (%s); using video-only MP4.", e)
return video_mp4_bytes
try:
import imageio_ffmpeg
ffmpeg_exe = imageio_ffmpeg.get_ffmpeg_exe()
except Exception:
ffmpeg_exe = "ffmpeg"
import tempfile
try:
with tempfile.TemporaryDirectory(prefix="syn_vid_mux_") as tmp:
vid_path = os.path.join(tmp, "video.mp4")
wav_path = os.path.join(tmp, "audio.wav")
out_path = os.path.join(tmp, "out.mp4")
with open(vid_path, "wb") as f:
f.write(video_mp4_bytes)
sf.write(wav_path, audio_pcm, sample_rate, format="WAV", subtype="PCM_16")
cmd = [
ffmpeg_exe,
"-y",
"-nostdin",
"-hide_banner",
"-loglevel",
"error",
"-i",
vid_path,
"-i",
wav_path,
"-c:v",
"copy",
"-c:a",
"aac",
"-b:a",
"128k",
"-shortest",
"-movflags",
"+faststart",
out_path,
]
subprocess.run(
cmd,
check=True,
stdin=subprocess.DEVNULL,
timeout=300,
)
with open(out_path, "rb") as f:
return f.read()
except (
FileNotFoundError,
subprocess.CalledProcessError,
subprocess.TimeoutExpired,
OSError,
) as e:
logger.warning("Synthetic video: audio mux failed (%s); using video-only MP4.", e)
return video_mp4_bytes
def generate_synthetic_video(
width: int,
height: int,
num_frames: int,
save_to_file: bool = False,
*,
embed_audio: bool = False,
) -> dict[str, Any]:
"""Generate synthetic video with bouncing balls and base64 MP4.
When ``embed_audio`` is True, muxes mono AAC from :func:`generate_synthetic_audio`
(TTS + ffmpeg) into the MP4; otherwise returns video-only MP4 (faster when tests do
not need an audio track).
"""
import cv2
import imageio
@@ -762,13 +867,13 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
result = {
"np_array": video_array,
}
video_bytes = None
saved_file_path = None
fps = 30
buffer = io.BytesIO()
writer_kwargs = {
"format": "mp4",
"fps": 30,
"fps": fps,
"codec": "libx264",
"quality": 7,
"pixelformat": "yuv420p",
@@ -787,32 +892,31 @@ def generate_synthetic_video(width: int, height: int, num_frames: int, save_to_f
],
}
if save_to_file:
import datetime
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
output_path = f"video_{width}x{height}_{timestamp}.mp4"
try:
with imageio.get_writer(output_path, **writer_kwargs) as writer:
for frame in video_frames:
writer.append_data(frame)
saved_file_path = output_path
print(f"Video saved to: {saved_file_path}")
with open(output_path, "rb") as f:
video_bytes = f.read()
except Exception as e:
print(f"Warning: Failed to save video to file {output_path}: {e}")
save_to_file = False
if not save_to_file or video_bytes is None:
try:
with imageio.get_writer(buffer, **writer_kwargs) as writer:
for frame in video_frames:
writer.append_data(frame)
buffer.seek(0)
video_bytes = buffer.read()
video_only_bytes = buffer.read()
except Exception as e:
print(f"Warning: Failed to encode synthetic video: {e}")
raise
if embed_audio:
video_bytes = _mux_mp4_bytes_with_synthetic_audio(video_only_bytes, num_frames=num_frames, fps=float(fps))
else:
video_bytes = video_only_bytes
if save_to_file:
timestamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
output_path = f"video_{width}x{height}_{timestamp}.mp4"
try:
with open(output_path, "wb") as f:
f.write(video_bytes)
saved_file_path = output_path
print(f"Video saved to: {saved_file_path}")
except Exception as e:
print(f"Warning: Failed to save video to file {output_path}: {e}")
base64_video = base64.b64encode(video_bytes).decode("utf-8")
@@ -962,7 +1066,7 @@ def convert_audio_to_text(audio_data):
Convert base64 encoded audio data to text using speech recognition.
"""
audio_data = base64.b64decode(audio_data)
output_path = f"./test_{int(time.time())}.wav"
output_path = f"./test_{uuid.uuid4().hex}.wav"
with open(output_path, "wb") as audio_file:
audio_file.write(audio_data)
@@ -986,8 +1090,24 @@ def _merge_base64_audio_to_segment(base64_list: list[str]):
def _whisper_transcribe_in_current_process(output_path: str) -> str:
import whisper
# Keep Whisper on CPU to avoid consuming GPU memory in tests.
model = whisper.load_model("small", device="cpu")
# Multi-GPU: use last visible device to avoid colliding with default device 0; single device uses 0.
device_index = None
if current_omni_platform.is_available():
n = current_omni_platform.get_device_count()
if n == 1:
device_index = 0
elif n > 1:
device_index = n - 1
if device_index is not None:
torch_device = current_omni_platform.get_torch_device(device_index)
current_omni_platform.set_device(torch_device)
device = str(torch_device)
use_accelerator = True
else:
use_accelerator = False
device = "cpu"
model = whisper.load_model("small", device=device)
try:
text = model.transcribe(
output_path,
@@ -998,6 +1118,9 @@ def _whisper_transcribe_in_current_process(output_path: str) -> str:
finally:
del model
gc.collect()
if use_accelerator:
current_omni_platform.synchronize()
current_omni_platform.empty_cache()
return text or ""
@@ -1023,18 +1146,6 @@ def convert_audio_bytes_to_text(raw_bytes: bytes) -> str:
return text
def merge_base64_and_convert_to_text(base64_list):
"""
Merge a list of base64 encoded audio data and convert to text.
"""
merged_audio = _merge_base64_audio_to_segment(base64_list)
output_path = f"./test_{uuid.uuid4().hex}.wav"
merged_audio.export(output_path, format="wav")
print(f"audio data is saved: {output_path}")
text = convert_audio_file_to_text(output_path)
return text
def modify_stage_config(
yaml_path: str,
updates: dict[str, Any] = None,
@@ -1262,10 +1373,12 @@ def modify_stage_config(
# Direct top-level key
config[key] = value
# Save to new file with timestamp
timestamp = int(time.time())
# Unique suffix: multiple modify_stage_config calls in one process often run
# within the same second (e.g. test_qwen3_omni_expansion imports both
# get_chunk_config and get_batch_token_config). int(time.time()) would collide
# and the later write would overwrite the earlier YAML on disk.
base_name = yaml_path.rsplit(".", 1)[0] if "." in yaml_path else yaml_path
output_path = f"{base_name}_{timestamp}.yaml"
output_path = f"{base_name}_{time.time_ns()}.yaml"
with open(output_path, "w", encoding="utf-8") as f:
yaml.dump(config, f, default_flow_style=None, sort_keys=False, allow_unicode=True, indent=2)
@@ -1617,7 +1730,7 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
label = str(top.get("label", "")).lower()
conf = float(top.get("score", 0.0))
if conf < 0.6:
if conf < 0.5:
gender = "unknown"
# Some models use non-English labels (e.g., Russian). Normalize to 'male'/'female'.
elif ("female" in label) or ("жен" in label):
@@ -1646,6 +1759,34 @@ def _estimate_voice_gender_from_audio(audio_bytes: bytes) -> str:
return "unknown"
_PRESET_VOICE_GENDER_MAP: dict[str, str] = {
"serena": "female",
"uncle_fu": "male",
"chelsie": "female",
"clone": "female",
"ethan": "male",
}
def _assert_preset_voice_gender_from_audio(
audio_bytes: bytes | None,
voice_name: str | None,
) -> None:
"""If ``voice_name`` matches a known preset, assert classifier gender matches (skip when unknown)."""
if not voice_name or not audio_bytes:
return
key = str(voice_name).lower()
expected_gender = _PRESET_VOICE_GENDER_MAP.get(key)
if expected_gender is None:
return
estimated_gender = _estimate_voice_gender_from_audio(audio_bytes)
print(f"Preset voice gender check: preset={key!r}, estimated={estimated_gender!r}, expected={expected_gender!r}")
if estimated_gender != "unknown":
assert estimated_gender == expected_gender, (
f"{voice_name!r} is expected {expected_gender}, but estimated gender is {estimated_gender!r}"
)
# Threshold aligned with _compute_pcm_hnr_db docstring (clean clone vs distorted).
_MIN_PCM_SPEECH_HNR_DB = 1.0
@@ -1712,6 +1853,12 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
if "audio" in modalities:
assert response.audio_content is not None, "No audio output is generated"
print(f"audio content is: {response.audio_content}")
speaker = request_config.get("speaker")
if speaker:
_assert_preset_voice_gender_from_audio(
response.audio_bytes,
speaker,
)
if "text" in modalities:
assert response.text_content is not None, "No text output is generated"
@@ -1724,18 +1871,22 @@ def assert_omni_response(response: OmniResponse, request_config: dict[str, Any],
keywords = keywords_dict.get(word_type)
if "text" in modalities:
if keywords:
assert any(keyword in response.text_content.lower() for keyword in keywords), (
text_lower = response.text_content.lower()
assert any(str(kw).lower() in text_lower for kw in keywords), (
"The output does not contain any of the keywords."
)
else:
if keywords:
assert any(keyword in response.audio_content.lower() for keyword in keywords), (
audio_lower = response.audio_content.lower()
assert any(str(kw).lower() in audio_lower for kw in keywords), (
"The output does not contain any of the keywords."
)
# Verify similarity
# Verify similarity (Whisper transcript vs streamed/detokenized text)
if "text" in modalities and "audio" in modalities:
assert response.similarity > 0.9, "The audio content is not same as the text"
assert response.similarity is not None and response.similarity > 0.9, (
"The audio content is not same as the text"
)
print(f"similarity is: {response.similarity}")
@@ -1781,24 +1932,12 @@ def assert_audio_speech_response(
f"Transcript doesn't match input: similarity={similarity:.2f}, transcript='{transcript}'"
)
# Voice gender consistency check:
# Voice gender consistency check (preset names in ``_PRESET_VOICE_GENDER_MAP``).
# When the estimator returns 'unknown', we treat it as inconclusive and do NOT fail the test.
voice = (request_config.get("voice") or "").lower()
if voice and response.audio_bytes:
estimated_gender = _estimate_voice_gender_from_audio(response.audio_bytes)
voice_gender_map = {
# adjust this mapping to your actual voice names
"serena": "female",
"uncle_fu": "male",
"clone": "female",
}
expected_gender = voice_gender_map.get(voice)
if expected_gender is not None:
print(f"Estimated voice gender from audio: {estimated_gender} (voice='{voice}')")
if estimated_gender != "unknown":
assert estimated_gender == expected_gender, (
f"Voice '{voice}' is expected {expected_gender}, but estimated gender is '{estimated_gender}'"
)
_assert_preset_voice_gender_from_audio(
response.audio_bytes,
request_config.get("voice"),
)
def assert_diffusion_response(response: DiffusionResponse, request_config: dict[str, Any], run_level: str = None):
@@ -1914,7 +2053,11 @@ class OpenAIClientHandler:
if audio_data or text_content:
if audio_data:
audio_content = merge_base64_and_convert_to_text(audio_data)
merged_seg = _merge_base64_audio_to_segment(audio_data)
wav_buf = BytesIO()
merged_seg.export(wav_buf, format="wav")
result.audio_bytes = wav_buf.getvalue()
audio_content = convert_audio_bytes_to_text(result.audio_bytes)
if audio_content and text_content:
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
@@ -1969,7 +2112,8 @@ class OpenAIClientHandler:
if audio_data or text_content:
if audio_data:
audio_content = convert_audio_to_text(audio_data)
result.audio_bytes = base64.b64decode(audio_data)
audio_content = convert_audio_bytes_to_text(result.audio_bytes)
if audio_content and text_content:
similarity = cosine_similarity_text(audio_content.lower(), text_content.lower())
@@ -2135,7 +2279,12 @@ class OpenAIClientHandler:
Send OpenAI requests.
Args:
request_config: Request configuration dictionary containing parameters like model, messages, stream
request_config: Request configuration dictionary containing parameters like model, messages, stream.
Optional ``use_audio_in_video`` (bool): when true, sets
``extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True}`` for Qwen-Omni video+audio
extraction.
Optional top-level ``speaker`` (str): Qwen3-Omni preset TTS speaker name; sent as
``extra_body["speaker"]`` to ``chat.completions.create``.
request_num: Number of requests, defaults to 1 (single request)
Returns:
@@ -2146,14 +2295,27 @@ class OpenAIClientHandler:
stream = request_config.get("stream", False)
modalities = request_config.get("modalities", ["text", "audio"])
extra_body: dict[str, Any] = {}
if "speaker" in request_config:
extra_body["speaker"] = request_config["speaker"]
if request_config.get("use_audio_in_video"):
mm = dict(extra_body.get("mm_processor_kwargs") or {})
mm["use_audio_in_video"] = True
extra_body["mm_processor_kwargs"] = mm
extra_body_arg: dict[str, Any] | None = extra_body if extra_body else None
create_kwargs: dict[str, Any] = {
"model": request_config.get("model"),
"messages": request_config.get("messages"),
"stream": stream,
"modalities": modalities,
}
if extra_body_arg is not None:
create_kwargs["extra_body"] = extra_body_arg
if request_num == 1:
# Send single request
chat_completion = self.client.chat.completions.create(
model=request_config.get("model"),
messages=request_config.get("messages"),
stream=stream,
modalities=modalities,
)
chat_completion = self.client.chat.completions.create(**create_kwargs)
if stream:
response = self._process_stream_omni_response(chat_completion)
@@ -2167,12 +2329,15 @@ class OpenAIClientHandler:
# Send concurrent requests: run create + process in worker so e2e_latency includes full round-trip.
def _one_omni_request():
start = time.perf_counter()
chat_completion = self.client.chat.completions.create(
model=request_config.get("model"),
messages=request_config.get("messages"),
modalities=modalities,
stream=stream,
)
worker_kwargs: dict[str, Any] = {
"model": request_config.get("model"),
"messages": request_config.get("messages"),
"modalities": modalities,
"stream": stream,
}
if extra_body_arg is not None:
worker_kwargs["extra_body"] = extra_body_arg
chat_completion = self.client.chat.completions.create(**worker_kwargs)
if stream:
response = self._process_stream_omni_response(chat_completion)
else:
View File
Whitespace-only changes.
+237
View File
@@ -0,0 +1,237 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for shared AdaLayerNorm layers used by FLUX and other models."""
import os
import pytest
import torch
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
@pytest.fixture(autouse=True)
def _init_distributed():
"""Initialize the minimal distributed environment required by
ReplicatedLinear (tensor-parallel group must exist)."""
from vllm.distributed.parallel_state import (
cleanup_dist_env_and_memory,
init_distributed_environment,
initialize_model_parallel,
)
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29501")
init_distributed_environment(
world_size=1,
rank=0,
local_rank=0,
distributed_init_method="env://",
)
initialize_model_parallel()
yield
cleanup_dist_env_and_memory()
@pytest.fixture(autouse=True)
def _force_default_gemm(monkeypatch):
"""Force CPU-compatible GEMM dispatch for tests using CPU tensors.
vLLM's dispatch_unquantized_gemm() selects the backend by platform
(e.g. rocm_unquantized_gemm on AMD machines), not by tensor device.
CPU test tensors crash with NotImplementedError on ROCm. Monkeypatch
the dispatcher to always return the default (torch.nn.functional.linear)
implementation which works on any device."""
from vllm.model_executor.layers.utils import default_unquantized_gemm
monkeypatch.setattr(
"vllm.model_executor.layers.linear.dispatch_unquantized_gemm",
lambda: default_unquantized_gemm,
)
def test_adalayernorm_import_from_shared_module():
"""Verify imports work from the shared adalayernorm module."""
from vllm_omni.diffusion.layers.adalayernorm import ( # noqa: F401
AdaLayerNormContinuous,
AdaLayerNormZero,
AdaLayerNormZeroSingle,
)
def test_adalayernorm_zero_forward_shape():
"""AdaLayerNormZero produces correct output shapes (x, gate, shift, scale, gate)."""
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
dim = 64
batch = 2
seq_len = 4
norm = AdaLayerNormZero(dim)
x = torch.randn(batch, seq_len, dim)
emb = torch.randn(batch, dim)
out_x, gate_msa, shift_mlp, scale_mlp, gate_mlp = norm(x, emb)
assert out_x.shape == (batch, seq_len, dim)
assert gate_msa.shape == (batch, dim)
assert shift_mlp.shape == (batch, dim)
assert scale_mlp.shape == (batch, dim)
assert gate_mlp.shape == (batch, dim)
def test_adalayernorm_zero_single_forward_shape():
"""AdaLayerNormZeroSingle produces (x, gate) with correct shapes."""
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
dim = 64
batch = 2
seq_len = 4
norm = AdaLayerNormZeroSingle(dim)
x = torch.randn(batch, seq_len, dim)
emb = torch.randn(batch, dim)
out_x, gate = norm(x, emb)
assert out_x.shape == (batch, seq_len, dim)
assert gate.shape == (batch, dim)
def test_adalayernorm_continuous_forward_shape():
"""AdaLayerNormContinuous produces correct output shape."""
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
dim = 64
cond_dim = 64
batch = 2
seq_len = 4
norm = AdaLayerNormContinuous(dim, cond_dim)
x = torch.randn(batch, seq_len, dim)
conditioning = torch.randn(batch, cond_dim)
out = norm(x, conditioning)
assert out.shape == (batch, seq_len, dim)
def test_adalayernorm_zero_accepts_quant_config():
"""Constructor accepts quant_config=None and prefix='test' without error."""
from vllm_omni.diffusion.layers.adalayernorm import (
AdaLayerNormContinuous,
AdaLayerNormZero,
AdaLayerNormZeroSingle,
)
# Should not raise with quant_config=None and prefix
AdaLayerNormZero(64, quant_config=None, prefix="test.norm1")
AdaLayerNormZeroSingle(64, quant_config=None, prefix="test.norm")
AdaLayerNormContinuous(64, 64, quant_config=None, prefix="test.norm_out")
def test_adalayernorm_uses_replicated_linear():
"""Verify .linear is a ReplicatedLinear instance (not nn.Linear)."""
from vllm.model_executor.layers.linear import ReplicatedLinear
from vllm_omni.diffusion.layers.adalayernorm import (
AdaLayerNormContinuous,
AdaLayerNormZero,
AdaLayerNormZeroSingle,
)
norm_zero = AdaLayerNormZero(64)
assert isinstance(norm_zero.linear, ReplicatedLinear)
norm_zero_single = AdaLayerNormZeroSingle(64)
assert isinstance(norm_zero_single.linear, ReplicatedLinear)
norm_continuous = AdaLayerNormContinuous(64, 64)
assert isinstance(norm_continuous.linear, ReplicatedLinear)
# ── Numerical equivalence tests against diffusers originals ──
def _copy_weights(src_linear, dst_replicated_linear):
"""Copy weights from nn.Linear to ReplicatedLinear for comparison."""
dst_replicated_linear.weight.data.copy_(src_linear.weight.data)
if src_linear.bias is not None and dst_replicated_linear.bias is not None:
dst_replicated_linear.bias.data.copy_(src_linear.bias.data)
def test_adalayernorm_zero_matches_diffusers():
"""Verify AdaLayerNormZero produces identical output to diffusers original."""
from diffusers.models.normalization import (
AdaLayerNormZero as DiffusersAdaLayerNormZero,
)
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZero
dim = 64
torch.manual_seed(42)
ours = AdaLayerNormZero(dim)
ref = DiffusersAdaLayerNormZero(dim)
# Copy weights: nn.Linear -> ReplicatedLinear
_copy_weights(ref.linear, ours.linear)
x = torch.randn(2, 4, dim)
emb = torch.randn(2, dim)
out_ours = ours(x, emb)
out_ref = ref(x, emb=emb)
for o, r in zip(out_ours, out_ref):
torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
def test_adalayernorm_zero_single_matches_diffusers():
"""Verify AdaLayerNormZeroSingle produces identical output to diffusers original."""
from diffusers.models.normalization import (
AdaLayerNormZeroSingle as DiffusersAdaLayerNormZeroSingle,
)
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormZeroSingle
dim = 64
torch.manual_seed(42)
ours = AdaLayerNormZeroSingle(dim)
ref = DiffusersAdaLayerNormZeroSingle(dim)
_copy_weights(ref.linear, ours.linear)
x = torch.randn(2, 4, dim)
emb = torch.randn(2, dim)
out_ours = ours(x, emb)
out_ref = ref(x, emb=emb)
for o, r in zip(out_ours, out_ref):
torch.testing.assert_close(o, r, atol=1e-5, rtol=1e-5)
def test_adalayernorm_continuous_matches_diffusers():
"""Verify AdaLayerNormContinuous produces identical output to diffusers original."""
from diffusers.models.normalization import (
AdaLayerNormContinuous as DiffusersAdaLayerNormContinuous,
)
from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNormContinuous
dim = 64
cond_dim = 64
torch.manual_seed(42)
# Match constructor args: diffusers defaults elementwise_affine=True, eps=1e-5
ours = AdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
ref = DiffusersAdaLayerNormContinuous(dim, cond_dim, elementwise_affine=False, eps=1e-6)
_copy_weights(ref.linear, ours.linear)
x = torch.randn(2, 4, dim)
cond = torch.randn(2, cond_dim)
out_ours = ours(x, cond)
out_ref = ref(x, cond)
torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5)
Whitespace-only changes.
@@ -0,0 +1,134 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests that FLUX transformer blocks correctly propagate `quant_config` and
`prefix` through all sub-layers.
The tests instantiate blocks with a known prefix and verify that all quantization-
aware sub-layers (AdaLayerNorm, FeedForward, Attention projections) receive the
prefix rooted at the block prefix. This is critical for quantized weight loading
to match checkpoint keys to the correct model parameters.
"""
import os
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
# Standard dimensions for a minimal FLUX block
_DIM = 64
_HEADS = 2
_HEAD_DIM = _DIM // _HEADS
@pytest.fixture(autouse=True)
def _init_distributed():
"""Initialize the minimal distributed environment required by
vLLM parallel linear layers (tensor-parallel group must exist)."""
from vllm.distributed.parallel_state import (
cleanup_dist_env_and_memory,
init_distributed_environment,
initialize_model_parallel,
)
os.environ.setdefault("MASTER_ADDR", "localhost")
os.environ.setdefault("MASTER_PORT", "29502")
init_distributed_environment(
world_size=1,
rank=0,
local_rank=0,
distributed_init_method="env://",
)
initialize_model_parallel()
yield
cleanup_dist_env_and_memory()
def _param_names(module) -> set[str]:
"""Return the set of all parameter names in a module."""
return {name for name, _ in module.named_parameters()}
def test_flux_transformer_block_passes_prefix():
"""FluxTransformerBlock propagates prefix to norm1, norm1_context, attn, ff, ff_context."""
from vllm_omni.diffusion.models.flux.flux_transformer import FluxTransformerBlock
prefix = "transformer_blocks.0"
block = FluxTransformerBlock(
dim=_DIM,
num_attention_heads=_HEADS,
attention_head_dim=_HEAD_DIM,
quant_config=None,
prefix=prefix,
)
params = _param_names(block)
# norm1 and norm1_context (AdaLayerNormZero) should have linear weights
assert any(name.startswith("norm1.linear.") for name in params), (
f"norm1.linear.* not found in params: {sorted(params)}"
)
assert any(name.startswith("norm1_context.linear.") for name in params), (
f"norm1_context.linear.* not found in params: {sorted(params)}"
)
# attn should have QKV projections
assert any(name.startswith("attn.to_qkv.") for name in params), (
f"attn.to_qkv.* not found in params: {sorted(params)}"
)
# ff and ff_context should have net layers
assert any(name.startswith("ff.net.") for name in params), f"ff.net.* not found in params: {sorted(params)}"
assert any(name.startswith("ff_context.net.") for name in params), (
f"ff_context.net.* not found in params: {sorted(params)}"
)
def test_flux_single_transformer_block_passes_prefix():
"""FluxSingleTransformerBlock propagates prefix to norm, proj_mlp, attn."""
from vllm_omni.diffusion.models.flux.flux_transformer import FluxSingleTransformerBlock
prefix = "single_transformer_blocks.0"
block = FluxSingleTransformerBlock(
dim=_DIM,
num_attention_heads=_HEADS,
attention_head_dim=_HEAD_DIM,
quant_config=None,
prefix=prefix,
)
params = _param_names(block)
# norm (AdaLayerNormZeroSingle) should have linear weights
assert any(name.startswith("norm.linear.") for name in params), (
f"norm.linear.* not found in params: {sorted(params)}"
)
# proj_mlp (ReplicatedLinear) should have weight
assert any(name.startswith("proj_mlp.") for name in params), f"proj_mlp.* not found in params: {sorted(params)}"
# attn should have QKV projection
assert any(name.startswith("attn.to_qkv.") for name in params), (
f"attn.to_qkv.* not found in params: {sorted(params)}"
)
def test_flux_feedforward_passes_prefix():
"""FeedForward propagates prefix to net.0 (GELU proj) and net.2 (output proj)."""
from vllm_omni.diffusion.models.flux.flux_transformer import FeedForward
prefix = "transformer_blocks.0.ff"
ff = FeedForward(
dim=_DIM,
dim_out=_DIM,
quant_config=None,
prefix=prefix,
)
params = _param_names(ff)
# net.0 is ColumnParallelApproxGELU which wraps a ColumnParallelLinear
assert any("net.0" in name for name in params), f"net.0 not found in params: {sorted(params)}"
# net.2 is RowParallelLinear
assert any("net.2" in name for name in params), f"net.2 not found in params: {sorted(params)}"
@@ -0,0 +1,129 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Unit tests for LayerwiseOffloadHook."""
import gc
import os
import socket
from contextlib import contextmanager
import pytest
import torch
import torch.distributed as dist
from torch import nn
from torch.distributed.tensor import DeviceMesh, DTensor, Replicate
import vllm_omni.diffusion.offloader.layerwise_backend as layerwise_backend_module
from vllm_omni.diffusion.offloader.layerwise_backend import LayerwiseOffloadHook
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
class DummyStream:
def wait_stream(self, _stream) -> None:
return None
def wait_event(self, _event) -> None:
return None
class DummyEvent:
def record(self, _stream) -> None:
return None
@contextmanager
def dummy_stream(_stream):
yield None
def _find_free_port() -> int:
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as s:
s.bind(("127.0.0.1", 0))
return int(s.getsockname()[1])
def _set_dist_env(*, rank: int, world_size: int, master_port: int) -> None:
os.environ["RANK"] = str(rank)
os.environ["LOCAL_RANK"] = str(rank)
os.environ["WORLD_SIZE"] = str(world_size)
os.environ["MASTER_ADDR"] = "127.0.0.1"
os.environ["MASTER_PORT"] = str(master_port)
def _cleanup_distributed() -> None:
if dist.is_initialized():
dist.destroy_process_group()
for key in ["MASTER_ADDR", "MASTER_PORT", "RANK", "WORLD_SIZE", "LOCAL_RANK"]:
os.environ.pop(key, None)
gc.collect()
if current_omni_platform.is_available():
current_omni_platform.empty_cache()
current_omni_platform.synchronize()
@pytest.fixture(scope="module")
def dist_group():
master_port = _find_free_port()
_set_dist_env(rank=0, world_size=1, master_port=master_port)
dist.init_process_group("gloo", rank=0, world_size=1)
try:
yield
finally:
_cleanup_distributed()
@pytest.fixture
def patched_offload_runtime(mocker):
mocker.patch.object(layerwise_backend_module.current_omni_platform, "Stream", DummyStream)
mocker.patch.object(layerwise_backend_module.current_omni_platform, "Event", DummyEvent)
mocker.patch.object(layerwise_backend_module.current_omni_platform, "current_stream", lambda: DummyStream())
mocker.patch.object(layerwise_backend_module.current_omni_platform, "stream", dummy_stream)
class TinyBlock(nn.Module):
def __init__(self, values: torch.Tensor):
super().__init__()
mesh = DeviceMesh("cpu", [0])
dtensor = DTensor.from_local(values, mesh, [Replicate()])
self.weight = nn.Parameter(dtensor)
def _make_values(start: float) -> torch.Tensor:
return torch.arange(start, start + 4, dtype=torch.float32)
class TestLayerwiseOffloadHook:
def test_dtensor_wrapper_is_preserved_across_prefetch_and_offload(self, dist_group, patched_offload_runtime):
current_block = TinyBlock(_make_values(1.0))
next_block = TinyBlock(_make_values(10.0))
hook = LayerwiseOffloadHook(
next_block=next_block,
device=torch.device("cpu"),
stream=DummyStream(),
pin_memory=False,
)
hook.initialize_hook(current_block)
assert isinstance(next_block.weight, DTensor)
assert next_block.weight.to_local().is_meta
assert next_block.weight.to_local().shape == torch.Size([4])
assert hook.dtype_metadata[next_block.weight.dtype][0]["shape"] == torch.Size([4])
hook.prefetch_layer(non_blocking=False)
assert isinstance(next_block.weight, DTensor)
assert torch.equal(next_block.weight.to_local(), _make_values(10.0))
assert next_block.weight.to_local().shape == torch.Size([4])
hook.offload_layer()
assert isinstance(current_block.weight, DTensor)
assert current_block.weight.to_local().is_meta
assert current_block.weight.to_local().shape == torch.Size([4])
assert not hook.is_materialized
@@ -0,0 +1,147 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for INC/AutoRound quantization via the unified framework."""
import pytest
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
def test_build_quant_config_autoround():
"""build_quant_config("auto-round", ...) should produce an INCConfig."""
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm_omni.quantization import build_quant_config
config = build_quant_config(
"auto-round",
bits=4,
group_size=128,
sym=True,
packing_format="auto_round:auto_gptq",
)
assert config is not None
assert isinstance(config, INCConfig)
assert config.weight_bits == 4
assert config.group_size == 128
def test_build_quant_config_inc():
"""build_quant_config("inc", ...) should also produce an INCConfig."""
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm_omni.quantization import build_quant_config
config = build_quant_config("inc", bits=4, group_size=128)
assert isinstance(config, INCConfig)
assert config.weight_bits == 4
def test_build_quant_config_autoround_dict():
"""Dict-style config with method=auto-round should work."""
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm_omni.quantization import build_quant_config
config = build_quant_config(
{
"method": "auto-round",
"bits": 4,
"group_size": 128,
"sym": True,
"packing_format": "auto_round:auto_gptq",
}
)
assert isinstance(config, INCConfig)
assert config.weight_bits == 4
def test_build_quant_config_autoround_filters_metadata():
"""Checkpoint metadata keys (autoround_version, batch_size, iters)
should be silently filtered out instead of causing TypeError."""
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm_omni.quantization import build_quant_config
config = build_quant_config(
"auto-round",
bits=4,
group_size=128,
sym=True,
packing_format="auto_round:auto_gptq",
block_name_to_quantize="transformer_blocks,single_transformer_blocks",
autoround_version="0.12.0", # metadata — must be filtered
batch_size=1, # metadata — must be filtered
iters=0, # metadata — must be filtered
)
assert isinstance(config, INCConfig)
assert config.weight_bits == 4
assert config.group_size == 128
def test_build_quant_config_bits_to_weight_bits_mapping():
"""The 'bits' key from checkpoints should be mapped to 'weight_bits'."""
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm_omni.quantization import build_quant_config
# If weight_bits is already provided, bits should be ignored
config = build_quant_config("auto-round", weight_bits=4, group_size=128)
assert isinstance(config, INCConfig)
assert config.weight_bits == 4
def test_autoround_in_supported_methods():
"""auto-round and inc should appear in SUPPORTED_QUANTIZATION_METHODS."""
from vllm_omni.quantization import SUPPORTED_QUANTIZATION_METHODS
assert "auto-round" in SUPPORTED_QUANTIZATION_METHODS
assert "inc" in SUPPORTED_QUANTIZATION_METHODS
def test_integration_autoround_via_omni_diffusion_config():
"""OmniDiffusionConfig with auto-round quantization dict should resolve."""
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm_omni.diffusion.data import OmniDiffusionConfig
config = OmniDiffusionConfig(
model="test",
quantization_config={
"method": "auto-round",
"bits": 4,
"group_size": 128,
"sym": True,
},
)
assert isinstance(config.quantization_config, INCConfig)
assert config.quantization_config.weight_bits == 4
def test_integration_autodetect_from_transformer_config():
"""When TransformerConfig has quant_config, OmniDiffusionConfig should
auto-detect it even without explicit quantization_config."""
from vllm.model_executor.layers.quantization.inc import INCConfig
from vllm_omni.diffusion.data import OmniDiffusionConfig, TransformerConfig
tf_config = TransformerConfig.from_dict(
{
"quantization_config": {
"quant_method": "auto-round",
"bits": 4,
"group_size": 128,
"sym": True,
"packing_format": "auto_round:auto_gptq",
"autoround_version": "0.12.0",
"batch_size": 1,
"iters": 0,
}
}
)
assert tf_config.quant_method == "auto-round"
assert isinstance(tf_config.quant_config, INCConfig)
od_config = OmniDiffusionConfig(model="test", tf_model_config=tf_config)
assert isinstance(od_config.quantization_config, INCConfig)
assert od_config.quantization_config.weight_bits == 4
@@ -118,8 +118,12 @@ def _generate_image(omni, config: QualityTestConfig):
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
return req_out.images[0], peak_mem
if hasattr(first, "images") and first.images:
return first.images[0], peak_mem
inner = first.request_output
if inner is not None and hasattr(inner, "images") and inner.images:
return inner.images[0], peak_mem
raise ValueError("Could not extract image from output.")
def _generate_video(omni, config: QualityTestConfig):
+358 -21
View File
@@ -1,12 +1,15 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
import queue
import threading
from types import SimpleNamespace
from unittest.mock import Mock, patch
import pytest
import torch
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.data import DiffusionOutput, DiffusionRequestAbortedError
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import (
@@ -14,6 +17,7 @@ from vllm_omni.diffusion.sched import (
RequestScheduler,
Scheduler,
SchedulerInterface,
StepScheduler,
)
from vllm_omni.diffusion.sched.interface import CachedRequestData, NewRequestData
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
@@ -29,9 +33,46 @@ def _make_request(req_id: str) -> OmniDiffusionRequest:
)
def _make_request_output(req_id: str, *, error: str | None = None) -> DiffusionOutput:
del req_id
return DiffusionOutput(output=None, error=error)
def _make_request_output(req_id: str, *, error: str | None = None, finished: bool = True):
return SimpleNamespace(
req_id=req_id,
step_index=None,
finished=finished,
result=DiffusionOutput(output=None, error=error),
)
def _make_step_output(
req_id: str,
step_index: int,
*,
finished: bool = False,
error: str | None = None,
):
return SimpleNamespace(
req_id=req_id,
step_index=step_index,
finished=finished,
result=DiffusionOutput(output=None, error=error) if error is not None else None,
)
def _make_step_request(
req_id: str,
*,
num_inference_steps: int = 4,
step_index: int | None = None,
sampling_params: OmniDiffusionSamplingParams | None = None,
) -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompts=[f"prompt_{req_id}"],
sampling_params=sampling_params
or OmniDiffusionSamplingParams(
num_inference_steps=num_inference_steps,
step_index=step_index,
),
request_ids=[req_id],
)
def _new_ids(sched_output) -> list[str]:
@@ -43,7 +84,7 @@ def _cached_ids(sched_output) -> list[str]:
class _StubScheduler(SchedulerInterface):
def __init__(self, request: OmniDiffusionRequest, output: DiffusionOutput) -> None:
def __init__(self, request: OmniDiffusionRequest, output) -> None:
self._request = request
self._output = output
self.initialized_with = None
@@ -75,9 +116,10 @@ class _StubScheduler(SchedulerInterface):
is_empty=False,
)
def update_from_output(self, sched_output, output: DiffusionOutput) -> set[str]:
def update_from_output(self, sched_output, output) -> set[str]:
del sched_output
assert output is self._output
self._state.status = DiffusionRequestStatus.FINISHED_COMPLETED
return {self._sched_req_id}
def has_requests(self) -> bool:
@@ -185,9 +227,14 @@ class TestRequestScheduler:
state_b = self.scheduler.get_request_state(req_id_b)
assert state_b.status == DiffusionRequestStatus.FINISHED_ABORTED
first = self.scheduler.schedule()
assert first.finished_req_ids == {req_id_b}
# A should still run normally.
output_a = self.scheduler.schedule()
assert _new_ids(output_a) == [req_id_a]
assert _new_ids(first) == [req_id_a]
# B is already marked finished aborted, scheduling again should not pull it.
second = self.scheduler.schedule()
assert second.finished_req_ids == set()
# Abort running request.
self.scheduler.finish_requests(req_id_a, DiffusionRequestStatus.FINISHED_ABORTED)
@@ -233,33 +280,33 @@ class TestDiffusionEngine:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(Mock())
engine.executor = Mock()
engine._rpc_lock = threading.Lock()
engine._rpc_lock = threading.RLock()
engine.abort_queue = queue.Queue()
request = _make_request("engine")
expected = DiffusionOutput(output=None)
engine.executor.add_req.return_value = expected
runner_output = _make_request_output("engine")
engine.execute_fn = Mock(return_value=runner_output)
output = engine.add_req_and_wait_for_response(request)
assert output is expected
engine.executor.add_req.assert_called_once_with(request)
assert output is runner_output.result
engine.execute_fn.assert_called_once()
def test_supports_scheduler_interface_injection(self) -> None:
request = _make_request("engine_iface")
expected = DiffusionOutput(output=None)
scheduler = _StubScheduler(request, expected)
runner_output = _make_request_output("engine_iface")
scheduler = _StubScheduler(request, runner_output)
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = scheduler
engine.executor = Mock()
engine.executor.add_req = Mock(return_value=expected)
engine._rpc_lock = threading.Lock()
engine._rpc_lock = threading.RLock()
engine.abort_queue = queue.Queue()
engine.execute_fn = Mock(return_value=runner_output)
output = engine.add_req_and_wait_for_response(request)
assert output is expected
engine.executor.add_req.assert_called_once_with(request)
assert output is runner_output.result
engine.execute_fn.assert_called_once()
def test_initializes_injected_scheduler(self) -> None:
request = _make_request("init")
@@ -289,6 +336,59 @@ class TestDiffusionEngine:
assert req_id in finished
assert scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
def test_step_raises_aborted_error(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.pre_process_func = None
engine.add_req_and_wait_for_response = Mock(
return_value=DiffusionOutput(aborted=True, abort_message="Request req-abort aborted.")
)
with pytest.raises(DiffusionRequestAbortedError, match="Request req-abort aborted"):
engine.step(_make_request("req-abort"))
def test_abort_queue_marks_request_finished_aborted(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(Mock())
engine.abort_queue = queue.Queue()
req_id = engine.scheduler.add_request(_make_request("req-abort"))
engine.abort("req-abort")
engine._process_aborts_queue()
assert engine.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED
def test_finalize_finished_request_returns_aborted_output(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.scheduler = RequestScheduler()
engine.scheduler.initialize(Mock())
req_id = engine.scheduler.add_request(_make_request("req-finalize"))
engine.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
output = engine._finalize_finished_request(req_id)
assert output.aborted is True
assert output.abort_message == "Request req-finalize aborted."
def test_initializes_step_scheduler_when_step_execution_enabled(self) -> None:
od_config = Mock(model_class_name="mock_model")
od_config.step_execution = True
fake_executor = Mock()
fake_executor_cls = Mock(return_value=fake_executor)
with (
patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_post_process_func", return_value=None),
patch("vllm_omni.diffusion.diffusion_engine.get_diffusion_pre_process_func", return_value=None),
patch("vllm_omni.diffusion.diffusion_engine.DiffusionExecutor.get_class", return_value=fake_executor_cls),
patch.object(DiffusionEngine, "_dummy_run", return_value=None),
):
engine = DiffusionEngine(od_config)
assert isinstance(engine.scheduler, StepScheduler)
assert engine.execute_fn is fake_executor.execute_step
fake_executor_cls.assert_called_once_with(od_config)
def test_dummy_run_raises_on_output_error(self) -> None:
engine = DiffusionEngine.__new__(DiffusionEngine)
engine.od_config = Mock(model_class_name="mock_model")
@@ -297,3 +397,240 @@ class TestDiffusionEngine:
with pytest.raises(RuntimeError, match="Dummy run failed: boom"):
engine._dummy_run()
class TestStepScheduler:
def setup_method(self) -> None:
self.scheduler: StepScheduler = StepScheduler()
self.scheduler.initialize(Mock())
def test_single_request_step_lifecycle(self) -> None:
request = _make_step_request("step", num_inference_steps=3)
req_id = self.scheduler.add_request(request)
first = self.scheduler.schedule()
assert _new_ids(first) == [req_id]
assert _cached_ids(first) == []
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 0
finished = self.scheduler.update_from_output(first, _make_step_output(req_id, step_index=1))
assert finished == set()
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.RUNNING
assert request.sampling_params.step_index == 1
assert self.scheduler.has_requests() is True
second = self.scheduler.schedule()
assert _new_ids(second) == []
assert _cached_ids(second) == [req_id]
assert second.num_running_reqs == 1
assert second.num_waiting_reqs == 0
finished = self.scheduler.update_from_output(second, _make_step_output(req_id, step_index=2))
assert finished == set()
assert request.sampling_params.step_index == 2
third = self.scheduler.schedule()
assert _new_ids(third) == []
assert _cached_ids(third) == [req_id]
finished = self.scheduler.update_from_output(
third,
_make_step_output(req_id, step_index=3, finished=True),
)
assert finished == {req_id}
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
assert request.sampling_params.step_index == 3
assert self.scheduler.has_requests() is False
def test_fifo_single_request_scheduling(self) -> None:
req_id_a = self.scheduler.add_request(_make_step_request("a", num_inference_steps=2))
req_id_b = self.scheduler.add_request(_make_step_request("b", num_inference_steps=2))
first = self.scheduler.schedule()
assert _new_ids(first) == [req_id_a]
assert _cached_ids(first) == []
assert first.num_running_reqs == 1
assert first.num_waiting_reqs == 1
finished = self.scheduler.update_from_output(first, _make_step_output(req_id_a, step_index=1))
assert finished == set()
second = self.scheduler.schedule()
assert _new_ids(second) == []
assert _cached_ids(second) == [req_id_a]
assert second.num_running_reqs == 1
assert second.num_waiting_reqs == 1
finished = self.scheduler.update_from_output(
second,
_make_step_output(req_id_a, step_index=2, finished=True),
)
assert finished == {req_id_a}
third = self.scheduler.schedule()
assert _new_ids(third) == [req_id_b]
assert _cached_ids(third) == []
assert third.num_running_reqs == 1
assert third.num_waiting_reqs == 0
def test_error_output_marks_finished_error(self) -> None:
req_id = self.scheduler.add_request(_make_step_request("err", num_inference_steps=3))
sched_output = self.scheduler.schedule()
assert _new_ids(sched_output) == [req_id]
finished = self.scheduler.update_from_output(
sched_output,
_make_step_output(req_id, step_index=1, finished=True, error="worker failed"),
)
assert finished == {req_id}
state = self.scheduler.get_request_state(req_id)
assert state.status == DiffusionRequestStatus.FINISHED_ERROR
assert state.error == "worker failed"
assert self.scheduler.has_requests() is False
def test_missing_step_index_marks_finished_error(self) -> None:
req_id = self.scheduler.add_request(_make_step_request("missing", num_inference_steps=3))
sched_output = self.scheduler.schedule()
finished = self.scheduler.update_from_output(
sched_output,
SimpleNamespace(
req_id=req_id,
step_index=None,
finished=True,
result=None,
),
)
assert finished == {req_id}
state = self.scheduler.get_request_state(req_id)
assert state.status == DiffusionRequestStatus.FINISHED_ERROR
assert state.error == "Missing step_index in RunnerOutput"
def test_abort_request_for_waiting_and_running(self) -> None:
req_id_a = self.scheduler.add_request(_make_step_request("a", num_inference_steps=2))
req_id_b = self.scheduler.add_request(_make_step_request("b", num_inference_steps=2))
self.scheduler.finish_requests(req_id_b, DiffusionRequestStatus.FINISHED_ABORTED)
assert self.scheduler.get_request_state(req_id_b).status == DiffusionRequestStatus.FINISHED_ABORTED
running = self.scheduler.schedule()
assert _new_ids(running) == [req_id_a]
self.scheduler.finish_requests(req_id_a, DiffusionRequestStatus.FINISHED_ABORTED)
assert self.scheduler.get_request_state(req_id_a).status == DiffusionRequestStatus.FINISHED_ABORTED
assert self.scheduler.has_requests() is False
def test_has_requests_state_transition(self) -> None:
assert self.scheduler.has_requests() is False
req_id = self.scheduler.add_request(_make_step_request("has", num_inference_steps=2))
assert self.scheduler.has_requests() is True
sched_output = self.scheduler.schedule()
assert self.scheduler.has_requests() is True
finished = self.scheduler.update_from_output(
sched_output,
_make_step_output(req_id, step_index=2, finished=True),
)
assert finished == {req_id}
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
assert self.scheduler.has_requests() is False
def test_scheduled_request_aborted_before_update_is_returned_finished(self) -> None:
req_id = self.scheduler.add_request(_make_step_request("abort-late", num_inference_steps=2))
sched_output = self.scheduler.schedule()
self.scheduler.finish_requests(req_id, DiffusionRequestStatus.FINISHED_ABORTED)
finished = self.scheduler.update_from_output(
sched_output,
_make_step_output(req_id, step_index=1),
)
assert finished == {req_id}
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_ABORTED
def test_preempt_request_preserves_step_index(self) -> None:
request = _make_step_request("preempt", num_inference_steps=3)
req_id = self.scheduler.add_request(request)
first = self.scheduler.schedule()
assert self.scheduler.update_from_output(first, _make_step_output(req_id, step_index=1)) == set()
assert request.sampling_params.step_index == 1
second = self.scheduler.schedule()
assert _cached_ids(second) == [req_id]
assert self.scheduler.preempt_request(req_id) is True
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.PREEMPTED
assert request.sampling_params.step_index == 1
third = self.scheduler.schedule()
assert _cached_ids(third) == [req_id]
assert request.sampling_params.step_index == 1
@pytest.mark.parametrize(
("sampling_params", "expected_steps"),
[
(
OmniDiffusionSamplingParams(
timesteps=torch.tensor([1.0, 0.5, 0.0]),
sigmas=[1.0, 0.5, 0.25, 0.0],
num_inference_steps=5,
),
3,
),
(
OmniDiffusionSamplingParams(
sigmas=[1.0, 0.5],
num_inference_steps=5,
),
2,
),
(
OmniDiffusionSamplingParams(
num_inference_steps=4,
),
4,
),
],
)
def test_total_steps_priority(self, sampling_params: OmniDiffusionSamplingParams, expected_steps: int) -> None:
request = _make_step_request("priority", sampling_params=sampling_params)
req_id = self.scheduler.add_request(request)
for _ in range(expected_steps - 1):
sched_output = self.scheduler.schedule()
assert sched_output.scheduled_req_ids == [req_id]
next_step = request.sampling_params.step_index + 1
assert (
self.scheduler.update_from_output(
sched_output,
_make_step_output(req_id, step_index=next_step),
)
== set()
)
final_output = self.scheduler.schedule()
assert final_output.scheduled_req_ids == [req_id]
assert self.scheduler.update_from_output(
final_output,
_make_step_output(req_id, step_index=expected_steps, finished=True),
) == {req_id}
assert self.scheduler.get_request_state(req_id).status == DiffusionRequestStatus.FINISHED_COMPLETED
@pytest.mark.parametrize(
"sampling_params",
[
OmniDiffusionSamplingParams(num_inference_steps=0),
OmniDiffusionSamplingParams(num_inference_steps=3, step_index=3),
OmniDiffusionSamplingParams(num_inference_steps=3, step_index=-1),
],
)
def test_rejects_invalid_initial_step_state(self, sampling_params: OmniDiffusionSamplingParams) -> None:
request = _make_step_request("invalid", sampling_params=sampling_params)
with pytest.raises(ValueError):
self.scheduler.add_request(request)
+256 -1
View File
@@ -1,10 +1,13 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Tests for step-level diffusion runner and worker execution."""
"""Tests for step-level diffusion execution across runner / worker / executor / engine."""
import os
import queue
import threading
from contextlib import contextmanager
from types import SimpleNamespace
from unittest.mock import Mock
import pytest
import torch
@@ -12,6 +15,7 @@ import torch
import vllm_omni.diffusion.worker.diffusion_model_runner as model_runner_module
from tests.utils import hardware_test
from vllm_omni.diffusion.data import DiffusionOutput
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
from vllm_omni.diffusion.distributed.cfg_parallel import CFGParallelMixin
from vllm_omni.diffusion.distributed.comm import RingComm, SeqAllToAll4D
from vllm_omni.diffusion.distributed.parallel_state import (
@@ -20,10 +24,13 @@ from vllm_omni.diffusion.distributed.parallel_state import (
init_distributed_environment,
initialize_model_parallel,
)
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
from vllm_omni.diffusion.ipc import (
pack_diffusion_output_shm,
unpack_diffusion_output_shm,
)
from vllm_omni.diffusion.request import OmniDiffusionRequest
from vllm_omni.diffusion.sched import StepScheduler
from vllm_omni.diffusion.sched.interface import (
CachedRequestData,
DiffusionSchedulerOutput,
@@ -32,6 +39,8 @@ from vllm_omni.diffusion.sched.interface import (
from vllm_omni.diffusion.worker.diffusion_model_runner import DiffusionModelRunner
from vllm_omni.diffusion.worker.diffusion_worker import DiffusionWorker
from vllm_omni.diffusion.worker.utils import RunnerOutput
from vllm_omni.engine.async_omni_engine import AsyncOmniEngine
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.platforms import current_omni_platform
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion]
@@ -86,6 +95,23 @@ class _StepPipeline:
return DiffusionOutput(output=torch.tensor([state.step_index], dtype=torch.float32))
class _InterruptingStepPipeline(_StepPipeline):
interrupt = True
def denoise_step(self, state, **kwargs):
del state, kwargs
self.denoise_calls += 1
return None
def step_scheduler(self, state, noise_pred, **kwargs):
del state, noise_pred, kwargs
raise AssertionError("step_scheduler should not run after interrupt")
def post_decode(self, state, **kwargs):
del state, kwargs
raise AssertionError("post_decode should not run after interrupt")
class _IdentityNoiseTransformer(torch.nn.Module):
def forward(self, x: torch.Tensor, **kwargs):
del kwargs
@@ -188,6 +214,21 @@ def _make_step_request(num_inference_steps: int = 2):
)
def _assert_aborted_output(output: DiffusionOutput, request_id: str) -> None:
assert output.output is None
assert output.error is None
assert output.aborted is True
assert output.abort_message == f"Request {request_id} aborted."
def _make_engine_request(req_id: str = "req-1", num_inference_steps: int = 2) -> OmniDiffusionRequest:
return OmniDiffusionRequest(
prompts=[f"prompt-{req_id}"],
sampling_params=OmniDiffusionSamplingParams(num_inference_steps=num_inference_steps),
request_ids=[req_id],
)
def _make_runner():
runner = object.__new__(DiffusionModelRunner)
runner.vllm_config = object()
@@ -242,6 +283,18 @@ def _make_cached_scheduler_output(sched_req_id="req-1", step_id=1, finished_req_
)
def _make_engine(scheduler, execute_fn=None) -> DiffusionEngine:
engine = object.__new__(DiffusionEngine)
engine.od_config = SimpleNamespace(model_class_name="QwenImagePipeline")
engine.pre_process_func = None
engine.post_process_func = None
engine.scheduler = scheduler
engine.execute_fn = execute_fn
engine._rpc_lock = threading.RLock()
engine.abort_queue = queue.Queue()
return engine
def _expected_output_for_mode(mode: str) -> torch.Tensor:
if mode == "cfg":
return torch.tensor([[3.0]])
@@ -322,6 +375,52 @@ class TestRunner:
assert runner.pipeline.scheduler_calls == 2
assert runner.pipeline.decode_calls == 1
def test_rejects_multi_request_step_batch(self):
runner = _make_runner()
req_1 = _make_step_request()
req_2 = _make_step_request()
req_2.request_ids = ["req-2"]
scheduler_output = DiffusionSchedulerOutput(
step_id=0,
scheduled_new_reqs=[
NewRequestData(sched_req_id="req-1", req=req_1),
NewRequestData(sched_req_id="req-2", req=req_2),
],
scheduled_cached_reqs=CachedRequestData.make_empty(),
finished_req_ids=set(),
num_running_reqs=2,
num_waiting_reqs=0,
)
with pytest.raises(ValueError, match="batch_size=1"):
DiffusionModelRunner.execute_stepwise(runner, scheduler_output)
def test_rejects_missing_cached_state(self):
runner = _make_runner()
with pytest.raises(ValueError, match="Missing cached state"):
DiffusionModelRunner.execute_stepwise(runner, _make_cached_scheduler_output(sched_req_id="req-missing"))
def test_interrupt_marks_request_finished_and_clears_state(self, monkeypatch):
runner = _make_runner()
runner.pipeline = _InterruptingStepPipeline()
req = _make_step_request()
monkeypatch.setattr(model_runner_module, "set_forward_context", _noop_forward_context)
output = DiffusionModelRunner.execute_stepwise(runner, _make_scheduler_output(req, step_id=0))
assert output.req_id == "req-1"
assert output.step_index == 0
assert output.finished is True
assert output.result is not None
assert output.result.error == "stepwise denoise interrupted"
assert "req-1" not in runner.state_cache
assert runner.pipeline.prepare_calls == 1
assert runner.pipeline.denoise_calls == 1
assert runner.pipeline.scheduler_calls == 0
assert runner.pipeline.decode_calls == 0
def test_load_model_rejects_unsupported_step_execution(self, monkeypatch):
class _RequestOnlyPipeline:
pass
@@ -439,6 +538,153 @@ class TestWorker:
DiffusionWorker.execute_stepwise(worker, scheduler_output)
@pytest.mark.cpu
class TestExecutor:
"""MultiprocDiffusionExecutor.execute_step"""
def test_execute_step_passes_through_runner_output(self):
executor = object.__new__(MultiprocDiffusionExecutor)
executor._ensure_open = lambda: None
expected = RunnerOutput(req_id="req-step", step_index=1, finished=False, result=None)
executor.collective_rpc = Mock(return_value=expected)
request = _make_engine_request("req-step", num_inference_steps=2)
scheduler_output = _make_scheduler_output(request, sched_req_id="req-step")
output = MultiprocDiffusionExecutor.execute_step(executor, scheduler_output)
assert output is expected
@pytest.mark.cpu
class TestEngine:
"""Step-execution paths in DiffusionEngine.add_req_and_wait_for_response"""
@pytest.mark.parametrize(
("execute_fn", "expected_error"),
[
(
lambda _: RunnerOutput(
req_id="req-error",
step_index=1,
finished=True,
result=DiffusionOutput(error="boom"),
),
"boom",
),
(
lambda _: (_ for _ in ()).throw(RuntimeError("gpu on fire")),
"gpu on fire",
),
],
)
def test_step_engine_returns_error(self, execute_fn, expected_error):
scheduler = StepScheduler()
scheduler.initialize(Mock())
engine = _make_engine(scheduler, execute_fn=execute_fn)
output = engine.add_req_and_wait_for_response(_make_engine_request("req-error", num_inference_steps=2))
assert output.output is None
assert expected_error in output.error
def test_step_execution_completes(self):
scheduler = StepScheduler()
scheduler.initialize(Mock())
engine = _make_engine(scheduler)
request = _make_engine_request("req-step", num_inference_steps=2)
call_count = {"n": 0}
def execute_fn(_):
call_count["n"] += 1
finished = call_count["n"] == 2
return RunnerOutput(
req_id="req-step",
step_index=call_count["n"],
finished=finished,
result=(DiffusionOutput(output=torch.tensor([2.0])) if finished else None),
)
engine.execute_fn = execute_fn
output = engine.add_req_and_wait_for_response(request)
assert call_count["n"] == 2
assert output.error is None
assert torch.equal(output.output, torch.tensor([2.0]))
def test_step_abort_stops_rescheduling_after_first_step(self):
scheduler = StepScheduler()
scheduler.initialize(Mock())
engine = _make_engine(scheduler)
request = _make_engine_request("req-stop", num_inference_steps=4)
step = {"n": 0}
def execute_fn(_):
step["n"] += 1
engine.abort("req-stop")
return RunnerOutput(
req_id="req-stop",
step_index=1,
finished=False,
result=None,
)
engine.execute_fn = execute_fn
output = engine.add_req_and_wait_for_response(request)
assert step["n"] == 1
_assert_aborted_output(output, "req-stop")
def test_step_abort_after_reschedule_returns_aborted_output(self):
scheduler = StepScheduler()
scheduler.initialize(Mock())
engine = _make_engine(scheduler)
request = _make_engine_request("req-mid", num_inference_steps=4)
step = {"n": 0}
def execute_fn(sched_output):
step["n"] += 1
if step["n"] == 2:
assert sched_output == _make_cached_scheduler_output("req-mid", step_id=1)
engine.abort("req-mid")
return RunnerOutput(
req_id="req-mid",
step_index=step["n"],
finished=False,
result=None,
)
engine.execute_fn = execute_fn
output = engine.add_req_and_wait_for_response(request)
assert step["n"] == 2
_assert_aborted_output(output, "req-mid")
def test_finished_step_without_result_returns_error(self):
scheduler = StepScheduler()
scheduler.initialize(Mock())
engine = _make_engine(
scheduler,
execute_fn=lambda _: RunnerOutput(
req_id="req-missing",
step_index=1,
finished=True,
result=None,
),
)
output = engine.add_req_and_wait_for_response(_make_engine_request("req-missing", num_inference_steps=1))
assert output.output is None
assert output.error == "Diffusion execution finished without a final output."
@pytest.mark.cpu
class TestIPC:
def test_pack_unpack_runner_output_shm(self):
@@ -458,6 +704,15 @@ class TestIPC:
class TestSupportedPipelines:
"""Step-execution protocol checks for supported pipelines."""
def test_default_stage_config_includes_step_execution(self):
stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg(
{
"step_execution": True,
}
)[0]
assert stage_cfg["engine_args"]["step_execution"] is True
def test_qwen_image_supports_step_execution(self):
from vllm_omni.diffusion.models.interface import SupportsStepExecution, supports_step_execution
from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import QwenImagePipeline
@@ -66,7 +66,9 @@ def _make_engine(num_gpus: int = 1):
sched.initialize(Mock())
engine.scheduler = sched
engine.executor = executor
engine._rpc_lock = threading.Lock()
engine._rpc_lock = threading.RLock()
engine.abort_queue = queue.Queue()
engine.execute_fn = executor.execute_request
return engine, executor, req_q, res_q
@@ -80,7 +82,7 @@ def _start_worker(req_q, res_q, count=2):
req = req_q.get(timeout=10)
method = req.get("method", "")
args = req.get("args", ())
if method == "generate" and args and hasattr(args[0], "request_ids"):
if method in {"generate", "execute_model"} and args and hasattr(args[0], "request_ids"):
tag = f"result_for_{args[0].request_ids[0]}"
elif args:
tag = f"result_for_{args[0]}"
@@ -116,11 +118,11 @@ def _inject_interleave(executor):
return a_enqueued, b_complete
# ──────────────────── bug-reproduction: concurrent add_req ────────────────
# ───────────────── concurrent request execution ─────────────────
class TestConcurrentAddReqBug:
"""Two concurrent ``add_req_and_wait_for_response()`` calls swap results."""
class TestConcurrentRequestExecution:
"""Concurrent request execution should not swap results."""
def test_results_are_correctly_routed(self):
engine, executor, req_q, res_q = _make_engine()
@@ -151,11 +153,11 @@ class TestConcurrentAddReqBug:
assert results["B"].error == "result_for_B"
# ──────────────── bug-reproduction: concurrent collective_rpc ─────────────
# ───────────────── concurrent collective RPC ─────────────────
class TestConcurrentCollectiveRpcBug:
"""Two concurrent ``collective_rpc()`` calls swap results."""
class TestConcurrentCollectiveRpc:
"""Concurrent ``collective_rpc()`` calls should not swap results."""
def test_results_are_correctly_routed(self):
engine, executor, req_q, res_q = _make_engine()
@@ -192,11 +194,11 @@ class TestConcurrentCollectiveRpcBug:
assert results["B"].error == "result_for_call_B"
# ──────── bug-reproduction: add_req vs collective_rpc concurrently ────────
# ──────────── concurrent request execution and collective RPC ────────────
class TestConcurrentAddReqVsCollectiveRpcBug:
"""``add_req`` and ``collective_rpc`` running concurrently swap results."""
class TestConcurrentRequestExecutionAndCollectiveRpc:
"""Request execution and ``collective_rpc()`` should not swap results."""
def test_results_are_correctly_routed(self):
engine, executor, req_q, res_q = _make_engine()
@@ -205,7 +207,7 @@ class TestConcurrentAddReqVsCollectiveRpcBug:
results: dict[str, object] = {}
def _a(): # add_req path
def _a(): # request execution path
results["A"] = engine.add_req_and_wait_for_response(_mock_request("A"))
def _b(): # collective_rpc path
@@ -230,10 +232,10 @@ class TestConcurrentAddReqVsCollectiveRpcBug:
assert results["B"].error == "result_for_call_B"
# ─────────────── backward-compatibility (serial) tests ────────────────────
# ─────────────────────── serial operation coverage ───────────────────────
class TestSerialOperations:
class TestSerialEngineOperations:
"""Verify correct behaviour for single-threaded (serial) usage.
These tests must pass both **before** and **after** any concurrency fix
@@ -385,18 +387,18 @@ class TestCollectiveRpcTimeoutWhileLockHeld:
executor._result_mq.dequeue = _hanging_dequeue
# Thread running add_req — acquires the lock, enqueues, then
# Thread running request execution — acquires the lock, enqueues, then
# blocks on dequeue forever (worker hang).
def _stalled_add_req():
def _stalled_request_execution():
try:
engine.add_req_and_wait_for_response(_mock_request("stalled"))
except Exception:
pass
t = threading.Thread(target=_stalled_add_req, daemon=True)
t = threading.Thread(target=_stalled_request_execution, daemon=True)
t.start()
# Wait until add_req is truly inside the lock and blocking.
# Wait until request execution is truly inside the lock and blocking.
add_req_blocked.wait(5)
# collective_rpc should time out at lock acquisition, not hang.
+24
View File
@@ -47,6 +47,19 @@ def pytest_addoption(parser):
help="Balanced sample count per GEdit task group",
)
group.addoption("--accuracy-workers", action="store", type=int, default=1, help="Worker count for accuracy benches")
group.addoption(
"--wan22-i2v-image-source",
action="store",
default=None,
help="Image source for Wan2.2 I2V accuracy tests. Can be local path or remote URL.",
)
group.addoption(
"--wan22-i2v-online-timeout-seconds",
action="store",
type=int,
default=1200,
help="Online serving timeout in seconds for Wan2.2 I2V accuracy tests.",
)
def _hf_cache_root() -> Path:
@@ -142,6 +155,17 @@ def accuracy_workers(request: pytest.FixtureRequest) -> int:
return int(request.config.getoption("accuracy_workers"))
@pytest.fixture(scope="session")
def wan22_i2v_image_source(request: pytest.FixtureRequest) -> str | None:
value = request.config.getoption("wan22_i2v_image_source")
return str(value) if value else None
@pytest.fixture(scope="session")
def wan22_i2v_online_timeout_seconds(request: pytest.FixtureRequest) -> int:
return int(request.config.getoption("wan22_i2v_online_timeout_seconds"))
@pytest.fixture(scope="session")
def gebench_samples_per_type(request: pytest.FixtureRequest) -> int:
return int(request.config.getoption("gebench_samples_per_type"))
Whitespace-only changes.
@@ -0,0 +1,165 @@
from __future__ import annotations
import argparse
import base64
import json
from io import BytesIO
from pathlib import Path
import requests
import torch
from diffusers import UniPCMultistepScheduler, WanImageToVideoPipeline
from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
from diffusers.utils import export_to_video, load_image
from PIL import Image
from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import BOUNDARY_RATIO
def _parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Run Wan2.2 I2V diffusers offline generation.")
parser.add_argument("--model", required=True)
parser.add_argument("--image-source", required=True)
parser.add_argument("--prompt", required=True)
parser.add_argument("--negative-prompt", required=True)
parser.add_argument("--size", required=True)
parser.add_argument("--fps", type=int, required=True)
parser.add_argument("--num-frames", type=int, required=True)
parser.add_argument("--guidance-scale", type=float, required=True)
parser.add_argument("--guidance-scale-2", type=float, required=True)
parser.add_argument("--flow-shift", type=float, required=True)
parser.add_argument("--num-inference-steps", type=int, required=True)
parser.add_argument("--seed", type=int, required=True)
parser.add_argument("--output", required=True)
parser.add_argument("--metadata-output", required=True)
return parser.parse_args()
def _parse_size(size: str) -> tuple[int, int]:
width_str, height_str = size.lower().split("x", 1)
return int(width_str), int(height_str)
class _IdentityFtfy:
@staticmethod
def fix_text(text: str) -> str:
return text
def _ensure_wan_ftfy_fallback() -> None:
if not hasattr(wan_i2v_module, "ftfy"):
wan_i2v_module.ftfy = _IdentityFtfy()
def _offline_cuda_device() -> torch.device:
return torch.device("cuda:0")
def _load_input_image(source: str) -> Image.Image:
if source.startswith("data:image"):
_, encoded = source.split(",", 1)
image = Image.open(BytesIO(base64.b64decode(encoded)))
image.load()
return image.convert("RGB")
source_path = Path(source)
if source_path.exists():
image = Image.open(source_path)
image.load()
return image.convert("RGB")
image = load_image(source)
if isinstance(image, Image.Image):
image.load()
return image.convert("RGB")
response = requests.get(source, timeout=60)
response.raise_for_status()
image = Image.open(BytesIO(response.content))
image.load()
return image.convert("RGB")
def _resize_to_target(image: Image.Image, *, width: int, height: int) -> Image.Image:
return image.resize((width, height), Image.Resampling.LANCZOS)
def _configure_scheduler(pipe: WanImageToVideoPipeline, *, flow_shift: float) -> None:
pipe.scheduler = UniPCMultistepScheduler.from_config(
pipe.scheduler.config,
flow_shift=flow_shift,
)
def _write_metadata(
path: Path,
*,
args: argparse.Namespace,
width: int,
height: int,
frame_count: int,
) -> None:
payload = {
"model": args.model,
"image_source": args.image_source,
"size": args.size,
"width": width,
"height": height,
"fps": args.fps,
"num_frames": args.num_frames,
"actual_frame_count": frame_count,
"guidance_scale": args.guidance_scale,
"guidance_scale_2": args.guidance_scale_2,
"boundary_ratio": BOUNDARY_RATIO,
"flow_shift": args.flow_shift,
"num_inference_steps": args.num_inference_steps,
"seed": args.seed,
"world_size": 1,
}
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(payload, ensure_ascii=False, indent=2), encoding="utf-8")
def main() -> int:
args = _parse_args()
device = _offline_cuda_device()
torch.cuda.set_device(device)
_ensure_wan_ftfy_fallback()
pipe = WanImageToVideoPipeline.from_pretrained(args.model, torch_dtype=torch.bfloat16)
pipe.register_to_config(boundary_ratio=BOUNDARY_RATIO)
_configure_scheduler(pipe, flow_shift=args.flow_shift)
pipe.to(device)
pipe.set_progress_bar_config(disable=False)
input_image = _load_input_image(args.image_source)
width, height = _parse_size(args.size)
resized_image = _resize_to_target(input_image, width=width, height=height)
generator = torch.Generator(device=device.type).manual_seed(args.seed)
frames = pipe(
image=resized_image,
prompt=args.prompt,
negative_prompt=args.negative_prompt,
height=height,
width=width,
num_frames=args.num_frames,
guidance_scale=args.guidance_scale,
guidance_scale_2=args.guidance_scale_2,
num_inference_steps=args.num_inference_steps,
generator=generator,
).frames[0]
output_path = Path(args.output)
metadata_path = Path(args.metadata_output)
output_path.parent.mkdir(parents=True, exist_ok=True)
export_to_video(frames, str(output_path), fps=args.fps)
_write_metadata(metadata_path, args=args, width=width, height=height, frame_count=len(frames))
if hasattr(pipe, "maybe_free_model_hooks"):
pipe.maybe_free_model_hooks()
return 0
if __name__ == "__main__":
raise SystemExit(main())
@@ -0,0 +1,672 @@
from __future__ import annotations
import json
import os
import re
import shutil
import subprocess
import sys
import time
from base64 import b64decode, b64encode
from fractions import Fraction
from hashlib import sha1
from io import BytesIO
from mimetypes import guess_type
from pathlib import Path
from typing import Any
from urllib.parse import urlparse
import pytest
import requests
import torch
from diffusers import UniPCMultistepScheduler
from PIL import Image
from tests.conftest import OmniServerParams
from tests.e2e.accuracy.wan22_i2v.run_wan22_i2v_diffusers_cp import (
_configure_scheduler,
_ensure_wan_ftfy_fallback,
_IdentityFtfy,
_offline_cuda_device,
_resize_to_target,
)
from tests.e2e.accuracy.wan22_i2v.wan22_i2v_video_similarity_common import (
FLOW_SHIFT,
FPS,
GUIDANCE_SCALE,
GUIDANCE_SCALE_2,
HEIGHT,
MODEL_NAME,
NEGATIVE_PROMPT,
NUM_FRAMES,
NUM_INFERENCE_STEPS,
PROMPT,
PSNR_THRESHOLD,
RABBIT_IMAGE_URL,
SEED,
SIZE,
SSIM_THRESHOLD,
WIDTH,
)
from tests.utils import hardware_test
def test_parse_video_metadata_extracts_dimensions_and_fps() -> None:
payload = {
"streams": [
{
"width": 832,
"height": 480,
"avg_frame_rate": "16/1",
"nb_read_frames": "93",
}
]
}
metadata = _parse_video_metadata(payload)
assert metadata["width"] == 832
assert metadata["height"] == 480
assert metadata["fps"] == 16.0
assert metadata["frame_count"] == 93
def test_parse_ssim_summary_extracts_all_score() -> None:
output = """
[Parsed_ssim_0 @ 000001] SSIM Y:0.971903 (15.512007) U:0.965077 (14.569044) V:0.962414 (14.252637) All:0.968311 (15.035654)
"""
assert _parse_ssim_score(output) == 0.968311
def test_parse_psnr_summary_extracts_average_score() -> None:
output = """
[Parsed_psnr_0 @ 000001] PSNR y:32.670157 u:31.844621 v:31.513839 average:32.148004 min:31.563744 max:33.201457
"""
assert _parse_psnr_score(output) == 32.148004
def test_build_diffusers_command_uses_python_runner_path(tmp_path: Path) -> None:
runner_path = tmp_path / "run_wan22_i2v_diffusers_cp.py"
command = _build_diffusers_command(
runner_path=runner_path,
image_source=RABBIT_IMAGE_URL,
output_path=tmp_path / "offline.mp4",
metadata_path=tmp_path / "offline.json",
)
assert command[:2] == [
sys.executable,
str(runner_path),
]
assert "--output" in command
assert "--metadata-output" in command
def test_resolve_image_source_prefers_existing_local_path(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None:
image_path = tmp_path / "rabbit.png"
Image.new("RGB", (8, 8), color=(255, 128, 64)).save(image_path)
assert _resolve_image_source(str(image_path)) == str(image_path.resolve())
def test_build_online_image_reference_uses_data_url_for_local_path(tmp_path: Path) -> None:
image_path = tmp_path / "rabbit.png"
Image.new("RGB", (4, 2), color=(10, 20, 30)).save(image_path)
reference = _build_online_image_reference(str(image_path))
assert reference.startswith("data:image/png;base64,")
encoded = reference.split(",", 1)[1]
assert len(b64decode(encoded)) > 0
def test_send_video_request_with_timeout_uses_requested_timeout(monkeypatch: pytest.MonkeyPatch) -> None:
class DummyResponse:
def raise_for_status(self) -> None:
return None
def json(self) -> dict[str, str]:
return {"id": "video_123"}
class DummyClient:
def __init__(self) -> None:
self.run_level = "core_model"
self.wait_args: tuple[str, int, int] | None = None
def _build_url(self, path: str) -> str:
return f"http://localhost:8000{path}"
def _wait_until_video_completed(
self,
video_id: str,
poll_interval_seconds: int = 2,
timeout_seconds: int = 300,
) -> None:
self.wait_args = (video_id, poll_interval_seconds, timeout_seconds)
def _download_video_content(self, video_id: str) -> bytes:
return b"video-bytes"
monkeypatch.setattr(requests, "post", lambda *args, **kwargs: DummyResponse())
client = DummyClient()
result = _send_video_request_with_timeout(
client,
{
"form_data": {"prompt": "test"},
"image_reference": RABBIT_IMAGE_URL,
},
timeout_seconds=1200,
)
assert result == b"video-bytes"
assert client.wait_args == ("video_123", 2, 1200)
def test_online_timeout_defaults_to_1200() -> None:
assert _online_timeout_seconds(None) == 1200
def test_artifact_dir_is_under_repo_result_folder(tmp_path: Path) -> None:
artifact_dir = _artifact_dir(str(tmp_path / "rabbit.png"))
assert artifact_dir.parent == Path(__file__).parent / "result"
assert artifact_dir.name.startswith("rabbit-")
def test_offline_cuda_device_uses_indexed_cuda_device() -> None:
assert _offline_cuda_device() == torch.device("cuda:0")
def test_ensure_wan_ftfy_fallback_sets_identity(monkeypatch: pytest.MonkeyPatch) -> None:
from diffusers.pipelines.wan import pipeline_wan_i2v as wan_i2v_module
monkeypatch.delattr(wan_i2v_module, "ftfy", raising=False)
_ensure_wan_ftfy_fallback()
assert hasattr(wan_i2v_module, "ftfy")
assert isinstance(wan_i2v_module.ftfy, _IdentityFtfy)
assert wan_i2v_module.ftfy.fix_text("abc") == "abc"
def test_resize_to_target_matches_requested_dimensions() -> None:
image = Image.new("RGB", (100, 50), color=(1, 2, 3))
resized = _resize_to_target(image, width=832, height=480)
assert resized.size == (832, 480)
def test_configure_scheduler_switches_to_unipc() -> None:
class DummyPipe:
def __init__(self) -> None:
self.scheduler = UniPCMultistepScheduler(num_train_timesteps=1000, flow_shift=1.0)
pipe = DummyPipe()
_configure_scheduler(pipe, flow_shift=5.0)
assert isinstance(pipe.scheduler, UniPCMultistepScheduler)
assert pipe.scheduler.config.flow_shift == 5.0
REPO_ROOT = Path(__file__).resolve().parents[4]
WORKSPACE_ROOT = REPO_ROOT.parent
RUNNER_PATH = Path(__file__).with_name("run_wan22_i2v_diffusers_cp.py")
RESULT_ROOT = Path(__file__).parent / "result"
VIDEO_TIMEOUT_SECONDS = 60 * 60
_SSIM_RE = re.compile(r"All:(?P<score>[0-9.]+)")
_PSNR_RE = re.compile(r"average:(?P<score>[0-9.]+)")
SERVER_CASES = [
pytest.param(
OmniServerParams(
model=MODEL_NAME,
server_args=[
"--usp",
"2",
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
use_omni=True,
),
id="wan22_i2v_usp2_hsdp2",
)
]
def _parse_video_metadata(payload: dict[str, Any]) -> dict[str, int | float]:
streams = payload.get("streams")
if not isinstance(streams, list) or not streams:
raise ValueError(f"ffprobe payload did not include video streams: {payload}")
stream = streams[0]
width = int(stream["width"])
height = int(stream["height"])
fps = float(Fraction(stream["avg_frame_rate"]))
frame_count_value = stream.get("nb_read_frames") or stream.get("nb_frames")
if frame_count_value is None:
raise ValueError(f"ffprobe payload did not include frame count: {payload}")
frame_count = int(frame_count_value)
return {
"width": width,
"height": height,
"fps": fps,
"frame_count": frame_count,
}
def _parse_ssim_score(output: str) -> float:
match = _SSIM_RE.search(output)
if match is None:
raise ValueError(f"Could not parse SSIM score from ffmpeg output:\n{output}")
return float(match.group("score"))
def _parse_psnr_score(output: str) -> float:
match = _PSNR_RE.search(output)
if match is None:
raise ValueError(f"Could not parse PSNR score from ffmpeg output:\n{output}")
return float(match.group("score"))
def _probe_binary(binary: str) -> str:
resolved = shutil.which(binary)
if resolved is None:
pytest.skip(f"{binary} is required for Wan2.2 video similarity e2e test.")
return resolved
def _build_diffusers_command(
*,
runner_path: Path,
image_source: str,
output_path: Path,
metadata_path: Path,
) -> list[str]:
return [
sys.executable,
str(runner_path),
"--model",
MODEL_NAME,
"--image-source",
image_source,
"--prompt",
PROMPT,
"--negative-prompt",
NEGATIVE_PROMPT,
"--size",
SIZE,
"--fps",
str(FPS),
"--num-frames",
str(NUM_FRAMES),
"--guidance-scale",
str(GUIDANCE_SCALE),
"--guidance-scale-2",
str(GUIDANCE_SCALE_2),
"--flow-shift",
str(FLOW_SHIFT),
"--num-inference-steps",
str(NUM_INFERENCE_STEPS),
"--seed",
str(SEED),
"--output",
str(output_path),
"--metadata-output",
str(metadata_path),
]
def _runner_env() -> dict[str, str]:
env = os.environ.copy()
pythonpath_parts = [
str(REPO_ROOT),
str(WORKSPACE_ROOT / "diffusers" / "src"),
]
existing_pythonpath = env.get("PYTHONPATH")
if existing_pythonpath:
pythonpath_parts.append(existing_pythonpath)
env["PYTHONPATH"] = os.pathsep.join(pythonpath_parts)
env.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
return env
def _resolve_image_source(configured: str | None) -> str:
configured = configured or RABBIT_IMAGE_URL
candidate = Path(configured)
if candidate.exists():
return str(candidate.resolve())
return configured
def _is_remote_image_source(source: str) -> bool:
return source.startswith("http://") or source.startswith("https://")
def _validate_image_source(source: str) -> None:
if _is_remote_image_source(source):
response = requests.get(source, timeout=60)
response.raise_for_status()
return
image_path = Path(source)
if not image_path.exists():
raise FileNotFoundError(f"Local image source does not exist: {image_path}")
def _build_online_image_reference(source: str) -> str:
if _is_remote_image_source(source) or source.startswith("data:image"):
return source
image_path = Path(source)
mime_type = guess_type(image_path.name)[0] or "image/png"
encoded = b64encode(image_path.read_bytes()).decode("ascii")
return f"data:{mime_type};base64,{encoded}"
def _probe_video(path: Path) -> dict[str, int | float]:
_probe_binary("ffprobe")
result = subprocess.run(
[
"ffprobe",
"-v",
"error",
"-select_streams",
"v:0",
"-count_frames",
"-show_entries",
"stream=width,height,avg_frame_rate,nb_read_frames",
"-of",
"json",
str(path),
],
capture_output=True,
text=True,
check=True,
)
return _parse_video_metadata(json.loads(result.stdout))
def _run_ffmpeg_similarity(filter_name: str, first: Path, second: Path) -> str:
_probe_binary("ffmpeg")
result = subprocess.run(
[
"ffmpeg",
"-hide_banner",
"-i",
str(first),
"-i",
str(second),
"-lavfi",
f"[0:v][1:v]{filter_name}",
"-f",
"null",
"-",
],
capture_output=True,
text=True,
check=True,
)
return result.stderr
def _online_timeout_seconds(configured: int | None) -> int:
return int(configured) if configured is not None else 1200
def _artifact_dir(image_source: str) -> Path:
if _is_remote_image_source(image_source):
source_name = Path(urlparse(image_source).path).stem or "remote"
else:
source_name = Path(image_source).stem or "local"
safe_name = re.sub(r"[^A-Za-z0-9_.-]+", "_", source_name)
digest = sha1(image_source.encode("utf-8")).hexdigest()[:8]
return RESULT_ROOT / f"{safe_name}-{digest}"
def _artifact_paths(image_source: str) -> tuple[Path, Path, Path]:
artifact_dir = _artifact_dir(image_source)
artifact_dir.mkdir(parents=True, exist_ok=True)
return (
artifact_dir / "online.mp4",
artifact_dir / "offline.mp4",
artifact_dir / "offline_metadata.json",
)
def _send_video_request_with_timeout(
openai_client,
request_config: dict[str, Any],
*,
timeout_seconds: int,
) -> bytes:
form_data = request_config.get("form_data")
if not isinstance(form_data, dict):
raise ValueError("Video request_config must contain 'form_data'")
if not form_data.get("prompt"):
raise ValueError("Video request_config['form_data'] must contain 'prompt'")
normalized_form_data = {key: str(value) for key, value in form_data.items() if value is not None}
files: dict[str, tuple[str, BytesIO, str]] = {}
image_reference = request_config.get("image_reference")
if image_reference:
if image_reference.startswith("data:image"):
header, encoded = image_reference.split(",", 1)
content_type = header.split(";")[0].removeprefix("data:")
extension = content_type.split("/")[-1]
file_data = b64decode(encoded)
files["input_reference"] = (
f"reference.{extension}",
BytesIO(file_data),
content_type,
)
else:
normalized_form_data["image_reference"] = json.dumps({"image_url": image_reference})
start_time = time.perf_counter()
create_url = openai_client._build_url("/v1/videos")
response = requests.post(
create_url,
data=normalized_form_data,
files=files,
headers={"Accept": "application/json"},
timeout=60,
)
response.raise_for_status()
job_data = response.json()
video_id = job_data["id"]
openai_client._wait_until_video_completed(video_id, timeout_seconds=timeout_seconds)
video_content = openai_client._download_video_content(video_id)
print(f"online_video_e2e_latency_s={time.perf_counter() - start_time:.3f}")
return video_content
def _generate_online_video(
*,
omni_server,
openai_client,
image_source: str,
online_timeout_seconds: int,
) -> Path:
online_path, _, _ = _artifact_paths(image_source)
request_config = {
"model": omni_server.model,
"form_data": {
"prompt": PROMPT,
"negative_prompt": NEGATIVE_PROMPT,
"size": SIZE,
"fps": FPS,
"num_frames": NUM_FRAMES,
"guidance_scale": GUIDANCE_SCALE,
"guidance_scale_2": GUIDANCE_SCALE_2,
"flow_shift": FLOW_SHIFT,
"num_inference_steps": NUM_INFERENCE_STEPS,
"seed": SEED,
},
"image_reference": _build_online_image_reference(image_source),
}
online_video_bytes = _send_video_request_with_timeout(
openai_client,
request_config,
timeout_seconds=_online_timeout_seconds(online_timeout_seconds),
)
online_path.write_bytes(online_video_bytes)
return online_path
def _generate_offline_video(*, image_source: str) -> tuple[Path, Path]:
_, offline_path, offline_metadata_path = _artifact_paths(image_source)
command = _build_diffusers_command(
runner_path=RUNNER_PATH,
image_source=image_source,
output_path=offline_path,
metadata_path=offline_metadata_path,
)
subprocess.run(
command,
cwd=REPO_ROOT,
env=_runner_env(),
check=True,
timeout=VIDEO_TIMEOUT_SECONDS,
)
return offline_path, offline_metadata_path
@pytest.mark.advanced_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=1)
def test_wan22_i2v_diffusers_offline_generates_video(
wan22_i2v_image_source: str | None,
) -> None:
if not torch.cuda.is_available():
pytest.skip("Wan2.2 I2V diffusers offline test requires CUDA.")
_probe_binary("ffprobe")
if not RUNNER_PATH.exists():
raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")
image_source = _resolve_image_source(wan22_i2v_image_source)
_validate_image_source(image_source)
offline_path, offline_metadata_path = _generate_offline_video(image_source=image_source)
assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
offline_metadata = _probe_video(offline_path)
assert offline_metadata["width"] == WIDTH
assert offline_metadata["height"] == HEIGHT
assert offline_metadata["fps"] == float(FPS)
assert offline_metadata["frame_count"] == NUM_FRAMES
@pytest.mark.advanced_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=2)
@pytest.mark.parametrize("omni_server", SERVER_CASES, indirect=True)
def test_wan22_i2v_online_serving_generates_video(
omni_server,
openai_client,
wan22_i2v_image_source: str | None,
wan22_i2v_online_timeout_seconds: int,
) -> None:
if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
_probe_binary("ffprobe")
image_source = _resolve_image_source(wan22_i2v_image_source)
_validate_image_source(image_source)
online_path = _generate_online_video(
omni_server=omni_server,
openai_client=openai_client,
image_source=image_source,
online_timeout_seconds=wan22_i2v_online_timeout_seconds,
)
assert online_path.exists(), f"Expected online video artifact at {online_path}"
online_metadata = _probe_video(online_path)
assert online_metadata["width"] == WIDTH
assert online_metadata["height"] == HEIGHT
assert online_metadata["fps"] == float(FPS)
assert online_metadata["frame_count"] == NUM_FRAMES
@pytest.mark.advanced_model
@pytest.mark.benchmark
@pytest.mark.diffusion
@hardware_test(res={"cuda": "H100"}, num_cards=2)
def test_wan22_i2v_serving_matches_diffusers_video_similarity(
wan22_i2v_image_source: str | None,
) -> None:
if not torch.cuda.is_available() or torch.cuda.device_count() < 2:
pytest.skip("Wan2.2 I2V similarity e2e test requires >= 2 CUDA GPUs.")
_probe_binary("ffmpeg")
_probe_binary("ffprobe")
if not RUNNER_PATH.exists():
raise AssertionError(f"Offline diffusers runner does not exist: {RUNNER_PATH}")
image_source = _resolve_image_source(wan22_i2v_image_source)
_validate_image_source(image_source)
online_path, offline_path, offline_metadata_path = _artifact_paths(image_source)
if not online_path.exists():
pytest.skip(f"Missing online artifact from prerequisite test: {online_path}")
if not offline_path.exists() or not offline_metadata_path.exists():
pytest.skip(f"Missing offline artifacts from prerequisite test: {offline_path}, {offline_metadata_path}")
assert online_path.exists(), f"Expected online video artifact at {online_path}"
assert offline_path.exists(), f"Expected offline video artifact at {offline_path}"
assert offline_metadata_path.exists(), f"Expected offline metadata artifact at {offline_metadata_path}"
online_metadata = _probe_video(online_path)
offline_metadata = _probe_video(offline_path)
assert online_metadata == offline_metadata, (
f"Video metadata mismatch:\n"
f"online={online_metadata}\n"
f"offline={offline_metadata}\n"
f"online_path={online_path}\n"
f"offline_path={offline_path}"
)
assert online_metadata["width"] == WIDTH
assert online_metadata["height"] == HEIGHT
assert online_metadata["fps"] == float(FPS)
assert online_metadata["frame_count"] == NUM_FRAMES
ssim_output = _run_ffmpeg_similarity("ssim", online_path, offline_path)
psnr_output = _run_ffmpeg_similarity("psnr", online_path, offline_path)
ssim_score = _parse_ssim_score(ssim_output)
psnr_score = _parse_psnr_score(psnr_output)
print("wan22_i2v_similarity metrics:")
print(
" SSIM:"
f" value={ssim_score:.6f},"
f" threshold>={SSIM_THRESHOLD:.6f},"
" range=[0, 1],"
" higher_is_better=True,"
" interpretation=structural_similarity"
)
print(
" PSNR:"
f" value={psnr_score:.6f} dB,"
f" threshold>={PSNR_THRESHOLD:.6f} dB,"
" range=[0, +inf),"
" higher_is_better=True,"
" interpretation=pixel_error_in_decibels"
)
print(f"online_video={online_path}")
print(f"offline_video={offline_path}")
print(f"offline_metadata={offline_metadata_path}")
assert ssim_score >= SSIM_THRESHOLD, (
f"SSIM below threshold: got {ssim_score:.6f}, expected >= {SSIM_THRESHOLD:.6f}. "
f"online={online_path} offline={offline_path}"
)
assert psnr_score >= PSNR_THRESHOLD, (
f"PSNR below threshold: got {psnr_score:.6f}, expected >= {PSNR_THRESHOLD:.6f}. "
f"online={online_path} offline={offline_path}"
)
@@ -0,0 +1,24 @@
from __future__ import annotations
RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png"
MODEL_NAME = "Wan-AI/Wan2.2-I2V-A14B-Diffusers"
SIZE = "832x480"
WIDTH = 832
HEIGHT = 480
FPS = 16
NUM_FRAMES = 93
GUIDANCE_SCALE = 3.5
GUIDANCE_SCALE_2 = 3.5
FLOW_SHIFT = 5.0
NUM_INFERENCE_STEPS = 50
SEED = 42
BOUNDARY_RATIO = 0.875
SSIM_THRESHOLD = 0.94
PSNR_THRESHOLD = 28.0
PROMPT = """一只棕色野兔的正面特写镜头,采用低角度仰拍视角,营造亲密而庄严的视觉冲击。兔子一双圆润漆黑的大眼睛直视镜头深处,眼神中交织着野生动物的警觉与一丝难以言喻的温柔好奇,仿佛在与观者建立跨越物种的静默对话。它毛色呈现层次丰富的棕褐渐变,从浅奶油色腹部过渡到深棕背部,每根毛发纹理清晰可辨,在侧光下泛着丝绸般的光泽。细长洁白的胡须共有三对,随呼吸节奏微微颤动,偶尔因捕捉气流信息而轻轻摇摆。
两只标志性的长耳完全竖立,耳廓外侧覆盖短密棕毛,内侧则露出粉嫩的血管网络,薄如蝉翼的皮肤下血液流动隐约可见,耳朵以细微幅度不时转动,精准定位声源方向。背景是一片澄澈的蔚蓝天空,形态蓬松的白色积云以缓慢速度横向漂移,云影在兔子头顶交替变化,光线随之明暗流转。晴朗天气的明媚阳光从画面左上方45度角倾泻而下,在兔脸右侧形成柔和的伦勃朗式阴影,强化了面部立体感和皮毛质感。
兔子湿润的黑鼻子持续进行每秒三至四次的快速抽动,这是它们感知化学信号的本能动作,粉色三瓣嘴随之轻启,露出正在反刍的洁白门齿,下颌以稳定节奏左右研磨。摄影采用大光圈浅景深,焦点牢牢锁定在兔子双眼连线所在的焦平面,背景天空和远景绿色植被虚化成圆润的彩色光斑,前景几根嫩绿草叶闯入画面边缘,以缓慢弧线随风摇曳,暗示着和煦的春日微风。
整个场景弥漫着宁静致远的田园诗意,色彩温暖饱和,充满生命力。兔子在持续五秒的对视后,以典型 lagomorph 特征完成一次完整的瞬膜眨眼——第三眼睑从内侧横向滑过眼球,继而缓缓歪头向右十五度,这个行为在动物行为学中代表认知加工和好奇心表达,耳朵随之向同一方向倾斜,最终恢复正视姿态,胡须舒展,完成这段短暂而珍贵的自然纪录。"""
NEGATIVE_PROMPT = """色调艳丽,过曝,静态,细节模糊不清,字幕,风格,作品,画作,画面,静止,整体发灰,最差质量,低质量,JPEG压缩残留,丑陋的,残缺的,多余的手指,画得不好的手部,画得不好的脸部,畸形的,毁容的,形态畸形的肢体,手指融合,静止不动的画面,杂乱的背景,三条腿,背景人很多,倒着走"""
@@ -79,19 +79,17 @@ def _find_free_port() -> int:
return port
def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
"""Configure sampling parameters for Bagel img2img generation.
Args:
omni: The Omni instance to get default params from.
max_tokens: Maximum tokens for the first stage.
num_inference_steps: Number of inference steps for the diffusion stage.
Returns:
Configured sampling params list.
"""
params_list = omni.default_sampling_params_list
params_list[0].max_tokens = max_tokens # type: ignore
if len(params_list) > 1:
params_list[1].num_inference_steps = num_inference_steps # type: ignore
params_list[1].extra_args = { # type: ignore
@@ -80,19 +80,17 @@ def _find_free_port() -> int:
return port
def _configure_sampling_params(omni: Omni, max_tokens: int = 1, num_inference_steps: int = 15) -> list:
def _configure_sampling_params(omni: Omni, num_inference_steps: int = 15) -> list:
"""Configure sampling parameters for Bagel text2img generation.
Args:
omni: The Omni instance to get default params from.
max_tokens: Maximum tokens for the first stage.
num_inference_steps: Number of inference steps for the diffusion stage.
Returns:
Configured sampling params list.
"""
params_list = omni.default_sampling_params_list
params_list[0].max_tokens = max_tokens # type: ignore
if len(params_list) > 1:
params_list[1].num_inference_steps = num_inference_steps # type: ignore
params_list[1].extra_args = { # type: ignore
@@ -231,6 +229,24 @@ def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool:
return False
def _is_mooncake_master_available() -> bool:
"""Check if mooncake_master binary is present and can actually execute."""
import shutil
binary = shutil.which("mooncake_master")
if binary is None:
return False
try:
result = subprocess.run(
[binary, "--help"],
capture_output=True,
timeout=5,
)
return result.returncode != 127
except (subprocess.TimeoutExpired, OSError):
return True
def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None:
"""Clean up any existing mooncake_master processes.
@@ -294,6 +310,8 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str:
@hardware_test(res={"cuda": "H100"})
def test_bagel_text2img_mooncake_connector(run_level):
"""Test Bagel text2img with Mooncake connector for inter-stage communication."""
if not _is_mooncake_master_available():
pytest.skip("mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)")
MOONCAKE_HOST = "127.0.0.1"
MOONCAKE_RPC_PORT = _find_free_port()
MOONCAKE_HTTP_PORT = _find_free_port()
@@ -0,0 +1,127 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""E2E tests for FLUX AutoRound W4A16 quantized inference.
These tests require:
- A CUDA GPU
- The quantized model checkpoint (vllm-project-org/FLUX.1-dev-AutoRound-w4a16)
"""
import gc
import sys
from pathlib import Path
import pytest
import torch
from vllm.distributed.parallel_state import cleanup_dist_env_and_memory
from tests.utils import DeviceMemoryMonitor, hardware_test
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
from vllm_omni.outputs import OmniRequestOutput
from vllm_omni.platforms import current_omni_platform
# ruff: noqa: E402
REPO_ROOT = Path(__file__).resolve().parents[2]
if str(REPO_ROOT) not in sys.path:
sys.path.insert(0, str(REPO_ROOT))
from vllm_omni import Omni
QUANTIZED_MODEL = "vllm-project-org/FLUX.1-dev-AutoRound-w4a16"
BASELINE_MODEL = "black-forest-labs/FLUX.1-dev"
# Allow overriding via environment for local testing
import os as _os
QUANTIZED_MODEL = _os.environ.get("FLUX_AUTOROUND_MODEL", QUANTIZED_MODEL)
BASELINE_MODEL = _os.environ.get("FLUX_BASELINE_MODEL", BASELINE_MODEL)
# Small resolution to keep GPU memory & time manageable
HEIGHT = 256
WIDTH = 256
NUM_STEPS = 2 # minimal for smoke-test
def _generate_image(model_name: str, **extra_kwargs) -> tuple[list, float]:
"""Load a FLUX model, generate one image, return (images, peak_memory_mb)."""
gc.collect()
current_omni_platform.empty_cache()
device_index = current_omni_platform.current_device()
current_omni_platform.reset_peak_memory_stats()
monitor = DeviceMemoryMonitor(device_index=device_index, interval=0.02)
monitor.start()
m = Omni(model=model_name, enforce_eager=True, **extra_kwargs)
current_omni_platform.reset_peak_memory_stats()
outputs = m.generate(
"a photo of a cat sitting on a laptop keyboard",
OmniDiffusionSamplingParams(
height=HEIGHT,
width=WIDTH,
num_inference_steps=NUM_STEPS,
guidance_scale=0.0,
generator=torch.Generator(device=current_omni_platform.device_type).manual_seed(42),
),
)
peak = monitor.peak_used_mb
monitor.stop()
first_output = outputs[0]
assert first_output.final_output_type == "image"
req_out = first_output.request_output
assert isinstance(req_out, OmniRequestOutput) and hasattr(req_out, "images")
images = req_out.images
del m
gc.collect()
current_omni_platform.empty_cache()
return images, peak
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"})
def test_flux_autoround_w4a16_generates_image():
"""Load the W4A16 quantized FLUX model and verify it produces a valid image."""
images, _ = _generate_image(QUANTIZED_MODEL)
assert len(images) >= 1, "Expected at least one generated image"
img = images[0]
assert img.width == WIDTH, f"Expected width {WIDTH}, got {img.width}"
assert img.height == HEIGHT, f"Expected height {HEIGHT}, got {img.height}"
# Sanity: image should not be blank (all-zero)
import numpy as np
arr = np.array(img)
assert arr.std() > 1.0, "Generated image appears blank (std ≈ 0)"
@pytest.mark.advanced_model
@pytest.mark.diffusion
@hardware_test(res={"cuda": "L4"})
def test_flux_autoround_w4a16_memory_savings():
"""Compare peak GPU memory of quantized vs FP16 baseline.
The W4A16 model should use meaningfully less memory than the
BF16/FP16 baseline since weights are 4-bit instead of 16-bit.
"""
quant_images, quant_peak = _generate_image(QUANTIZED_MODEL)
cleanup_dist_env_and_memory()
_, baseline_peak = _generate_image(BASELINE_MODEL)
print(f"Quantized (W4A16) peak memory: {quant_peak:.0f} MB")
print(f"Baseline (BF16) peak memory: {baseline_peak:.0f} MB")
print(f"Savings: {baseline_peak - quant_peak:.0f} MB")
# W4A16 weights are 4x smaller than BF16/FP16. FLUX.1-dev transformer
# is ~12 GB in BF16, so we expect ~9 GB savings on weights alone.
# Use a conservative threshold to account for activations and overhead.
min_savings_mb = 2000
assert quant_peak + min_savings_mb < baseline_peak, (
f"Quantized model ({quant_peak:.0f} MB) should use at least "
f"{min_savings_mb} MB less than baseline ({baseline_peak:.0f} MB)"
)
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E offline tests for OmniVoice TTS model with text input and audio output.
Uses GPUGenerationWorker for both stages (iterative unmasking + DAC decoder).
"""
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
from pathlib import Path
import numpy as np
import pytest
from tests.utils import hardware_test
MODEL = "k2-fsa/OmniVoice"
def get_stage_config():
return str(
Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml"
)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_omnivoice_text_to_audio() -> None:
"""
Test OmniVoice text-to-audio generation via offline Omni runner.
Deploy Setting: omnivoice.yaml (enforce_eager=true)
Input Modal: text
Output Modal: audio
"""
from vllm_omni.entrypoints.omni import Omni
omni = Omni(
model=MODEL,
stage_configs_path=get_stage_config(),
trust_remote_code=True,
log_stats=True,
)
try:
prompts = {"prompt": "Hello, this is a test for text to audio."}
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
sampling_params_list = [OmniDiffusionSamplingParams()]
outputs = list(omni.generate(prompts, sampling_params_list=sampling_params_list))
assert len(outputs) > 0, "No outputs generated"
# Check final output has audio
final_output = outputs[-1]
ro = final_output.request_output
assert ro is not None, "No request_output"
mm = getattr(ro, "multimodal_output", None)
if not mm and ro.outputs:
mm = getattr(ro.outputs[0], "multimodal_output", None)
assert mm is not None, "No multimodal_output"
assert "audio" in mm, f"No 'audio' key in multimodal_output: {mm.keys()}"
audio = mm["audio"]
if isinstance(audio, np.ndarray):
audio_np = audio
else:
audio_np = audio.cpu().numpy().squeeze()
assert audio_np.size > 0, "Audio output is empty"
rms = np.sqrt(np.mean(audio_np**2))
assert rms > 0.01, f"Audio RMS too low ({rms:.4f}), likely silence"
print(f"Generated audio: {len(audio_np) / 24000:.2f}s, rms={rms:.4f}")
finally:
omni.close()
@@ -120,7 +120,6 @@ def _generate_bagel_image(
torch.cuda.reset_peak_memory_stats()
params_list = omni.default_sampling_params_list
params_list[0].max_tokens = 1 # type: ignore
if len(params_list) > 1:
params_list[1].num_inference_steps = num_inference_steps # type: ignore
params_list[1].extra_args = { # type: ignore
@@ -57,7 +57,7 @@ def get_question(prompt_type="mix"):
return prompts.get(prompt_type, prompts["mix"])
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
@@ -88,7 +88,7 @@ def test_mix_to_audio(omni_runner, omni_runner_handler) -> None:
omni_runner_handler.send_request(request_config)
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325", "xpu": "B60"}, num_cards={"cuda": 4, "rocm": 2, "xpu": 3})
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
@@ -56,7 +56,7 @@ def get_question(prompt_type="video"):
return prompts.get(prompt_type, prompts["video"])
@pytest.mark.core_model
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
@@ -509,7 +509,7 @@ def test_diffusion_batching_async_explicit_batch(model_name: str):
all prompts in a single engine call and returns a single combined result.
The list-prompt path routes through the orchestrator's
``add_batch_request_async`` → ``AsyncOmniDiffusion.generate_batch``
``add_batch_request_async`` → ``AsyncOmni.generate_batch``
and yields ONE ``OmniRequestOutput`` with ALL images combined.
"""
@@ -0,0 +1,124 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E Online tests for CosyVoice3 TTS model with voice cloning.
These tests verify the /v1/audio/speech endpoint works correctly with
the CosyVoice3 model, which requires reference audio for voice cloning.
"""
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
from pathlib import Path
import pytest
from tests.conftest import OmniServerParams
from tests.utils import hardware_test
MODEL = "FunAudioLLM/Fun-CosyVoice3-0.5B-2512"
# Official CosyVoice zero-shot prompt audio and its transcript
REF_AUDIO_URL = "https://raw.githubusercontent.com/FunAudioLLM/CosyVoice/main/asset/zero_shot_prompt.wav"
REF_TEXT = "希望你以后能够做的比我还好呦。"
def get_stage_config(name: str = "cosyvoice3.yaml"):
"""Get the stage config path from vllm_omni model_executor stage_configs."""
return str(Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / name)
def get_prompt(prompt_type="zh"):
prompts = {
"zh": "收到好友从远方寄来的生日礼物,那份意外的惊喜与深深的感动让我热泪盈眶。",
"en": "Hello, this is a voice cloning test with English text.",
}
return prompts.get(prompt_type, prompts["zh"])
tts_server_params = [
pytest.param(
OmniServerParams(
model=MODEL,
stage_config_path=get_stage_config(),
server_args=["--trust-remote-code", "--disable-log-stats"],
),
id="cosyvoice3",
)
]
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_voice_clone_zh_001(omni_server, openai_client) -> None:
"""
Test voice cloning TTS with Chinese text via OpenAI API.
Deploy Setting: default yaml
Input Modal: text + ref_audio + ref_text
Output Modal: audio
Input Setting: stream=False
Datasets: single request
"""
request_config = {
"model": omni_server.model,
"input": get_prompt("zh"),
"stream": False,
"response_format": "wav",
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
openai_client.send_audio_speech_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_voice_clone_zh_002(omni_server, openai_client) -> None:
"""
Test voice cloning TTS with Chinese text via OpenAI API.
Deploy Setting: default yaml
Input Modal: text + ref_audio + ref_text
Output Modal: audio
Input Setting: stream=True
Datasets: single request
"""
request_config = {
"model": omni_server.model,
"input": get_prompt("zh"),
"stream": True,
"response_format": "wav",
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
openai_client.send_audio_speech_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_voice_clone_en_001(omni_server, openai_client) -> None:
"""
Test voice cloning TTS with English text via OpenAI API.
Deploy Setting: default yaml
Input Modal: text + ref_audio + ref_text
Output Modal: audio
Input Setting: stream=False
Datasets: single request
"""
request_config = {
"model": omni_server.model,
"input": get_prompt("en"),
"stream": False,
"response_format": "wav",
"ref_audio": REF_AUDIO_URL,
"ref_text": REF_TEXT,
}
openai_client.send_audio_speech_request(request_config)
+21 -14
View File
@@ -63,20 +63,27 @@ def download_tokenizer():
# CI stage config for H100 / MI325
stage_configs = [get_chunk_config()]
tokenizer_path = download_tokenizer()
os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
# Guard module-level setup so test collection doesn't fail in environments
# where the model cache is read-only or models aren't available.
try:
stage_configs = [get_chunk_config()]
tokenizer_path = download_tokenizer()
os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path
# Create parameter combinations for model and stage config
test_params = [
OmniServerParams(
model=model,
stage_config_path=stage_config,
server_args=["--chat-template", CHAT_TEMPLATE_PATH],
test_params = [
OmniServerParams(
model=model,
stage_config_path=stage_config,
server_args=["--chat-template", CHAT_TEMPLATE_PATH],
)
for model in models
for stage_config in stage_configs
]
except Exception as exc:
pytest.skip(
f"MiMo-Audio online serving tests skipped: module setup failed ({type(exc).__name__}: {exc})",
allow_module_level=True,
)
for model in models
for stage_config in stage_configs
]
def get_prompt(prompt_type="text_only"):
@@ -95,7 +102,7 @@ def get_max_batch_size(size_type="few"):
@pytest.mark.advanced_model
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
"""
@@ -128,7 +135,7 @@ def test_audio_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=2)
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_text_to_text_001(omni_server, openai_client) -> None:
"""
@@ -0,0 +1,84 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
E2E Online tests for OmniVoice TTS model via /v1/audio/speech endpoint.
Tests verify that the OmniVoice model generates valid audio when
accessed through the standard OpenAI-compatible speech API.
"""
import os
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0"
from pathlib import Path
import httpx
import pytest
from tests.conftest import OmniServerParams
from tests.utils import hardware_test
MODEL = "k2-fsa/OmniVoice"
STAGE_CONFIG = str(
Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "omnivoice.yaml"
)
EXTRA_ARGS = [
"--trust-remote-code",
"--disable-log-stats",
]
TEST_PARAMS = [
OmniServerParams(
model=MODEL,
stage_config_path=STAGE_CONFIG,
server_args=EXTRA_ARGS,
)
]
MIN_AUDIO_BYTES = 5000
def make_speech_request(
host: str,
port: int,
text: str,
timeout: float = 180.0,
) -> httpx.Response:
"""Make a request to the /v1/audio/speech endpoint for OmniVoice."""
url = f"http://{host}:{port}/v1/audio/speech"
payload = {"input": text}
with httpx.Client(timeout=timeout) as client:
return client.post(url, json=payload)
def verify_wav_audio(content: bytes) -> bool:
"""Verify that content is valid WAV audio data."""
if len(content) < 44:
return False
return content[:4] == b"RIFF" and content[8:12] == b"WAVE"
@pytest.mark.parametrize("omni_server", TEST_PARAMS, indirect=True)
class TestOmniVoiceTTS:
"""E2E tests for OmniVoice TTS model."""
@pytest.mark.core_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
def test_speech_auto_voice(self, omni_server) -> None:
"""Test auto voice TTS generation (text only, no reference audio)."""
response = make_speech_request(
host=omni_server.host,
port=omni_server.port,
text="Hello, this is a test of the OmniVoice text to speech system.",
)
assert response.status_code == 200, f"Request failed: {response.text}"
assert response.headers.get("content-type") == "audio/wav"
assert verify_wav_audio(response.content), "Response is not valid WAV audio"
assert len(response.content) > MIN_AUDIO_BYTES, (
f"Audio too small ({len(response.content)} bytes), expected > {MIN_AUDIO_BYTES}"
)
@@ -23,10 +23,10 @@ from tests.conftest import (
)
from tests.utils import hardware_test
models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"]
model = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
AUDIO_KEY = ["test"]
IMAGE_KEY = ["square", "quadrate"]
IMAGE_KEY = ["square", "quadrate", "rectangle"]
VIDEO_KEY = ["sphere", "globe", "circle", "round", "ball"]
@@ -49,16 +49,32 @@ def get_chunk_config(default_path):
return path
def get_batch_token_config(default_path):
path = modify_stage_config(
default_path,
updates={
"stage_args": {1: {"engine_args.max_num_batched_tokens": 64}},
},
)
return path
# CI stage config for 2*H100-80G GPUs
default_path = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml")
stage_configs = [default_path, get_chunk_config(default_path)]
if current_omni_platform.is_xpu():
stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")]
default_path = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")
# Create parameter combinations for model and stage config
test_params = [
OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs
pytest.param(OmniServerParams(model=model, stage_config_path=default_path), id="default"),
pytest.param(OmniServerParams(model=model, stage_config_path=get_chunk_config(default_path)), id="async_chunk"),
]
test_token_params = [
pytest.param(
OmniServerParams(model=model, stage_config_path=get_batch_token_config(default_path)), id="batch_token_64"
)
]
@@ -85,6 +101,9 @@ def get_prompt(prompt_type="text_only"):
"text_video": "What is in this video? ",
"text_image": "What is in this image? ",
"text_audio": "What is in this audio? ",
"text_audio_video": "First, what is in this audio? Then, what is in this video? ",
"one_word": "What is the capital of UK? Answer in one word",
"text_chinese": "北京,中国的首都,是一座融合了长城等历史地点与现代建筑的国际化大都市,充满了独特的文化与活力。请重复这句话。",
}
return prompts.get(prompt_type, prompts["text_only"])
@@ -121,7 +140,7 @@ def test_text_to_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_text_to_text_audio_001(omni_server, openai_client) -> None:
"""
Input Modal: text
@@ -288,7 +307,7 @@ def test_video_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
"""
Input Modal: text, audio
@@ -313,7 +332,7 @@ def test_text_audio_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
"""
Input Modal: text, image
@@ -339,7 +358,7 @@ def test_text_image_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
"""
Input Modal: text, video
@@ -367,7 +386,7 @@ def test_text_video_to_text_audio_001(omni_server, openai_client) -> None:
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
@pytest.mark.parametrize("omni_server", test_params + test_token_params, indirect=True)
def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
"""
Input Modal: text, audio, image, video
@@ -393,3 +412,213 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None:
"key_words": {"audio": AUDIO_KEY, "image": IMAGE_KEY, "video": VIDEO_KEY},
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_audio_in_video_001(omni_server, openai_client) -> None:
"""
Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
Output Modal: text, audio
Input Setting: stream=False
Datasets: single request
"""
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300, embed_audio=True)['base64']}"
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
video_data_url=video_data_url,
content_text=get_prompt("text_audio_video"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": False,
"use_audio_in_video": True,
"key_words": {"video": VIDEO_KEY, "audio": AUDIO_KEY + ["beep", "electronic"]},
}
openai_client.send_omni_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_audio_in_video_002(omni_server, openai_client) -> None:
"""
Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video).
Output Modal: text, audio
Input Setting: stream=True
Datasets: few requests
"""
video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300, embed_audio=True)['base64']}"
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
video_data_url=video_data_url,
content_text=get_prompt("text_audio_video"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"use_audio_in_video": True,
"key_words": {"video": VIDEO_KEY},
}
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_one_word_prompt_001(omni_server, openai_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("one_word"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"key_words": {"text": ["london"]},
}
# Retry only when assert_omni_response fails on text/audio cosine similarity (see tests/conftest.py).
_similarity_assert_msg = "The audio content is not same as the text"
_max_retries = 3
for attempt in range(_max_retries):
try:
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
break
except AssertionError as e:
if _similarity_assert_msg not in str(e) or attempt == _max_retries - 1:
raise
print(f"Similarity assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_speaker_001(omni_server, openai_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("text"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"speaker": "Chelsie",
"key_words": {"text": ["beijing"]},
}
openai_client.send_omni_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_speaker_002(omni_server, openai_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("text"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"speaker": "Ethan",
"key_words": {"text": ["beijing"]},
}
# Retry only when assert_omni_response fails on preset voice gender (see tests/conftest.py).
_gender_assert_substr = "estimated gender"
_max_retries = 3
for attempt in range(_max_retries):
try:
openai_client.send_omni_request(request_config, request_num=get_max_batch_size())
break
except AssertionError as e:
if _gender_assert_substr not in str(e) or attempt == _max_retries - 1:
raise
print(f"Gender assertion failed, retrying {attempt + 2}/{_max_retries}: {e!r}")
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_speaker_003(omni_server, openai_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("text"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"speaker": "CHELSIE",
"key_words": {"text": ["beijing"]},
}
openai_client.send_omni_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
def test_language_001(omni_server, openai_client) -> None:
"""
Input Modal: text only (one-word answer constraint).
Output Modal: text, audio (default ``modalities``); ``key_words`` only assert on text.
Input Setting: stream=True
Datasets: single request
"""
messages = dummy_messages_from_mix_data(
system_prompt=get_system_prompt(),
content_text=get_prompt("text_chinese"),
)
request_config = {
"model": omni_server.model,
"messages": messages,
"stream": True,
"key_words": {"text": ["北京"]},
}
openai_client.send_omni_request(request_config)
@@ -120,6 +120,31 @@ def test_voice_002(omni_server, openai_client) -> None:
openai_client.send_audio_speech_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
def test_voice_003(omni_server, openai_client) -> None:
"""
Test text input processing and audio output via OpenAI API.
Deploy Setting: default yaml
Input Modal: text
Output Modal: audio
Input Setting: stream=False, language=chinese
Datasets: few requests
"""
request_config = {
"model": omni_server.model,
"input": get_prompt(),
"stream": False,
"response_format": "wav",
"task_type": "CustomVoice",
"voice": "SERENA",
}
openai_client.send_audio_speech_request(request_config)
@pytest.mark.advanced_model
@pytest.mark.omni
@hardware_test(res={"cuda": "L4"}, num_cards=1)
@@ -98,6 +98,18 @@ def _get_diffusion_feature_cases(model: str):
id="parallel_004",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_005",
marks=PARALLEL_FEATURE_MARKS,
),
]
@@ -28,6 +28,11 @@ PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
def _get_diffusion_feature_cases(model: str):
return [
pytest.param(
OmniServerParams(model=model, server_args=["--step-execution"]),
id="step_execution",
marks=SINGLE_CARD_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(model=model, server_args=["--cache-backend", "tea_cache"]),
id="cache_tea_cache",
@@ -102,6 +107,18 @@ def _get_diffusion_feature_cases(model: str):
id="vae_patch_parallel_2",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=model,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_hsdp",
marks=PARALLEL_FEATURE_MARKS,
),
]
@@ -7,7 +7,7 @@ and are supported by Qwen-Image-Layered model.
Kept cases (maximum feature coverage, no cpu-offload):
sp_001 : cache_dit + Ulysses-SP 2 (2×H100)
cfg_parallel_002 : cache_dit + CFG-Parallel 2 (2×H100)
layers_guard_001 : layerwise offload + layers=2 (1×H100, issue #1969 guard)
layers_guard_001 : layerwise offload + layers=3 (1×H100, issue #1969 guard)
Total distinct features covered: cache_dit, Ulysses-SP, CFG-Parallel, layerwise-offload.
"""
@@ -62,6 +62,18 @@ FEATURE_CASES = [
id="cfg_parallel_001",
marks=PARALLEL_FEATURE_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_hsdp",
marks=PARALLEL_FEATURE_MARKS,
),
]
@@ -102,8 +114,8 @@ LAYERS_GUARD_CASES = [
model=MODEL,
server_args=["--enable-layerwise-offload"],
),
2,
id="layers_guard_001_layers2",
3,
id="layers_guard_001_layers3",
marks=SINGLE_CARD_FEATURE_MARKS,
),
]
@@ -0,0 +1,161 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""
Comprehensive e2e tests of diffusion features for Wan2.1-VACE in online serving mode.
Wan2.1-VACE supports: Cache-DiT, Ulysses-SP, Ring, CFG-Parallel, TP,
VAE-Patch-Parallel, HSDP. TeaCache is NOT supported for this model, so
Cache-DiT is used in place of TeaCache for single-card and CFG tests.
Uses the 1.3B variant for faster CI testing.
Coverage:
Single GPU:
- Cache-DiT + layerwise CPU offload
Two GPUs:
- Cache-DiT + Ulysses-SP = 2
- Cache-DiT + Ring = 2
- Cache-DiT + CFG-Parallel = 2
- Cache-DiT + TP = 2 + VAE-Patch-Parallel = 2
- Cache-DiT + HSDP = 2 + VAE-Patch-Parallel = 2
"""
import pytest
from tests.conftest import (
OmniServer,
OmniServerParams,
OpenAIClientHandler,
)
from tests.utils import hardware_marks
MODEL = "Wan-AI/Wan2.1-VACE-1.3B-diffusers"
PROMPT = "A cat walking slowly across a sunlit garden path"
SINGLE_CARD_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"})
PARALLEL_FEATURE_MARKS = hardware_marks(res={"cuda": "H100"}, num_cards=2)
def _get_vace_feature_cases():
return [
# Single GPU: Cache-DiT + layerwise CPU offload
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--cache-backend",
"cache_dit",
"--enable-layerwise-offload",
"--vae-use-tiling",
],
),
id="single_card_001",
marks=SINGLE_CARD_FEATURE_MARKS,
),
# 2 GPUs: Cache-DiT + Ulysses-SP = 2
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--cache-backend",
"cache_dit",
"--ulysses-degree",
"2",
"--vae-use-tiling",
],
),
id="parallel_001",
marks=PARALLEL_FEATURE_MARKS,
),
# 2 GPUs: Cache-DiT + Ring = 2
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--cache-backend",
"cache_dit",
"--ring",
"2",
"--vae-use-tiling",
],
),
id="parallel_002",
marks=PARALLEL_FEATURE_MARKS,
),
# 2 GPUs: Cache-DiT + CFG-Parallel = 2
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--cache-backend",
"cache_dit",
"--cfg-parallel-size",
"2",
"--vae-use-tiling",
],
),
id="parallel_003",
marks=PARALLEL_FEATURE_MARKS,
),
# 2 GPUs: Cache-DiT + TP = 2 + VAE-Patch-Parallel = 2
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--cache-backend",
"cache_dit",
"--tensor-parallel-size",
"2",
"--vae-patch-parallel-size",
"2",
"--vae-use-tiling",
],
),
id="parallel_004",
marks=PARALLEL_FEATURE_MARKS,
),
# 2 GPUs: Cache-DiT + HSDP = 2 + VAE-Patch-Parallel = 2
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--cache-backend",
"cache_dit",
"--hsdp-shard-size",
"2",
"--vae-patch-parallel-size",
"2",
"--vae-use-tiling",
],
),
id="parallel_005",
marks=PARALLEL_FEATURE_MARKS,
),
]
@pytest.mark.advanced_model
@pytest.mark.diffusion
@pytest.mark.parametrize(
"omni_server",
_get_vace_feature_cases(),
indirect=True,
)
def test_wan_2_1_vace(omni_server: OmniServer, openai_client: OpenAIClientHandler):
"""Test VACE T2V generation with all supported diffusion acceleration features."""
openai_client.send_video_diffusion_request(
{
"model": MODEL,
"form_data": {
"prompt": PROMPT,
"height": 480,
"width": 320,
"num_frames": 5,
"fps": 8,
"num_inference_steps": 2,
"guidance_scale": 5.0,
"seed": 42,
},
}
)
@@ -60,6 +60,18 @@ def _get_diffusion_feature_cases():
id="parallel_teacache_fp8_ulysses2_ring2",
marks=FOUR_CARD_MARKS,
),
pytest.param(
OmniServerParams(
model=MODEL,
server_args=[
"--use-hsdp",
"--hsdp-shard-size",
"2",
],
),
id="parallel_hsdp",
marks=[*FOUR_CARD_MARKS, pytest.mark.skip(reason="issue #2435")],
),
]
+2 -2
View File
@@ -33,7 +33,7 @@ stage_args:
temperature: 0.4
top_p: 0.9
top_k: 1
max_tokens: 100
max_tokens: 150
seed: 42
ignore_eos: False
detokenize: True
@@ -74,7 +74,7 @@ stage_args:
devices: "1"
engine_args:
model_stage: code2wav
max_num_seqs: 1
max_num_seqs: 5
model_arch: Qwen3OmniMoeForConditionalGeneration
worker_type: generation
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
@@ -60,6 +60,7 @@ async def generate(
@pytest.mark.core_model
@pytest.mark.omni
@pytest.mark.real_hf_config
@hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1)
@pytest.mark.asyncio
async def test_abort():
@@ -61,3 +61,31 @@ def test_build_add_request_message_preserves_additional_information():
assert request.additional_information.entries["text"].list_data == ["hello world"]
assert request.additional_information.entries["speaker"].list_data == ["vivian"]
output_processor.add_request.assert_called_once()
def test_build_add_request_message_with_resumable_streaming():
engine = object.__new__(AsyncOmniEngine)
params = SamplingParams(max_tokens=8)
engine.default_sampling_params_list = [params]
engine.stage_metadata = [{"stage_type": "llm"}]
engine.supported_tasks = ("generate",)
input_processor = Mock()
input_processor.process_inputs.return_value = _make_engine_core_request()
engine.input_processor = input_processor
output_processor = Mock()
engine.output_processors = [output_processor]
msg = engine._build_add_request_message(
request_id="req-stream",
prompt={"prompt_token_ids": [1, 2, 3]},
sampling_params_list=[params],
final_stage_id=0,
resumable=True,
message_type="streaming_update",
)
assert msg["type"] == "streaming_update"
input_processor.process_inputs.assert_called_once()
assert input_processor.process_inputs.call_args.kwargs["resumable"] is True
@@ -113,6 +113,7 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch):
executor_class=object,
engine_manager=object(),
coordinator=object(),
proc=None,
addresses=types.SimpleNamespace(
inputs=["inproc://input"],
outputs=["inproc://output"],
@@ -106,7 +106,7 @@ def test_encode_image_base64():
class MockGenerationResult:
"""Mock result object from AsyncOmniDiffusion.generate()"""
"""Mock result object from AsyncOmni.generate()"""
def __init__(self, images):
self.images = images
Loaded 100 of 238 files, more files were not shown because too many files have changed in this diff. Show more