Compare commits
51
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
191b9a8dbf | ||
|
|
2804a85e58 | ||
|
|
4c031580cf | ||
|
+11 |
f50c5a413f | ||
|
|
0e83ebe1d4 | ||
|
|
10db95f9a9 | ||
|
|
515d15ef87 | ||
|
|
cd71567b06 | ||
|
|
6dc61c9a20 | ||
|
|
563f73b78a | ||
|
|
7fb86d51cd | ||
|
|
fa275fd1be | ||
|
|
6be5d05a7c | ||
|
|
6211413677 | ||
|
|
728cf6d023 | ||
|
|
50bb47a629 | ||
|
|
ca02351a1e | ||
|
|
458f402323 | ||
|
|
e2892ef6c1 | ||
|
|
c1d2dcc3ac | ||
|
|
900f6aa837 | ||
|
|
d3daafbe41 | ||
|
|
ebc9a8d875 | ||
|
|
9c2a576301 | ||
|
|
9595be59ae | ||
|
|
bbae904f1d | ||
|
|
70a62651b9 | ||
|
|
4e4bbc42a6 | ||
|
|
6ef0e907af | ||
|
|
3def008b32 | ||
|
|
bf5bd0a4c0 | ||
|
|
3fd4a4dc27 | ||
|
|
d40840b114 | ||
|
|
08cb436d4e | ||
|
|
c3376a466b | ||
|
|
080593c364 | ||
|
|
183775ec36 | ||
|
|
7b965a7b6c | ||
|
|
7274e15840 | ||
|
|
dd0b6fd4e8 | ||
|
|
369f301a1b | ||
|
|
f8d0bf5389 | ||
|
|
642f16949d | ||
|
|
a5bf33f25a | ||
|
|
7cb7b3e8ad | ||
|
|
de1ac5ee88 | ||
|
|
e68ca4aa5c | ||
|
|
0cbd1cbeb8 | ||
|
|
c1a978ae8d | ||
|
|
1734c1ec23 | ||
|
|
12a8f57934 |
No files matched your search
@@ -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"
|
||||
@@ -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
@@ -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
|
||||
Executable
+193
@@ -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}")
|
||||
Executable
+33
@@ -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
@@ -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
@@ -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
@@ -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
|
||||
@@ -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
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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(
|
||||
|
||||
@@ -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
@@ -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
@@ -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,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.
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
@@ -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 |
@@ -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
|
||||
|
||||
@@ -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 ⬜
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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,
|
||||
|
||||
@@ -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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
|
||||
@@ -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"},
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -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
@@ -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)
|
||||
@@ -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.
|
||||
|
||||
@@ -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"
|
||||
@@ -0,0 +1,4 @@
|
||||
-r common.txt
|
||||
# MUSA platform dependencies
|
||||
torchada>=0.1.46
|
||||
onnxruntime>=1.23.2
|
||||
@@ -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
@@ -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:
|
||||
|
||||
Whitespace-only changes.
@@ -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):
|
||||
|
||||
@@ -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)
|
||||
@@ -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.
|
||||
|
||||
@@ -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)
|
||||
@@ -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")],
|
||||
),
|
||||
]
|
||||
|
||||
|
||||
|
||||
@@ -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
Reference in new issue
Block a user