Compare commits
58
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
2b3a1c8200 | ||
|
|
08b2f694f3 | ||
|
|
170f98c847 | ||
|
|
3805852a71 | ||
|
|
d75a6433b4 | ||
|
|
6506ed3d4d | ||
|
|
d4f2a49a68 | ||
|
|
9bc8705c98 | ||
|
|
1733831f78 | ||
|
|
20b7a67b0d | ||
|
|
9577ccba50 | ||
|
|
54a9d241c0 | ||
|
|
6de0ca36bb | ||
|
|
293396737d | ||
|
|
aa59dc50c7 | ||
|
|
f44f3ef9f7 | ||
|
|
8f869c8ea5 | ||
|
|
1be1a71156 | ||
|
|
3fe3aa359b | ||
|
|
541f06c60d | ||
|
|
a3834e32cf | ||
|
|
152343d5fb | ||
|
|
03cc888485 | ||
|
|
e27fd46ab9 | ||
|
|
eb964dcec5 | ||
|
|
5414f78e5f | ||
|
|
f3d40b5402 | ||
|
|
3529b8280d | ||
|
|
e26f5cb918 | ||
|
|
0d3c259a76 | ||
|
|
74b5ae5ebe | ||
|
|
1fe267d6c9 | ||
|
|
14879d6619 | ||
|
|
18e083e512 | ||
|
|
a1a53e5f94 | ||
|
|
1b2ab23586 | ||
|
|
57ea047acf | ||
|
|
0342827dd7 | ||
|
|
362b3ab013 | ||
|
|
c2b4564a67 | ||
|
|
aa92a3a35d | ||
|
|
2fa4c7a4a9 | ||
|
|
39ae6e12f8 | ||
|
|
2d9e5e0515 | ||
|
|
d50bf1aa42 | ||
|
|
f1bf544208 | ||
|
|
57c8f3abfc | ||
|
|
6ae355b9a9 | ||
|
|
73df8326d7 | ||
|
|
95f633f830 | ||
|
|
e1b0b83238 | ||
|
|
57227dc7f3 | ||
|
|
ca75623ff3 | ||
|
|
3753b21c3d | ||
|
|
e0fde88a5f | ||
|
|
bf03413c36 | ||
|
|
f174ee54b0 | ||
|
|
6ddaf6188d |
No files matched your search
@@ -1,14 +1,14 @@
|
||||
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
|
||||
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
|
||||
# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# exception in upload_pipeline.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
|
||||
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline.py (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
|
||||
steps:
|
||||
- label: ":github: Resolve skip-ci & upload pipeline"
|
||||
key: upload-ci-pipeline
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
|
||||
- "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
key: upload-ready-pipeline
|
||||
if: __UPLOAD_READY_IF__
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -41,7 +41,7 @@ steps:
|
||||
key: upload-merge-pipeline
|
||||
if: __UPLOAD_MERGE_IF__
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -23,7 +23,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -0,0 +1,308 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
|
||||
|
||||
Bootstrap mode (pipeline.yml with __IMAGE_BUILD_IF__ placeholders):
|
||||
- Detect docs-only PR/main changes and substitute skip-ci ``if`` expressions.
|
||||
|
||||
Test pipeline mode (e.g. test-merge.yml):
|
||||
- Drop steps or groups whose ``source_file_dependencies`` do not match changed files.
|
||||
- ``source_file_dependencies`` may be set on a leaf step or on a group step.
|
||||
- Always strip ``source_file_dependencies`` from uploaded YAML (Buildkite does not
|
||||
recognize this uploader-only key).
|
||||
|
||||
Usage:
|
||||
python3 upload_pipeline.py [--upload] <pipeline.yml>
|
||||
|
||||
# Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
|
||||
python3 upload_pipeline.py --upload .buildkite/pipeline.yml
|
||||
|
||||
# Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
|
||||
python3 upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import yaml
|
||||
except ImportError:
|
||||
subprocess.check_call(
|
||||
[sys.executable, "-m", "pip", "install", "-q", "pyyaml"],
|
||||
)
|
||||
import yaml
|
||||
|
||||
LOG = "upload_pipeline"
|
||||
ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
DOC_SEP = "\n---\n"
|
||||
BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
|
||||
|
||||
|
||||
def _log(message: str) -> None:
|
||||
print(f"{LOG}: {message}", file=sys.stderr)
|
||||
|
||||
|
||||
def _git(*args: str) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
["git", *args],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
|
||||
def resolve_changed_files() -> list[str] | None:
|
||||
"""Return changed file paths, or None when diff cannot be resolved."""
|
||||
is_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false" and os.environ.get(
|
||||
"BUILDKITE_PULL_REQUEST", ""
|
||||
)
|
||||
commit = os.environ.get("BUILDKITE_COMMIT", "")
|
||||
|
||||
if is_pr:
|
||||
base_branch = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
|
||||
base_ref = f"origin/{base_branch}"
|
||||
if _git("rev-parse", "--verify", base_ref).returncode != 0:
|
||||
_log(f"origin/{base_branch} not found locally; trying fetch")
|
||||
_git("fetch", "--depth=200", "origin", base_branch)
|
||||
if _git("rev-parse", "--verify", base_ref).returncode != 0:
|
||||
if _git("rev-parse", "--verify", base_branch).returncode == 0:
|
||||
base_ref = base_branch
|
||||
else:
|
||||
_log(f"cannot resolve PR base {base_branch}; using safe defaults")
|
||||
return None
|
||||
diff_range = f"{base_ref}...{commit}"
|
||||
elif os.environ.get("BUILDKITE_BRANCH", "") == "main":
|
||||
if _git("rev-parse", "--verify", f"{commit}^").returncode != 0:
|
||||
_log("main commit has no parent; using safe defaults")
|
||||
return None
|
||||
diff_range = f"{commit}^..{commit}"
|
||||
else:
|
||||
_log("not PR/main build; using safe defaults")
|
||||
return None
|
||||
|
||||
result = _git("diff", "--name-only", diff_range)
|
||||
if result.returncode != 0:
|
||||
_log(f"git diff failed ({diff_range}); using safe defaults")
|
||||
return None
|
||||
|
||||
files = [line for line in result.stdout.splitlines() if line.strip()]
|
||||
_log(f"{len(files)} changed file(s)")
|
||||
return files
|
||||
|
||||
|
||||
def is_docs_only_change(changed_files: list[str]) -> bool:
|
||||
has_any = False
|
||||
for file_path in changed_files:
|
||||
if not file_path:
|
||||
continue
|
||||
has_any = True
|
||||
if file_path.startswith("docs/"):
|
||||
continue
|
||||
if file_path.endswith(".md"):
|
||||
continue
|
||||
if file_path == "mkdocs.yml":
|
||||
continue
|
||||
return False
|
||||
return has_any
|
||||
|
||||
|
||||
def resolve_skip_ci(changed_files: list[str] | None) -> bool:
|
||||
if changed_files is None:
|
||||
_log("skip-ci=0 (could not resolve changed files)")
|
||||
return False
|
||||
if is_docs_only_change(changed_files):
|
||||
_log("docs-only change detected; skip-ci=1")
|
||||
return True
|
||||
_log("non-doc changes detected; skip-ci=0")
|
||||
return False
|
||||
|
||||
|
||||
def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
|
||||
if DOC_SEP in text:
|
||||
_, continuation = text.split(DOC_SEP, 1)
|
||||
else:
|
||||
continuation = text
|
||||
|
||||
nightly_only = (
|
||||
'(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
|
||||
)
|
||||
if skip_ci:
|
||||
image_if = f"'{nightly_only}'"
|
||||
ready_if = "'false'"
|
||||
merge_if = "'false'"
|
||||
else:
|
||||
image_if = "'true'"
|
||||
ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
|
||||
merge_if = (
|
||||
'\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
|
||||
'&& build.env("WEEKLY") != "1") || '
|
||||
'(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
|
||||
)
|
||||
|
||||
return (
|
||||
continuation.replace("__IMAGE_BUILD_IF__", image_if)
|
||||
.replace("__UPLOAD_READY_IF__", ready_if)
|
||||
.replace("__UPLOAD_MERGE_IF__", merge_if)
|
||||
)
|
||||
|
||||
|
||||
def _matches_dependencies(changed_files: list[str], prefixes: list[str]) -> bool:
|
||||
for path in changed_files:
|
||||
for prefix in prefixes:
|
||||
normalized = prefix.rstrip("/")
|
||||
if path == normalized or path.startswith(f"{normalized}/"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _strip_source_file_dependencies(step: dict[str, Any]) -> dict[str, Any]:
|
||||
return {key: value for key, value in step.items() if key != "source_file_dependencies"}
|
||||
|
||||
|
||||
def _step_label(step: dict[str, Any]) -> str:
|
||||
return str(step.get("group") or step.get("label") or "<step>")
|
||||
|
||||
|
||||
def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
|
||||
filtered: list[Any] = []
|
||||
for step in steps:
|
||||
if not isinstance(step, dict):
|
||||
filtered.append(step)
|
||||
continue
|
||||
|
||||
deps = step.get("source_file_dependencies")
|
||||
if deps is not None and not isinstance(deps, list):
|
||||
raise ValueError(
|
||||
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
|
||||
)
|
||||
if deps is not None and not _matches_dependencies(changed_files, deps):
|
||||
_log(
|
||||
f"skip {_step_label(step)!r} (no changes under {deps})",
|
||||
)
|
||||
continue
|
||||
|
||||
nested = step.get("steps")
|
||||
if nested is not None:
|
||||
kept_nested = _filter_steps(nested, changed_files)
|
||||
if not kept_nested:
|
||||
_log(f"omit empty group {_step_label(step)!r}")
|
||||
continue
|
||||
new_step = _strip_source_file_dependencies(step)
|
||||
new_step["steps"] = kept_nested
|
||||
filtered.append(new_step)
|
||||
continue
|
||||
|
||||
if deps is not None:
|
||||
filtered.append(_strip_source_file_dependencies(step))
|
||||
else:
|
||||
filtered.append(step)
|
||||
|
||||
return filtered
|
||||
|
||||
|
||||
def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
|
||||
"""Remove uploader-only keys while keeping all steps (no diff filtering)."""
|
||||
stripped: list[Any] = []
|
||||
for step in steps:
|
||||
if not isinstance(step, dict):
|
||||
stripped.append(step)
|
||||
continue
|
||||
|
||||
deps = step.get("source_file_dependencies")
|
||||
if deps is not None and not isinstance(deps, list):
|
||||
raise ValueError(
|
||||
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
|
||||
)
|
||||
|
||||
nested = step.get("steps")
|
||||
new_step = _strip_source_file_dependencies(step)
|
||||
if nested is not None:
|
||||
new_step["steps"] = _strip_uploader_metadata_from_steps(nested)
|
||||
stripped.append(new_step)
|
||||
|
||||
return stripped
|
||||
|
||||
|
||||
def render_test_pipeline(
|
||||
doc: dict[str, Any],
|
||||
changed_files: list[str] | None,
|
||||
) -> dict[str, Any]:
|
||||
steps = doc.get("steps")
|
||||
if not isinstance(steps, list):
|
||||
return doc
|
||||
if changed_files is not None:
|
||||
steps = _filter_steps(steps, changed_files)
|
||||
else:
|
||||
steps = _strip_uploader_metadata_from_steps(steps)
|
||||
return {**doc, "steps": steps}
|
||||
|
||||
|
||||
def resolve_pipeline_path(arg: str) -> Path:
|
||||
path = Path(arg)
|
||||
if path.is_absolute():
|
||||
return path
|
||||
return ROOT / path
|
||||
|
||||
|
||||
def render_pipeline(path: Path) -> str:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
changed_files = resolve_changed_files()
|
||||
|
||||
if BOOTSTRAP_MARKER in text:
|
||||
rendered = render_bootstrap_pipeline(text, skip_ci=resolve_skip_ci(changed_files))
|
||||
return rendered
|
||||
|
||||
doc = yaml.safe_load(text)
|
||||
if not isinstance(doc, dict):
|
||||
raise ValueError(f"invalid pipeline YAML: {path}")
|
||||
|
||||
doc = render_test_pipeline(doc, changed_files)
|
||||
|
||||
return yaml.safe_dump(doc, sort_keys=False)
|
||||
|
||||
|
||||
def upload_to_buildkite(content: str) -> None:
|
||||
subprocess.run(
|
||||
["buildkite-agent", "pipeline", "upload"],
|
||||
input=content,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"pipeline",
|
||||
nargs="?",
|
||||
default=".buildkite/pipeline.yml",
|
||||
help="Pipeline YAML path (default: .buildkite/pipeline.yml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload",
|
||||
action="store_true",
|
||||
help="Pipe rendered YAML to buildkite-agent pipeline upload",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
path = resolve_pipeline_path(args.pipeline)
|
||||
if not path.is_file():
|
||||
_log(f"missing pipeline file: {path}")
|
||||
return 1
|
||||
|
||||
rendered = render_pipeline(path)
|
||||
if args.upload:
|
||||
upload_to_buildkite(rendered)
|
||||
else:
|
||||
sys.stdout.write(rendered)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,146 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
|
||||
# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
|
||||
#
|
||||
# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml]
|
||||
# pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml)
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||||
PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
|
||||
if [[ "${PIPELINE_ARG}" = /* ]]; then
|
||||
PIPELINE_YML="${PIPELINE_ARG}"
|
||||
else
|
||||
PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
|
||||
fi
|
||||
|
||||
# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
|
||||
is_docs_only_change() {
|
||||
local file_path
|
||||
local has_any=0
|
||||
|
||||
while IFS= read -r file_path; do
|
||||
[[ -z "${file_path}" ]] && continue
|
||||
has_any=1
|
||||
|
||||
if [[ "${file_path}" == docs/* ]]; then
|
||||
continue
|
||||
fi
|
||||
if [[ "${file_path}" == *.md ]]; then
|
||||
continue
|
||||
fi
|
||||
if [[ "${file_path}" == "mkdocs.yml" ]]; then
|
||||
continue
|
||||
fi
|
||||
return 1
|
||||
done
|
||||
|
||||
[[ "${has_any}" -eq 1 ]]
|
||||
}
|
||||
|
||||
resolve_skip_ci() {
|
||||
local is_pr_build=0
|
||||
local files
|
||||
local base_branch base_ref
|
||||
|
||||
if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
|
||||
is_pr_build=1
|
||||
fi
|
||||
|
||||
if [[ "${is_pr_build}" -eq 1 ]]; then
|
||||
base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
|
||||
if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
|
||||
echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
|
||||
git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
|
||||
fi
|
||||
|
||||
base_ref=""
|
||||
if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
|
||||
base_ref="origin/${base_branch}"
|
||||
elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
|
||||
base_ref="${base_branch}"
|
||||
else
|
||||
echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
|
||||
if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
|
||||
echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
|
||||
if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
|
||||
echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
|
||||
echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
else
|
||||
echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
|
||||
if is_docs_only_change <<< "${files}"; then
|
||||
echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
|
||||
echo -n 1
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
}
|
||||
|
||||
SKIP_CI="$(resolve_skip_ci)"
|
||||
|
||||
if [[ ! -f "${PIPELINE_YML}" ]]; then
|
||||
echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
export ROOT SKIP_CI PIPELINE_YML
|
||||
python3 <<'PY' | buildkite-agent pipeline upload
|
||||
import os
|
||||
import pathlib
|
||||
|
||||
path = pathlib.Path(os.environ["PIPELINE_YML"])
|
||||
text = path.read_text(encoding="utf-8")
|
||||
sep = "\n---\n"
|
||||
# Two supported layouts:
|
||||
# - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps.
|
||||
# - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap.
|
||||
if sep in text:
|
||||
_, continuation = text.split(sep, 1)
|
||||
else:
|
||||
continuation = text
|
||||
|
||||
skip = os.environ.get("SKIP_CI") == "1"
|
||||
# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
|
||||
# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
|
||||
nightly_only = (
|
||||
'(build.pull_request.labels includes "nightly-test") '
|
||||
'|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
|
||||
)
|
||||
# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
|
||||
if skip:
|
||||
rep = f"'{nightly_only}'"
|
||||
ready_rep = "'false'"
|
||||
merge_rep = "'false'"
|
||||
else:
|
||||
rep = "'true'"
|
||||
ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
|
||||
merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
|
||||
rendered = (
|
||||
continuation
|
||||
.replace("__IMAGE_BUILD_IF__", rep)
|
||||
.replace("__UPLOAD_READY_IF__", ready_rep)
|
||||
.replace("__UPLOAD_MERGE_IF__", merge_rep)
|
||||
)
|
||||
print(rendered, end="")
|
||||
PY
|
||||
+159
-148
@@ -1,54 +1,111 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
steps:
|
||||
- label: "Diffusion · LoRA Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
"
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
@@ -62,42 +119,6 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Engine Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
@@ -111,50 +132,65 @@ steps:
|
||||
# This model seems to be important, so I am trying to run it here.
|
||||
# Since this is from test-nightly.yml,
|
||||
# I keep this non-blocking.
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
steps:
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "Qwen3-TTS Base E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
- label: "Qwen3-TTS Base E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
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"
|
||||
'
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
@@ -195,28 +231,3 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
"
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
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"
|
||||
'
|
||||
+156
-144
@@ -1,42 +1,33 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
||||
|
||||
- label: "Diffusion · Batch Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -46,23 +37,92 @@ steps:
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
||||
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
steps:
|
||||
- label: "Diffusion · Batch Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Offloader Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Diffusion · Cache Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
@@ -76,56 +136,6 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
|
||||
- label: "Engine Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
@@ -149,43 +159,56 @@ steps:
|
||||
# - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
# - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "MiMo-Audio E2E Test with H100"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
steps:
|
||||
- label: "MiMo-Audio E2E Test with H100"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
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"
|
||||
'
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
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"
|
||||
'
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
|
||||
'
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
@@ -225,14 +248,3 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
|
||||
'
|
||||
+140
-28
@@ -149,7 +149,7 @@ steps:
|
||||
commands:
|
||||
- pytest -s -v tests/entrypoints/ -m "advanced_model and cuda and L4" --run-level "advanced_model"
|
||||
agents:
|
||||
queue: "gpu_4_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -164,7 +164,35 @@ steps:
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
depends_on: upload-merge-pipeline
|
||||
steps:
|
||||
- label: "TTS · Step-Audio2 E2E Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_step_audio2.py -m 'advanced_model' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
|
||||
- tests/e2e/offline_inference/test_qwen3_tts_customvoice.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -173,7 +201,7 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -186,6 +214,12 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Qwen3-TTS Base Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_base.py
|
||||
- tests/e2e/offline_inference/test_qwen3_tts_base.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -194,7 +228,34 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Ming-Omni-TTS Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_ming_tts.py
|
||||
- tests/e2e/offline_inference/test_ming_tts.py
|
||||
- vllm_omni/model_executor/models/ming_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/ming_tts.py
|
||||
- vllm_omni/deploy/ming_tts.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
pytest -s -v tests/e2e/online_serving/test_ming_tts.py tests/e2e/offline_inference/test_ming_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -207,6 +268,12 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Voxtral-TTS Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_voxtral_tts.py
|
||||
- tests/e2e/online_serving/test_voxtral_tts.py
|
||||
- vllm_omni/model_executor/models/voxtral_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/voxtral_tts.py
|
||||
- vllm_omni/deploy/voxtral_tts.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
@@ -249,6 +316,12 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · GLM-TTS Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_glm_tts.py
|
||||
- tests/e2e/offline_inference/test_glm_tts.py
|
||||
- vllm_omni/model_executor/models/glm_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/glm_tts.py
|
||||
- vllm_omni/deploy/glm_tts.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
@@ -291,6 +364,12 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- vllm_omni/model_executor/models/qwen3_omni/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -329,6 +408,11 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Mimo Audio Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_mimo_audio.py
|
||||
- vllm_omni/model_executor/models/mimo_audio/
|
||||
- vllm_omni/model_executor/stage_input_processors/mimo_audio.py
|
||||
- vllm_omni/deploy/mimo_audio.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -366,37 +450,19 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · VoxCPM2 Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · VoxCPM2 Online Test"
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_voxcpm2_tts.py
|
||||
- tests/e2e/online_serving/test_voxcpm2_tts.py
|
||||
- vllm_omni/model_executor/models/voxcpm2/
|
||||
- vllm_omni/deploy/voxcpm2.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
@@ -413,6 +479,12 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Higgs-Audio V2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_higgs_audio_v2.py
|
||||
- tests/e2e/offline_inference/test_higgs_audio_v2.py
|
||||
- vllm_omni/model_executor/models/higgs_audio_v2/
|
||||
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v2.py
|
||||
- vllm_omni/deploy/higgs_audio_v2.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
@@ -438,6 +510,9 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -476,6 +551,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Longcat Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_longcat_image_edit.py
|
||||
- vllm_omni/diffusion/models/longcat_image/
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_longcat_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -514,6 +592,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Longcat Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_longcat_image.py
|
||||
- vllm_omni/diffusion/models/longcat_image/
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_longcat_image.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -552,6 +633,13 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Bagel Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_bagel.py
|
||||
- tests/e2e/online_serving/test_bagel.py
|
||||
- vllm_omni/model_executor/models/bagel/
|
||||
- vllm_omni/diffusion/models/bagel/
|
||||
- vllm_omni/model_executor/stage_input_processors/bagel.py
|
||||
- vllm_omni/deploy/bagel.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -591,6 +679,9 @@ steps:
|
||||
|
||||
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c '
|
||||
@@ -632,6 +723,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Qwen Image Layered Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_layered.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -673,6 +767,10 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_wan22_t2v.py
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
- vllm_omni/diffusion/models/wan2_2/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -714,6 +812,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · HunyuanVideo 1.5 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_hunyuan_video_15.py
|
||||
- vllm_omni/diffusion/models/hunyuan_video/
|
||||
timeout_in_minutes: 40
|
||||
commands:
|
||||
- |
|
||||
@@ -755,6 +856,10 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Z Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_z_image.py
|
||||
- tests/e2e/online_serving/test_z_image.py
|
||||
- vllm_omni/diffusion/models/z_image/
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_z_image.py tests/e2e/online_serving/test_z_image.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
@@ -772,6 +877,10 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · FLUX.2 Klein Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_flux2_klein.py
|
||||
- tests/e2e/online_serving/test_flux2_klein.py
|
||||
- vllm_omni/diffusion/models/flux2_klein/
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_flux2_klein.py tests/e2e/online_serving/test_flux2_klein.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
@@ -789,6 +898,9 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · SD3 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_sd3.py
|
||||
- vllm_omni/diffusion/models/sd3/
|
||||
timeout_in_minutes: 45
|
||||
commands:
|
||||
- |
|
||||
|
||||
@@ -568,6 +568,55 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
|
||||
key: nightly-tts-performance-higgs-audio-v3
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
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
|
||||
|
||||
# Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
@@ -591,6 +640,7 @@ steps:
|
||||
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_flux_2_dev_expansion.py
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
tests/e2e/offline_inference/test_sensenova_u1_expansion.py
|
||||
-m "full_model and diffusion and H100 and not distributed_cuda"
|
||||
--run-level "full_model"
|
||||
agents:
|
||||
@@ -1462,6 +1512,7 @@ steps:
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-omni-performance-multi-replicas
|
||||
- nightly-tts-performance
|
||||
- nightly-tts-performance-higgs-audio-v3
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
@@ -1475,6 +1526,7 @@ steps:
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
|
||||
|
||||
+115
-64
@@ -215,10 +215,78 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- group: ":card_index_dividers: Entrypoints Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: "Entrypoints Test with H100"
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -sv tests/entrypoints/ -m 'core_model and cuda and H100' --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: "Entrypoints Test with L4"
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -sv tests/entrypoints/ -m 'core_model and cuda and L4' --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"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
- label: "RLHF · VeRL-Omni E2E Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/rlhf_test/test_verl_omni_e2e.py
|
||||
- tests/e2e/offline_inference/custom_pipeline/
|
||||
- vllm_omni/diffusion/
|
||||
- vllm_omni/engine/
|
||||
- vllm_omni/entrypoints/
|
||||
- vllm_omni/inputs/
|
||||
- vllm_omni/outputs/
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -239,6 +307,11 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- vllm_omni/model_executor/models/qwen3_omni/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -273,73 +346,18 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Qwen3-Omni Realtime WebSocket"
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -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 · VoxCPM2 Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.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"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · VoxCPM2 Online Test"
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_voxcpm2_tts.py
|
||||
- vllm_omni/model_executor/models/voxcpm2/
|
||||
- vllm_omni/deploy/voxcpm2.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
@@ -356,6 +374,11 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -363,7 +386,7 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -377,6 +400,11 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Qwen3-TTS Base Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_base.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -386,7 +414,7 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -400,6 +428,11 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Higgs-Audio-V3 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_higgs_audio_v3.py
|
||||
- vllm_omni/model_executor/models/higgs_audio_v3/
|
||||
- vllm_omni/model_executor/stage_input_processors/higgs_audio_v3.py
|
||||
- vllm_omni/deploy/higgs_multimodal_qwen3.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -409,7 +442,7 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
@@ -423,6 +456,9 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -458,6 +494,12 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Bagel Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_bagel.py
|
||||
- vllm_omni/model_executor/models/bagel/
|
||||
- vllm_omni/diffusion/models/bagel/
|
||||
- vllm_omni/model_executor/stage_input_processors/bagel.py
|
||||
- vllm_omni/deploy/bagel.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -499,6 +541,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
- vllm_omni/diffusion/models/wan2_2/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -541,6 +586,9 @@ steps:
|
||||
|
||||
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -582,6 +630,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Qwen Image Layered Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_layered.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
|
||||
@@ -6,6 +6,13 @@
|
||||
{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
|
||||
{% set default_working_dir = "/app/vllm-omni" %}
|
||||
|
||||
{% if env %}
|
||||
env:
|
||||
{% for name in env %}
|
||||
{{ name }}: "{{ env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
steps:
|
||||
- group: "AMD Tests"
|
||||
depends_on: ~
|
||||
steps:
|
||||
@@ -31,8 +38,15 @@
|
||||
agents:
|
||||
queue: amd-cpu
|
||||
|
||||
{% for step in steps %}
|
||||
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
|
||||
{% for step_group in steps %}
|
||||
{% if step_group.group %}
|
||||
- group: "{{ step_group.group }}"
|
||||
{% if step_group.depends_on %}
|
||||
depends_on: {{ step_group.depends_on }}
|
||||
{% endif %}
|
||||
steps:
|
||||
{% for step in step_group.steps %}
|
||||
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
|
||||
- label: "{{ step.agent_pool }}: {{ step.label }}"
|
||||
depends_on: amd-build
|
||||
agents:
|
||||
@@ -46,8 +60,23 @@
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
{{ name }}: "{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
TEST_COMMAND: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{% if env %}
|
||||
{% for name in env %}
|
||||
export {{ name }}="{{ env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
export {{ name }}="{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if "mi250" in step.agent_pool %}
|
||||
python3 -m pip uninstall -y amd-aiter
|
||||
{% endif %}
|
||||
@@ -61,5 +90,52 @@
|
||||
{% else %}
|
||||
soft_fail: true
|
||||
{% endif%}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
{% elif step_group.mirror_hardwares and mirror_hw in step_group.mirror_hardwares %}
|
||||
{% set step = step_group %}
|
||||
- label: "{{ step.agent_pool }}: {{ step.label }}"
|
||||
depends_on: amd-build
|
||||
agents:
|
||||
{% if step.agent_pool %}
|
||||
queue: amd_{{ step.agent_pool }}
|
||||
{% else %}
|
||||
queue: amd_mi325_1
|
||||
{% endif %}
|
||||
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
|
||||
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
{{ name }}: "{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
TEST_COMMAND: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{% if env %}
|
||||
{% for name in env %}
|
||||
export {{ name }}="{{ env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
export {{ name }}="{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if "mi250" in step.agent_pool %}
|
||||
python3 -m pip uninstall -y amd-aiter
|
||||
{% endif %}
|
||||
{{ indented_cmd | safe }}
|
||||
priority: 100
|
||||
{% if step.timeout_in_minutes %}
|
||||
timeout_in_minutes: {{ step.timeout_in_minutes }}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
{% if step.grade and step.grade == "Blocking" %}
|
||||
soft_fail: false
|
||||
{% else %}
|
||||
soft_fail: true
|
||||
{% endif%}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
@@ -0,0 +1,103 @@
|
||||
---
|
||||
name: quantization
|
||||
description: Work on vLLM-Omni quantization for diffusion, autoregressive, omni, or multi-stage models. Use when choosing or adding methods such as fp8, int8, gguf, mxfp8, mxfp4, mxfp4_dualscale, ModelOpt, AutoRound, INC, msModelSlim, awq, or gptq; debugging quantized loading; or validating memory, speed, and output quality.
|
||||
---
|
||||
|
||||
# vLLM-Omni Quantization
|
||||
|
||||
Use this skill for quantization work in `vllm-omni`. Start from the local
|
||||
code and docs, then pick the smallest validated path for the model, method,
|
||||
hardware, and quality target.
|
||||
|
||||
## First Checks
|
||||
|
||||
Before changing code or recommending a command, identify:
|
||||
|
||||
- Model and task: AR text, omni/TTS, text-to-image, text-to-video, or multi-stage diffusion.
|
||||
- Hardware backend: CUDA generation, Ascend NPU, Intel XPU, or another platform.
|
||||
- Quantization mode: online from BF16/FP16 weights, pre-quantized checkpoint, runtime KV/FA quantization, or per-component routing.
|
||||
- Scope: transformer/DiT, thinker language model, text encoder, VAE, audio/vision encoder, talker, or a specific stage.
|
||||
- Validation target: memory reduction, latency/throughput, output quality, or all three.
|
||||
|
||||
Do not assume a method is supported just because the CLI accepts the string.
|
||||
Check `docs/user_guide/quantization/` and the closest implementation first.
|
||||
|
||||
## Quick Decision
|
||||
|
||||
| Task | Start With |
|
||||
|------|------------|
|
||||
| Choose a method or command | `references/methods.md` and `references/modality-compat.md` |
|
||||
| Use `build_quant_config()` or per-component routing | `references/methods.md` |
|
||||
| Work on diffusion quantization | `references/diffusion.md` |
|
||||
| Add quantization to a new model | `references/adding-models.md` |
|
||||
| Convert or load ModelOpt FP8 checkpoints | `references/modelopt-fp8.md` |
|
||||
| Debug ModelOpt, GGUF, AutoRound, MXFP, or serialized Int8 loading | `references/diffusion.md` and method docs under `docs/user_guide/quantization/` |
|
||||
|
||||
## Current Runtime Shape
|
||||
|
||||
The unified entrypoint is:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
```
|
||||
|
||||
It supports method strings, flat method dictionaries, per-component
|
||||
dictionaries, existing `QuantizationConfig` objects, and `None`. The factory
|
||||
delegates generic methods to upstream `vllm` and keeps vLLM-Omni overrides for
|
||||
diffusion or omni-specific routing:
|
||||
|
||||
- `gguf`
|
||||
- `int8`
|
||||
- `mxfp8`
|
||||
- `mxfp4`
|
||||
- `mxfp4_dualscale`
|
||||
- `inc`, `auto-round`, `auto_round`
|
||||
- ModelOpt auto-detected configs: `modelopt`, `modelopt_fp4`, `modelopt_mixed`
|
||||
|
||||
Use `ComponentQuantizationConfig` when only one stage or component should be
|
||||
quantized. Pre-quantized ModelOpt-style checkpoints should not spill into
|
||||
vision/audio encoders that have no corresponding scale tensors.
|
||||
|
||||
## Ownership Boundary
|
||||
|
||||
- Upstream `vllm` owns generic quantization configs, kernels, loader semantics,
|
||||
hardware capability rules, and generic AR quantization methods.
|
||||
- `vllm-omni` owns unified config routing, diffusion-specific wrappers,
|
||||
component scoping, GGUF or ModelOpt checkpoint adapters, model-specific
|
||||
prefix mapping, docs, examples, and validation.
|
||||
|
||||
If a new method needs missing generic kernels or loader behavior, fix upstream
|
||||
`vllm` first. In `vllm-omni`, add thin integration and model wiring.
|
||||
|
||||
## Working Loop
|
||||
|
||||
1. Read the local method doc under `docs/user_guide/quantization/`.
|
||||
2. Inspect the active factory and config class in `vllm_omni/quantization/`.
|
||||
3. Inspect the target model pipeline and transformer for stable prefixes and
|
||||
whether every relevant vLLM linear layer receives `quant_config`.
|
||||
4. For pre-quantized checkpoints, inspect `transformer/config.json` and the
|
||||
checkpoint tensor names before running full generation.
|
||||
5. Validate with the same prompt, seed, size, scheduler, steps, dtype, eager or
|
||||
non-eager mode, and parallelism as the BF16 baseline.
|
||||
6. Report quality, latency, and memory separately. Do not call a method
|
||||
supported until quality is checked.
|
||||
|
||||
## Common Mistakes
|
||||
|
||||
| Symptom | Likely Cause | Fix |
|
||||
|---------|--------------|-----|
|
||||
| `--quantization` has no visible effect | Wrong scope or unsupported model path | Check component routing and method docs |
|
||||
| Some layers stay BF16 unexpectedly | `quant_config` was not threaded into all vLLM linear layers | Audit transformer constructors and prefixes |
|
||||
| Quality collapses but loading succeeds | Too many sensitive layers were quantized | Add model-specific `ignored_layers` and compare to BF16 |
|
||||
| ModelOpt checkpoint loads but output is corrupted | Prefixes, packed-module mapping, scale routing, or BF16 fallback is wrong | Use `references/modelopt-fp8.md` |
|
||||
| GGUF shape or tensor mismatch | Missing architecture-specific adapter | Add explicit adapter mapping; avoid generic fallback |
|
||||
| Online and offline MXFP4 disagree | Wrong `mxfp4` vs `mxfp4_dualscale` mode | Check `docs/user_guide/quantization/mxfp4.md` |
|
||||
| Quantized path is slower than BF16 | Kernel path, compile mode, dtype, or shape mismatch | Compare eager-to-eager and non-eager-to-non-eager |
|
||||
|
||||
## References
|
||||
|
||||
- General methods and config routing: [references/methods.md](references/methods.md)
|
||||
- Modality and model scope: [references/modality-compat.md](references/modality-compat.md)
|
||||
- Diffusion implementation patterns: [references/diffusion.md](references/diffusion.md)
|
||||
- Adding model support: [references/adding-models.md](references/adding-models.md)
|
||||
- ModelOpt FP8 conversion and loading: [references/modelopt-fp8.md](references/modelopt-fp8.md)
|
||||
@@ -0,0 +1,70 @@
|
||||
# Adding Quantization Support to a Model
|
||||
|
||||
Use this checklist when extending quantization to a new architecture or stage.
|
||||
|
||||
## Classify Ownership
|
||||
|
||||
| Case | Owner |
|
||||
|------|-------|
|
||||
| New generic AR method, kernel, or loader semantics | upstream `vllm` first |
|
||||
| Existing upstream method needs diffusion wiring | `vllm-omni` model and config integration |
|
||||
| GGUF or ModelOpt tensor mapping for a diffusion architecture | `vllm-omni` adapter |
|
||||
| NPU MXFP or msModelSlim diffusion path | `vllm-omni` platform/method wrapper and docs |
|
||||
| Multi-stage omni checkpoint quantizes only the thinker LM | `vllm-omni` component scoping |
|
||||
|
||||
Avoid building a private quantization stack in model code.
|
||||
|
||||
## Model Wiring Checklist
|
||||
|
||||
1. Find the closest supported model implementation.
|
||||
2. Confirm the model uses vLLM linear layers that accept `quant_config`.
|
||||
3. Add `quant_config` to the transformer constructor.
|
||||
4. Pass stable `prefix` values through every linear layer.
|
||||
5. Handle fused modules and packed mappings explicitly.
|
||||
6. Ensure `load_weights()` and any `WeightsMapper` remap `ignored_layers` or
|
||||
`modules_to_not_convert`.
|
||||
7. Keep tokenizer, scheduler, VAE, text encoder, audio/vision encoders, and
|
||||
talker BF16 unless method-specific docs validate them.
|
||||
8. Add or update a method doc under `docs/user_guide/quantization/`.
|
||||
|
||||
## Checkpoint Adapter Checklist
|
||||
|
||||
For pre-quantized checkpoints:
|
||||
|
||||
- Inspect `transformer/config.json` first.
|
||||
- Check `quant_method`, `method`, `producer.name`, `quant_algo`, serialized
|
||||
flags, `ignored_layers`, and any method-specific fields.
|
||||
- Verify tensor name mapping before generation.
|
||||
- Verify packed modules such as `to_qkv`, `add_kv_proj`, and `w13`.
|
||||
- Confirm scale tensor names and shapes.
|
||||
- Allow dequantization back to BF16 for intentionally skipped target layers.
|
||||
|
||||
## Sensitive Layer Discovery
|
||||
|
||||
Start conservative. Quantize the backbone first and keep output-side or routing
|
||||
layers BF16 until quality is proven.
|
||||
|
||||
Frequently sensitive areas:
|
||||
|
||||
- attention output projections such as `to_out`
|
||||
- FFN down projections such as `w2`
|
||||
- final projections such as `proj_out`, `final_layer`, and output heads
|
||||
- modulation, timestep, position, image, or text embedders
|
||||
- refiner blocks
|
||||
- MoE routers, experts selected by routing, and layernorm-adjacent modules
|
||||
- VAE, text/vision towers, audio encoders, and codec/talker stages
|
||||
|
||||
Do not reuse one model family's ignore list blindly.
|
||||
|
||||
## Tests and Evidence
|
||||
|
||||
Minimum evidence for a PR:
|
||||
|
||||
- loader/config test or smoke command
|
||||
- fixed-seed BF16 vs quantized output comparison
|
||||
- quality metrics and saved outputs for precision-changing work
|
||||
- memory and latency comparison under identical execution mode
|
||||
- failure behavior for unsupported models or missing checkpoint fields
|
||||
|
||||
When the PR adds a method to a support table, include the exact model,
|
||||
hardware, command, quality threshold, and artifact path used for validation.
|
||||
@@ -0,0 +1,113 @@
|
||||
# Diffusion Quantization Reference
|
||||
|
||||
Use this file for DiT quantization, config routing, checkpoint adapters, and
|
||||
quality or loader debugging.
|
||||
|
||||
## Supported Method Families
|
||||
|
||||
The current local diffusion quantization families are:
|
||||
|
||||
- CUDA/GPU online or checkpoint: `fp8`, `int8`, `gguf`, ModelOpt
|
||||
- Ascend NPU: `int8`, `mxfp8`, `mxfp4`, `mxfp4_dualscale`, msModelSlim paths
|
||||
- Intel XPU: AutoRound/INC and online MXFP8 paths where documented
|
||||
|
||||
Always cross-check the method doc before claiming model support. The CLI
|
||||
accepting a method string is not enough.
|
||||
|
||||
## Implementation Starting Points
|
||||
|
||||
| Area | Files |
|
||||
|------|-------|
|
||||
| Unified factory | `vllm_omni/quantization/factory.py` |
|
||||
| Per-component routing | `vllm_omni/quantization/component_config.py` |
|
||||
| GGUF diffusion method | `vllm_omni/quantization/gguf_config.py` |
|
||||
| Int8 diffusion method | `vllm_omni/quantization/int8_config.py` |
|
||||
| MXFP8 diffusion method | `vllm_omni/quantization/mxfp8_config.py` |
|
||||
| MXFP4 and dual-scale method | `vllm_omni/quantization/mxfp4_config.py` |
|
||||
| AutoRound/INC omni extension | `vllm_omni/quantization/inc_config.py` |
|
||||
| ModelOpt checkpoint adapters | `vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py` |
|
||||
| FP8 scaled-mm runtime patches | `vllm_omni/patch.py` |
|
||||
| Similarity tool | `vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` |
|
||||
|
||||
## Threading `quant_config`
|
||||
|
||||
For a diffusion transformer:
|
||||
|
||||
1. Accept `quant_config: QuantizationConfig | None` in the transformer.
|
||||
2. Pass `quant_config` and stable `prefix` values to every vLLM linear layer.
|
||||
3. Preserve packed-module mappings for fused QKV, FFN, or model-specific fused
|
||||
projections.
|
||||
4. Call `apply_vllm_mapper()` on configs that need HF-to-vLLM name remapping.
|
||||
5. For component configs, keep prefixes specific enough to avoid quantizing VAE
|
||||
or encoders by accident.
|
||||
|
||||
Common skipped areas are `proj_out`, output projections, modulation layers,
|
||||
embedders, refiner blocks, MoE routers, and layernorm-adjacent paths.
|
||||
|
||||
## Online vs Offline
|
||||
|
||||
Online quantization:
|
||||
|
||||
- starts from BF16/FP16 checkpoint weights
|
||||
- computes or stores quantized weights at load time
|
||||
- typically uses `--quantization <method>` or `quantization="<method>"`
|
||||
|
||||
Offline or pre-quantized checkpoints:
|
||||
|
||||
- include quantized weights and scales on disk
|
||||
- usually carry `quantization_config` in `transformer/config.json`
|
||||
- should auto-detect when the config is complete
|
||||
- may require model-specific adapter logic for tensor names and packed modules
|
||||
|
||||
Do not mix online flags with incompatible offline checkpoint configs. If the
|
||||
active config and disk config differ, `resolve_quant_config_from_disk()` should
|
||||
raise or rebuild rather than silently loading corrupted weights.
|
||||
|
||||
## GGUF Rules
|
||||
|
||||
- Require `quantization_config.gguf_model`.
|
||||
- Use the normal base model for tokenizer, scheduler, text encoder, and VAE.
|
||||
- Add explicit architecture adapters for tensor-name mapping.
|
||||
- Use `named_parameters()` and `named_buffers()`, not `state_dict()`, to
|
||||
discover loadable names.
|
||||
- Guard fused QKV/KV rewrites so tensors are not rewritten twice.
|
||||
- Flatten N-D activations around the GGUF matmul and restore shape.
|
||||
|
||||
Unsupported architectures should fail clearly instead of falling back to a fake
|
||||
generic mapper.
|
||||
|
||||
## MXFP Rules
|
||||
|
||||
- `mxfp8` is W8A8 microscaling, currently centered on NPU and documented XPU
|
||||
behavior.
|
||||
- `mxfp4` is online single-scale W4A4 on NPU.
|
||||
- `mxfp4_dualscale` is online/offline dual-scale W4A4 with BF16 fallback; the
|
||||
offline path is the production path when calibrated `mul_scale` exists.
|
||||
- Offline dual-scale checkpoints must inject per-transformer `ignored_layers`
|
||||
into config because cascade transformers can have different BF16 fallback
|
||||
sets.
|
||||
|
||||
Do not load an offline dual-scale checkpoint with the online single-scale path.
|
||||
|
||||
## Common Failures
|
||||
|
||||
| Symptom | Likely Cause | Fix |
|
||||
|---------|--------------|-----|
|
||||
| `Unknown quantization method` | Factory override or upstream registration is missing | Add the right override or upstream method |
|
||||
| Missing scale tensors | Wrong scope, wrong adapter, or BF16 layer not allowed to dequantize | Check config scope and adapter scale routing |
|
||||
| Output rank collapses from 3-D to 2-D | FP8 scaled-mm kernel ignored output shape | Check the patch in `vllm_omni/patch.py` |
|
||||
| Quality degrades silently | Sensitive layers quantized | Add `ignored_layers`; compare to BF16 |
|
||||
| Only one cascade transformer works | Disk config differs per transformer but active config is reused | Rebuild from each transformer's config |
|
||||
| Quantized path slower than BF16 | Kernel, dtype, compile, or benchmark mode mismatch | Compare eager-to-eager and non-eager-to-non-eager |
|
||||
|
||||
## Verification
|
||||
|
||||
For a new or changed diffusion quantization path, run:
|
||||
|
||||
1. Config and loader smoke test.
|
||||
2. One fixed-seed generation against BF16.
|
||||
3. Numeric comparison with the similarity tool or equivalent metrics.
|
||||
4. Serving test if the path is user-facing online.
|
||||
5. Memory and latency check under the same execution mode as the baseline.
|
||||
|
||||
Keep output artifacts when precision changes are part of the PR.
|
||||
@@ -0,0 +1,82 @@
|
||||
# General Quantization Methods
|
||||
|
||||
Use this file for method selection and unified `vllm-omni` config syntax. For
|
||||
diffusion-specific wiring, see `diffusion.md`. For ModelOpt checkpoint export
|
||||
and loading details, see `modelopt-fp8.md`.
|
||||
|
||||
## Unified Entry Point
|
||||
|
||||
`vllm_omni.quantization.build_quant_config()` accepts:
|
||||
|
||||
- `None` or `"none"` to disable quantization
|
||||
- a method string such as `"fp8"`, `"int8"`, `"mxfp8"`, `"awq"`, or `"auto-round"`
|
||||
- a flat dict with `"method"` or `"quant_method"`
|
||||
- a per-component dict such as `{"transformer": {"method": "fp8"}, "vae": None}`
|
||||
- an already-built `QuantizationConfig`
|
||||
|
||||
Per-component dicts are routed by `ComponentQuantizationConfig` using
|
||||
longest-prefix matching. Prefix stability matters: if a model remaps HF names
|
||||
to vLLM names, apply the mapper to `ignored_layers` or component prefixes.
|
||||
|
||||
## Local Overrides
|
||||
|
||||
The local factory delegates most methods to upstream `vllm`, but has vLLM-Omni
|
||||
overrides for methods that need diffusion, NPU, XPU, or component behavior:
|
||||
|
||||
| Method | Local Config | Main Use |
|
||||
|--------|--------------|----------|
|
||||
| `gguf` | `DiffusionGGUFConfig` | Pre-quantized diffusion transformer weights |
|
||||
| `int8` | `DiffusionInt8Config` | Online or serialized W8A8 diffusion transformers |
|
||||
| `mxfp8` | `DiffusionMXFP8Config` | W8A8 MXFP8 diffusion on NPU and online XPU path |
|
||||
| `mxfp4` | `DiffusionMXFP4Config` | Online W4A4 MXFP4 diffusion on NPU |
|
||||
| `mxfp4_dualscale` | `DiffusionMXFP4DualScaleMixedConfig` | Online/offline W4A4 dual-scale plus BF16 fallback on NPU |
|
||||
| `inc`, `auto-round`, `auto_round` | `OmniINCConfig` | AutoRound/INC checkpoint loading, including MXFP8 checkpoint metadata |
|
||||
| ModelOpt config dicts | upstream ModelOpt configs through auto-detect | FP8, NVFP4, or mixed ModelOpt checkpoints |
|
||||
|
||||
Check `SUPPORTED_QUANTIZATION_METHODS` in `vllm_omni/quantization/factory.py`
|
||||
before adding a new method name.
|
||||
|
||||
## AR and Generic Quantization
|
||||
|
||||
For AR-backed models, upstream `vllm` methods usually apply:
|
||||
|
||||
- `awq` and `gptq`: 4-bit or 8-bit weight quantization with calibration,
|
||||
commonly used for text and AR stages.
|
||||
- `fp8`: FP8 weight path or KV-cache FP8 depending on flags and hardware.
|
||||
- `inc` or `auto-round`: pre-quantized checkpoints with Intel Neural
|
||||
Compressor / AutoRound metadata.
|
||||
|
||||
Use these for the AR language model stage only when the checkpoint and model
|
||||
guide support it. Omni and TTS models often keep audio, vision, talker, and
|
||||
codec stages in BF16.
|
||||
|
||||
## Diffusion and Offline Methods
|
||||
|
||||
For diffusion transformers, prefer the method-specific docs:
|
||||
|
||||
- `docs/user_guide/quantization/fp8.md`
|
||||
- `docs/user_guide/quantization/int8.md`
|
||||
- `docs/user_guide/quantization/gguf.md`
|
||||
- `docs/user_guide/quantization/modelopt.md`
|
||||
- `docs/user_guide/quantization/mxfp8.md`
|
||||
- `docs/user_guide/quantization/mxfp4.md`
|
||||
- `docs/user_guide/quantization/autoround.md`
|
||||
- `docs/user_guide/quantization/msmodelslim.md`
|
||||
|
||||
Pre-quantized checkpoints should normally auto-detect from config. Runtime
|
||||
flags such as `--force-cutlass-fp8`, `--linear-backend cutlass`, and
|
||||
`--moe-backend cutlass` select kernels; they do not quantize a BF16 checkpoint.
|
||||
|
||||
## Validation
|
||||
|
||||
Use fixed-seed A/B comparisons and record:
|
||||
|
||||
- model path and quantized checkpoint path
|
||||
- method config and `ignored_layers`
|
||||
- prompt, size, frame count, steps, scheduler, dtype, and seed
|
||||
- eager/non-eager mode, compile flags, and parallelism
|
||||
- output metrics such as PSNR, MAE, cosine similarity, SSIM, LPIPS, plus visual artifacts
|
||||
- latency and peak memory, reported separately from quality
|
||||
|
||||
The helper `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
|
||||
is the default lightweight image/video comparison tool for diffusion paths.
|
||||
@@ -0,0 +1,47 @@
|
||||
# Quantization Compatibility by Modality
|
||||
|
||||
Use this as a routing guide. The source of truth is the local method docs under
|
||||
`docs/user_guide/quantization/` and the current model implementation.
|
||||
|
||||
## Summary
|
||||
|
||||
- AR and omni language-model quantization mostly follows upstream `vllm`, with
|
||||
vLLM-Omni scoping for multi-stage models.
|
||||
- Diffusion quantization is method-specific and model-specific. It usually
|
||||
targets the diffusion transformer only.
|
||||
- Pre-quantized checkpoints should quantize only the component that has
|
||||
quantized weights and scale tensors.
|
||||
|
||||
## Model Type Matrix
|
||||
|
||||
| Model Type | Typical Quantization Scope | Notes |
|
||||
|------------|----------------------------|-------|
|
||||
| AR text models | Language model | Upstream `vllm` methods usually apply |
|
||||
| Omni/TTS models | Thinker or AR language-model stage | Audio, vision, talker, and codec stages usually stay BF16 |
|
||||
| DiT image models | Diffusion transformer | Check FP8, Int8, GGUF, ModelOpt, AutoRound docs |
|
||||
| DiT video models | Diffusion transformer | HunyuanVideo FP8 and Wan2.2 MXFP paths are documented; validate other video paths |
|
||||
| Multi-stage diffusion | Selected stage only | Use per-component routing and per-stage validation |
|
||||
| Audio diffusion | Usually unsupported | Prefer dtype, offload, cache, or parallelism unless docs say otherwise |
|
||||
|
||||
## Method Hints
|
||||
|
||||
| Goal | Candidate |
|
||||
|------|-----------|
|
||||
| AR memory reduction | AWQ, GPTQ, AutoRound, ModelOpt where checkpoint supports it |
|
||||
| CUDA diffusion W8A8 | FP8 or Int8 |
|
||||
| Pre-quantized CUDA diffusion checkpoint | ModelOpt or GGUF |
|
||||
| Ascend NPU diffusion | Int8, MXFP8, MXFP4, MXFP4 DualScale, msModelSlim |
|
||||
| Intel XPU checkpoint | AutoRound/INC paths |
|
||||
| Conservative image quality | selective FP8/Int8 with `ignored_layers` |
|
||||
| Unsupported diffusion model | cache acceleration, CPU/layerwise offload, dtype, TP/SP/CFG parallelism |
|
||||
|
||||
## Validation Rule
|
||||
|
||||
Compatibility means more than successful loading. A method is compatible only
|
||||
after it has:
|
||||
|
||||
- correct component scope
|
||||
- complete checkpoint/config metadata if offline
|
||||
- fixed-seed quality comparison against BF16
|
||||
- acceptable latency and memory behavior
|
||||
- docs or examples that state model, method, hardware, and caveats
|
||||
@@ -0,0 +1,122 @@
|
||||
# ModelOpt FP8 Conversion and Loading
|
||||
|
||||
Use this when exporting or loading ModelOpt FP8 checkpoints for diffusion
|
||||
transformers. This reference captures the practical conversion lessons from the
|
||||
2026-04-21 ModelOpt FP8 paste and aligns them with the current
|
||||
`vllm-omni` ModelOpt adapter.
|
||||
|
||||
## Core Rule
|
||||
|
||||
Treat ModelOpt FP8 as selective checkpoint quantization, not one-click
|
||||
whole-model FP8. Quantize the main transformer or DiT backbone, keep sensitive
|
||||
or peripheral modules BF16, serialize a complete ModelOpt checkpoint, and let
|
||||
vLLM-Omni auto-detect it from config.
|
||||
|
||||
Do not pass `--quantization fp8` to force a ModelOpt checkpoint path. The
|
||||
checkpoint's `quantization_config` should select `modelopt`.
|
||||
|
||||
## Conversion Strategy
|
||||
|
||||
Prefer:
|
||||
|
||||
- backbone attention Q/K/V and FFN expansion projections
|
||||
- conservative per-model ignore lists
|
||||
- MHA quantizers disabled until basic linear quantization is stable
|
||||
- full module-name prefixes in ignore/routing logic
|
||||
- BF16 fallback for output-side and routing-sensitive layers
|
||||
|
||||
Avoid quantizing by default:
|
||||
|
||||
- VAE
|
||||
- text, vision, or audio encoders
|
||||
- embedders and modulation layers
|
||||
- final heads and output projections
|
||||
- refiner modules
|
||||
- MoE routers and layernorm-sensitive blocks
|
||||
|
||||
## Model-Specific Starting Points
|
||||
|
||||
| Model Family | Quantize First | Keep BF16 First |
|
||||
|--------------|----------------|-----------------|
|
||||
| Z-Image | attention q/k/v, FFN w1/w3 | `to_out`, `w2`, `noise_refiner`, `context_refiner`, embedders, modulation, final/output |
|
||||
| Qwen-Image | transformer backbone linear layers | `proj_out`, `img_in`, `txt_in`, norms, modulation, position/time/embed layers |
|
||||
| FLUX2-dev / FLUX2-klein | backbone linear layers | `proj_out`, context/x embedders, `norm_out`, stream modulation, time embeddings |
|
||||
| HunyuanImage-3.0 | `model.model` backbone | VAE, vision tower, token/time/image embedders, final image head, LM head, routers, layernorm-sensitive blocks |
|
||||
|
||||
These are starting points, not support claims. Validate with BF16 outputs.
|
||||
|
||||
## Required Checkpoint Metadata
|
||||
|
||||
`transformer/config.json` must include a complete `quantization_config`.
|
||||
At minimum it should identify ModelOpt and the algorithm:
|
||||
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"producer": {"name": "modelopt"},
|
||||
"quant_method": "modelopt",
|
||||
"quant_algo": "FP8"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Current vLLM-Omni auto-detects:
|
||||
|
||||
- `FP8` and `FP8_PER_CHANNEL_PER_TOKEN` -> `modelopt`
|
||||
- `NVFP4` -> `modelopt_fp4`
|
||||
- `MIXED_PRECISION` -> `modelopt_mixed`
|
||||
|
||||
Packed modules and scales must map correctly:
|
||||
|
||||
- `to_qkv` from `to_q`, `to_k`, `to_v`
|
||||
- `add_kv_proj` from `add_q_proj`, `add_k_proj`, `add_v_proj`
|
||||
- `w13` from `w1`, `w3`
|
||||
- `.input_scale`, `.weight_scale`, `.weight_scale_2`, `.weight_scale_inv`
|
||||
|
||||
For selective FP8, target layers that remain BF16 must be allowed to
|
||||
dequantize from FP8 weights using the matching `weight_scale`, or the adapter
|
||||
must skip the quantized scale tensors for that BF16 target.
|
||||
|
||||
## Runtime Signals
|
||||
|
||||
Healthy loading should show ModelOpt auto-detection and a ModelOpt linear
|
||||
kernel path. Useful signals include:
|
||||
|
||||
- auto-detected quantization from config
|
||||
- selected ModelOpt FP8 linear method
|
||||
- CUTLASS or FlashInfer FP8 scaled-mm backend when enabled and supported
|
||||
- non-eager path compiling the transformer when expected
|
||||
|
||||
If online serving fails while offline inference works, inspect warmup and
|
||||
serving initialization before changing checkpoint conversion.
|
||||
|
||||
## Validation Order
|
||||
|
||||
1. Check config completeness.
|
||||
2. Run an export checker such as `examples/quantization/check_modelopt_fp8_export.py`.
|
||||
3. Confirm actual FP8 tensors and expected disk-size reduction.
|
||||
4. Run one offline generation with fixed prompt, seed, size, and steps.
|
||||
5. Run serving with `vllm serve --omni` and verify `/v1/images/generations`.
|
||||
6. Compare against BF16 using visual review plus metrics such as LPIPS, PSNR,
|
||||
MAE, cosine similarity, and SSIM.
|
||||
7. Compare performance only under matched modes:
|
||||
- BF16 eager vs FP8 eager
|
||||
- BF16 non-eager vs FP8 non-eager
|
||||
|
||||
FP8 is not automatically faster. Check whether execution reaches FP8 kernels,
|
||||
whether compile mode is comparable, and whether the workload shape fits the
|
||||
kernel path.
|
||||
|
||||
## Triage
|
||||
|
||||
| Symptom | First Check |
|
||||
|---------|-------------|
|
||||
| Startup failure | `config.json`, `model_index.json`, `quantization_config`, missing scales |
|
||||
| Loads but output is corrupted | Over-quantized sensitive layers or wrong full-name prefixes |
|
||||
| Offline works but online fails | Dummy warmup, serving init, compile path |
|
||||
| FP8 is not faster | Kernel selection, dtype, compile mode, workload shape |
|
||||
| BF16 fallback fails | Missing `weight_scale` for dequantization or wrong packed-module mapping |
|
||||
|
||||
The success criterion is stable, correctly scoped checkpoint loading with
|
||||
quality preserved enough for the target use case, not maximizing the number of
|
||||
FP8 layers.
|
||||
@@ -21,6 +21,8 @@ include:
|
||||
- `diffusion-perf-opt`: guides diffusion model performance optimization,
|
||||
including profiling traces, parallel strategies, stage timing analysis, and
|
||||
benchmark-driven tuning
|
||||
- `quantization`: guides quantization method selection, model integration,
|
||||
checkpoint loading, and quality/performance validation for vLLM-Omni
|
||||
- `add-omni-model`: covers addition of new omni-modality model support
|
||||
- `add-tts-model`: covers integration of new TTS models and related serving
|
||||
workflows
|
||||
|
||||
@@ -1,24 +1,15 @@
|
||||
<!-- markdownlint-disable -->
|
||||
**Before submitting:** run the precheck-pr skill with code agent for a self-check against project conventions — catches dead code, missing benchmarks, and title format issues.
|
||||
|
||||
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
|
||||
PLEASE FILL IN THE PR DESCRIPTION HERE.
|
||||
|
||||
## Purpose
|
||||
|
||||
## Test Plan
|
||||
|
||||
**vLLM Version:**
|
||||
|
||||
**vLLM-Omni Commit:**
|
||||
|
||||
## Test Result
|
||||
|
||||
---
|
||||
<details>
|
||||
<summary> Essential Elements of an Effective PR Description Checklist </summary>
|
||||
|
||||
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
|
||||
- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
|
||||
- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
|
||||
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
|
||||
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
|
||||
</details>
|
||||
|
||||
**BEFORE SUBMITTING,** please read <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>
|
||||
**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
|
||||
(anything written below this line will be removed by GitHub Actions)
|
||||
@@ -1,3 +1,5 @@
|
||||
# Contributing to vLLM-Omni
|
||||
|
||||
You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
|
||||
|
||||
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
|
||||
@@ -57,8 +57,9 @@ vLLM-Omni is flexible and easy to use with:
|
||||
|
||||
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
|
||||
|
||||
- Omni-modality models (e.g. Qwen-Omni)
|
||||
- Multi-modality generation models (e.g. Qwen-Image)
|
||||
- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
|
||||
- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
|
||||
- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
|
||||
|
||||
## Getting Started
|
||||
|
||||
|
||||
@@ -33,6 +33,15 @@ models:
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
|
||||
bosonai/higgs-audio-v3-tts-4b:
|
||||
supported_tasks: [default_voice, voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
trust_remote_code: true
|
||||
task_extra_body:
|
||||
default_voice: {}
|
||||
voice_clone: {}
|
||||
|
||||
openbmb/VoxCPM2:
|
||||
supported_tasks: [voice_clone, default_voice]
|
||||
trust_remote_code: true
|
||||
|
||||
@@ -142,15 +142,12 @@ WORKDIR /workspace/vllm-omni
|
||||
COPY . .
|
||||
|
||||
ENV VLLM_OMNI_TARGET_DEVICE=xpu
|
||||
RUN uv pip install --no-cache-dir ".[dev]" --no-build-isolation
|
||||
|
||||
# FIX triton
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
uv pip uninstall triton triton-xpu && \
|
||||
uv pip install triton-xpu==3.7.0
|
||||
|
||||
# remove torch bundled oneccl to avoid conflicts
|
||||
RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
echo "triton @ https://example.com/blocked" > /tmp/vllm_omni_constraints.txt && \
|
||||
PIP_CONSTRAINT=/tmp/vllm_omni_constraints.txt \
|
||||
uv pip install --no-build-isolation \
|
||||
--extra-index-url https://mirrors.ustc.edu.cn/pypi/simple/ \
|
||||
".[dev]" && \
|
||||
uv pip uninstall oneccl oneccl-devel
|
||||
|
||||
FROM vllm-omni AS vllm-omni-openai
|
||||
|
||||
@@ -87,6 +87,8 @@ nav:
|
||||
- AutoRound: user_guide/quantization/autoround.md
|
||||
- msModelSlim: user_guide/quantization/msmodelslim.md
|
||||
- ComfyUI: features/comfyui.md
|
||||
- Fault Tolerance:
|
||||
- Failure Modes: user_guide/fault_injection_reliability_matrix.md
|
||||
- Developer Guide:
|
||||
- General:
|
||||
- contributing/README.md
|
||||
|
||||
@@ -118,17 +118,13 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
|
||||
### Pre-Check Before Submitting
|
||||
|
||||
Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
|
||||
|
||||
```
|
||||
/precheck-pr
|
||||
```
|
||||
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-review against project conventions:
|
||||
|
||||
The skill offers two modes:
|
||||
- **Quick (~3 min):** catches showstoppers — PR title format, missing benchmark claims, rebase status
|
||||
- **Full (~10 min):** thorough maintainer-grade review — dead code scan, copy-paste detection, import hygiene
|
||||
|
||||
The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the skill in `.claude/skills/precheck-pr/` for the full checklist.
|
||||
The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Model, and General. Each type has a tailored checklist that validates evidence quality (repro steps, A/B benchmarks, registry entries, etc.). See the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) for the full checklist.
|
||||
|
||||
### Local Test
|
||||
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
|
||||
|
||||
@@ -268,6 +268,120 @@ Before entering specific testing levels, the project establishes two common spec
|
||||
1. ***PR Checklist ([Tests Style](../ci/tests_style.md))***: This template defines the self-check items that must be completed before submitting a code review (Pull Request). It ensures that each code change meets basic requirements such as code style, dependency updates, and documentation synchronization before entering the automated testing pipeline, serving as the first manual line of defense for quality assurance.
|
||||
2. ***CI Failure Explanation ([CI Failures](../ci/failures.md))***: This document archives and explains common failure patterns in the Continuous Integration (CI) pipeline, error log interpretation, and preliminary troubleshooting steps. It helps developers and testers quickly diagnose the causes of automated test failures, improving problem-solving efficiency.
|
||||
|
||||
## Notes
|
||||
|
||||
### Diff-aware Buildkite uploads (`source_file_dependencies`)
|
||||
|
||||
L2 (`.buildkite/test-ready.yml`) and L3 (`.buildkite/test-merge.yml`) pipelines can **skip unrelated GPU jobs at upload time** based on the PR diff. This is implemented by `.buildkite/scripts/upload_pipeline.py`, which filters steps before calling `buildkite-agent pipeline upload`.
|
||||
|
||||
#### What `source_file_dependencies` is
|
||||
|
||||
- A **uploader-only** YAML key on a Buildkite step or group. **Buildkite does not understand it**; `upload_pipeline.py` always **removes** it from the YAML that is uploaded.
|
||||
- A list of path **prefixes** (directories or individual files). If **any** changed file in the diff equals a prefix or starts with `prefix/`, the step is kept; otherwise the step (or entire group) is **omitted** from the uploaded pipeline.
|
||||
|
||||
#### When filtering runs
|
||||
|
||||
| Build context | Changed files used for matching |
|
||||
| --- | --- |
|
||||
| Pull request | `git diff --name-only origin/<base>...<BUILDKITE_COMMIT>` |
|
||||
| `main` branch push | `git diff --name-only <commit>^..<commit>` |
|
||||
| Other (e.g. local dry-run, non-PR branch) | Diff cannot be resolved → **no filtering**; all steps are uploaded and `source_file_dependencies` is still stripped |
|
||||
|
||||
Docs-only PRs are handled earlier in bootstrap (`pipeline.yml`) via skip-ci logic; `source_file_dependencies` applies only to the **uploaded** L2/L3 test pipelines.
|
||||
|
||||
#### Where it is configured
|
||||
|
||||
| Level | Pipeline file | Upload entry |
|
||||
| --- | --- | --- |
|
||||
| L2 | `.buildkite/test-ready.yml` | `upload_pipeline.py --upload .buildkite/test-ready.yml` (from `pipeline.yml`) |
|
||||
| L3 | `.buildkite/test-merge.yml` | `upload_pipeline.py --upload .buildkite/test-merge.yml` (from `pipeline.yml`) |
|
||||
|
||||
Steps **without** `source_file_dependencies` are always uploaded (subject to the usual label conditions: `ready` for L2, `merge-test` for L3).
|
||||
|
||||
#### Current skip policy (L2 / L3)
|
||||
|
||||
To balance CI cost and coverage:
|
||||
|
||||
- **Always run** (no `source_file_dependencies`): baseline groups outside E2E Test—e.g. Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3).
|
||||
- **Diff-gated** (`source_file_dependencies` set): **every leaf job under the E2E Test group** in `test-ready.yml` and `test-merge.yml`, regardless of queue (`mithril-h100-pool`, `gpu_1_queue`, or `gpu_4_queue`). Each step lists the smallest set of prefixes that should trigger it—typically:
|
||||
- pytest file(s) exercised by the job (online and/or offline);
|
||||
- model code under `vllm_omni/model_executor/models/` or `vllm_omni/diffusion/models/`;
|
||||
- related `vllm_omni/model_executor/stage_input_processors/` and `vllm_omni/deploy/*.yaml` when applicable.
|
||||
|
||||
Adding a new E2E step: add `source_file_dependencies` on the leaf job with those prefixes. Prefer **per-step** deps rather than a broad group-level list unless every child shares the same paths.
|
||||
|
||||
#### YAML examples
|
||||
|
||||
H100 E2E (kubernetes / `mithril-h100-pool`):
|
||||
|
||||
```yaml
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen_image.py -m 'core_model' ...
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
```
|
||||
|
||||
Docker E2E (`gpu_1_queue` / `gpu_4_queue`)—same key, same upload-time filtering:
|
||||
|
||||
```yaml
|
||||
- label: "TTS · Qwen3-TTS CustomVoice Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_customvoice.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py ...
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
```
|
||||
|
||||
A **group** may also define `source_file_dependencies`; nested steps inherit filtering as a unit—the whole group is dropped if no prefix matches.
|
||||
|
||||
#### Local dry-run
|
||||
|
||||
```bash
|
||||
# Render filtered YAML to stdout (no upload)
|
||||
python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-ready.yml
|
||||
|
||||
# Confirm uploader-only keys are stripped
|
||||
python3 .buildkite/scripts/upload_pipeline.py .buildkite/test-merge.yml | grep source_file_dependencies
|
||||
# (no output expected)
|
||||
```
|
||||
|
||||
On a PR build, Buildkite logs from `upload_pipeline.py` include lines such as `skip '…' (no changes under …)` for omitted steps.
|
||||
|
||||
#### Related
|
||||
|
||||
- Implementation: `.buildkite/scripts/upload_pipeline.py`
|
||||
- L2/L3 diff skip does **not** replace label-based triggers (`ready`, `merge-test`); it only reduces which steps appear **after** the pipeline is already scheduled.
|
||||
|
||||
### Test helper environment variables
|
||||
|
||||
Some shared helpers under `tests/helpers/` honor optional environment variables for local debugging. These are **not** set in CI by default.
|
||||
|
||||
| Variable | Accepted values | Description |
|
||||
| -------- | --------------- | ----------- |
|
||||
| `VLLM_OMNI_KEEP_REQUEST_MEDIA` | `1`, `true`, `yes` (case-insensitive) | When enabled, temporary WAV files created by `tests.helpers.media.convert_audio_bytes_to_text` are **not** deleted when the pytest process exits. By default, each call writes a unique file under the system temp directory via `tempfile.mkstemp` and registers `atexit` cleanup. Use this when debugging audio output validation (Whisper transcription, keyword checks, text–audio similarity). The saved path is logged as `audio data is saved: <path>`. |
|
||||
|
||||
Example (Linux / macOS):
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_KEEP_REQUEST_MEDIA=1
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -k test_mix_to_text_audio
|
||||
```
|
||||
|
||||
Example (Windows PowerShell):
|
||||
|
||||
```powershell
|
||||
$env:VLLM_OMNI_KEEP_REQUEST_MEDIA = "1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -k test_mix_to_text_audio
|
||||
```
|
||||
|
||||
## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification
|
||||
|
||||
### 1.1 Testing Purpose
|
||||
@@ -302,6 +416,7 @@ A clear directory structure is key to managing test cases efficiently.
|
||||
### 1.4 Execution Method and Example
|
||||
|
||||
- ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered.
|
||||
- ***Diff-aware step skipping***: On L2, **E2E Test** jobs may be omitted at pipeline upload when the PR diff does not touch their [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) prefixes (H100 and docker queues alike); non-E2E groups still always upload. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
|
||||
- ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment.
|
||||
- ***Script Example***:
|
||||
|
||||
@@ -404,7 +519,8 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
### 2.4 Execution Method and Example
|
||||
|
||||
- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch.
|
||||
- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge).
|
||||
- ***Diff-aware step skipping***: Same [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) mechanism as L2—**all E2E Test** leaf steps are diff-gated; other L3 groups always upload. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
|
||||
- ***Execution Environment***: ***GPU*** servers.
|
||||
- ***Script Example***:
|
||||
|
||||
@@ -520,6 +636,8 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation.
|
||||
|
||||
**Audio output debugging**: Deep validation may transcribe returned audio via `convert_audio_bytes_to_text` (Whisper). If an audio keyword or text–audio similarity assertion fails, set `VLLM_OMNI_KEEP_REQUEST_MEDIA=1` before running pytest to keep the intermediate WAV files for inspection (see [Test helper environment variables](#test-helper-environment-variables)).
|
||||
|
||||
- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
|
||||
- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
|
||||
|
||||
@@ -69,6 +69,8 @@ If the failure is already tracked, leave a comment on your PR noting that the fa
|
||||
|
||||
If the failure appears to be new, investigate whether your changes caused it.
|
||||
|
||||
For resilience behavior under injected process signals/OOM pressure (validation scenarios, not root-cause equivalence), see [Failure Mode Injection Scenarios and Expected Behavior Matrix](../../user_guide/fault_injection_reliability_matrix.md).
|
||||
|
||||
### Reading Buildkite Logs
|
||||
|
||||
1. Click **Details** next to the Buildkite check on your PR.
|
||||
|
||||
@@ -8,6 +8,7 @@ Backend-specific setup lives in separate docs:
|
||||
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
|
||||
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
|
||||
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
|
||||
- [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md)
|
||||
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
|
||||
- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
|
||||
|
||||
@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
|
||||
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
|
||||
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
|
||||
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
|
||||
| Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. |
|
||||
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
|
||||
| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
|
||||
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
# MoriTransferEngineConnector
|
||||
|
||||
## When to Use
|
||||
|
||||
Currently supports intra-node deployment with Mori.
|
||||
|
||||
As noted in #1742, inter-node support will be added back in a future
|
||||
refactor.
|
||||
|
||||
## Mechanism
|
||||
|
||||
Uses Mori's `IOEngine` / `MemoryDesc` API for zero-copy RDMA transfers.
|
||||
|
||||
- Data Plane: RDMA (InfiniBand/RoCE) with managed memory pool.
|
||||
- Control Plane: ZMQ for pull-request handshake and async completion.
|
||||
|
||||
## Installation
|
||||
|
||||
See the [Mori repository](https://github.com/ROCm/mori) for installation instructions.
|
||||
|
||||
## Configuration
|
||||
|
||||
Mori is configured through the new deploy-config schema (see
|
||||
[`docs/configuration/stage_configs.md`](../../../configuration/stage_configs.md)).
|
||||
Define the connector at the top level of the deploy YAML and reference it
|
||||
by name from each stage's `input_connectors` / `output_connectors`:
|
||||
|
||||
```yaml
|
||||
connectors:
|
||||
mori_connector:
|
||||
name: MoriTransferEngineConnector
|
||||
extra:
|
||||
host: "auto"
|
||||
zmq_port: 50051
|
||||
device_name: ""
|
||||
memory_pool_size: 536870912
|
||||
memory_pool_device: "cuda"
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
output_connectors:
|
||||
to_stage_1: mori_connector
|
||||
|
||||
- stage_id: 1
|
||||
input_connectors:
|
||||
from_stage_0: mori_connector
|
||||
```
|
||||
|
||||
A ready-to-run intra-node example for Qwen3-Omni-MoE on AMD MI300X lives
|
||||
at `vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml` and can be loaded
|
||||
with:
|
||||
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
|
||||
--deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
|
||||
```
|
||||
|
||||
The yaml wires `MoriTransferEngineConnector` (with `backend_type: xgmi`) to
|
||||
the chunk_transfer_adapter path (`async_chunk: true`) so stage-to-stage
|
||||
hidden-state and codec-frame streams ship GPU-to-GPU over AMD Infinity
|
||||
Fabric instead of SHM. Qwen2.5-Omni + Mori is not yet functional on
|
||||
the chunk path: it needs `thinker2talker_async_chunk` /
|
||||
`talker2code2wav_async_chunk` input processors that do not exist yet
|
||||
(the orchestrator-level path the upstream PR originally targeted was
|
||||
lost during the entrypoints → engine refactor and removed in #1742).
|
||||
|
||||
Parameters:
|
||||
|
||||
- host: local RDMA IP (`"auto"` for auto-detect).
|
||||
- zmq_port: ZMQ base port for control-plane communication.
|
||||
- device_name: RDMA device (e.g., `"mlx5_0"`), empty for auto-detect.
|
||||
- memory_pool_size: RDMA memory pool size in bytes.
|
||||
- memory_pool_device: `"cpu"` (pinned) or `"cuda"` (GPUDirect / XGMI RDMA).
|
||||
|
||||
For more details, refer to the
|
||||
[Mori repository](https://github.com/ROCm/mori).
|
||||
@@ -33,7 +33,7 @@ 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | | | |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `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` | ✅︎ | ✅︎ | | |
|
||||
@@ -63,6 +63,7 @@ th {
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-VoiceDesign | `Qwen/Qwen3-TTS-12Hz-1.7B-VoiceDesign` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-0.6B-Base` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `MingTTSForConditionalGeneration` | Ming-omni-tts dense 0.5B | `inclusionAI/Ming-omni-tts-0.5B` | ✅︎ | | | |
|
||||
| `GLMTTSForConditionalGeneration` | GLM-TTS | `zai-org/GLM-TTS` | ✅︎ | | | |
|
||||
| `MossTTSNanoForCausalLM` | MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✅︎ | | | |
|
||||
| `MossTTSDelayModel` | MOSS-TTS, MOSS-TTSD, MOSS-SoundEffect, MOSS-VoiceGenerator | `OpenMOSS-Team/MOSS-VoiceGenerator` | ✅︎ | | | |
|
||||
|
||||
@@ -114,6 +114,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. |
|
||||
| `non_streaming_mode` | bool | null | Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |
|
||||
| `stream` | bool | false | Stream raw PCM chunks as they are decoded (requires `response_format="pcm"`) |
|
||||
|
||||
**Supported languages:** Only applicable to Qwen3-TTS. Derived from the model configuration (`talker_config.codec_language_id` in the checkpoint's `config.json`), plus `Auto`, which is always accepted. Official Qwen3-TTS checkpoints support: Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian.
|
||||
@@ -332,7 +333,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
"input": "Hello, this is a cloned voice",
|
||||
"task_type": "Base",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "Original transcript of the reference audio"
|
||||
"ref_text": "Original transcript of the reference audio",
|
||||
"non_streaming_mode": true
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
@@ -454,6 +456,7 @@ Content-Type: application/json
|
||||
| `ref_audio` | string | null | Default reference audio (Base task) |
|
||||
| `ref_text` | string | null | Default reference transcript (Base task) |
|
||||
| `max_new_tokens` | integer | null | Default max tokens |
|
||||
| `non_streaming_mode` | bool | null | Default Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |
|
||||
|
||||
Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
|
||||
|
||||
@@ -469,6 +472,7 @@ Each item in the `items` array requires only `input` (the text). All other field
|
||||
| `ref_audio` | string | Override reference audio |
|
||||
| `ref_text` | string | Override reference transcript |
|
||||
| `max_new_tokens` | integer | Override max tokens |
|
||||
| `non_streaming_mode` | bool | Override Qwen3-TTS prompt construction mode. Does not affect HTTP response streaming or async-chunk pipelining. When null, inherits the batch-level value (then the model default). |
|
||||
|
||||
### Response Format
|
||||
|
||||
|
||||
@@ -70,6 +70,7 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
| `input_reference` | file | null | Uploaded reference image or video for image-to-video/video-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` or `video_reference` |
|
||||
| `video_reference` | string | null | JSON-encoded reference video payload; do not combine with `input_reference` or `image_reference` |
|
||||
| `audio_reference` | string | null | JSON-encoded audio reference for speech-to-video: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `width` | integer | model default | Output video width |
|
||||
| `height` | integer | model default | Output video height |
|
||||
| `num_frames` | integer | model default | Number of generated frames |
|
||||
@@ -161,6 +162,25 @@ are not implemented yet. Models may expose additional V2V controls through
|
||||
`condition_frame_indexes_vision` and `condition_video_keep` to select which
|
||||
decoded reference frames are used as clean conditioning.
|
||||
|
||||
### Speech-to-Video
|
||||
|
||||
For models that support audio-driven generation (e.g., Wan2.2-S2V), pass both
|
||||
an image reference and an audio reference. The `audio_reference` field accepts a
|
||||
JSON string with `audio_url` pointing to an HTTP(s) URL or base64 data URL.
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=A person singing" \
|
||||
-F 'image_reference={"image_url": "https://example.com/face.png"}' \
|
||||
-F 'audio_reference={"audio_url": "https://example.com/speech.mp3"}' \
|
||||
-F "width=832" \
|
||||
-F "height=480" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.5" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
|
||||
### Synchronous Generation
|
||||
|
||||
```bash
|
||||
@@ -189,5 +209,7 @@ walkthroughs, see:
|
||||
|
||||
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
|
||||
- [Speech-to-Video](../user_guide/examples/online_serving/speech_to_video.md)
|
||||
for Wan2.2-S2V audio-driven lip-sync generation
|
||||
- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
|
||||
for model-specific video-to-video examples and conditioning controls
|
||||
@@ -148,7 +148,7 @@ The following tables show which models support each feature:
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
# Speech-To-Video
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/speech_to_video>.
|
||||
|
||||
This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Model ID |
|
||||
|-------|----------|
|
||||
| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` |
|
||||
|
||||
## Start Server
|
||||
|
||||
### Basic Start
|
||||
|
||||
```bash
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size 2 \
|
||||
--flow-shift 3.0 \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--cache-backend cache_dit \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
### Start with Script
|
||||
|
||||
```bash
|
||||
bash run_server.sh
|
||||
```
|
||||
|
||||
The script allows overriding:
|
||||
- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`)
|
||||
- `PORT` (default: `8091`)
|
||||
- `FLOW_SHIFT` (default: `3.0`)
|
||||
- `TP` (default: `2`)
|
||||
- `CACHE_BACKEND` (default: `cache_dit`)
|
||||
|
||||
### Start with 4 GPUs
|
||||
|
||||
```bash
|
||||
TP=4 bash run_server.sh
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
The server provides OpenAI-compatible video generation endpoints:
|
||||
|
||||
- `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 (Recommended for Testing)
|
||||
|
||||
`POST /v1/videos/sync` blocks until generation completes and returns the raw
|
||||
video bytes (`video/mp4`) directly in the response body.
|
||||
|
||||
```bash
|
||||
bash run_curl_speech_to_video.sh
|
||||
```
|
||||
|
||||
## Async Job API
|
||||
|
||||
`POST /v1/videos` creates an asynchronous job and returns immediately.
|
||||
|
||||
```bash
|
||||
no_proxy=127.0.0.1 \
|
||||
create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A person singing" \
|
||||
-F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
|
||||
-F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
|
||||
-F "width=832" -F "height=480" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.5" \
|
||||
-F "fps=16")
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
|
||||
# Poll until complete
|
||||
while true; do
|
||||
status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
# Download
|
||||
curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4
|
||||
```
|
||||
|
||||
## Request Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `prompt` | str | - | Text description of the desired video |
|
||||
| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `width` | int | None | Video width in pixels |
|
||||
| `height` | int | None | Video height in pixels |
|
||||
| `fps` | int | None | Output video frame rate |
|
||||
| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) |
|
||||
| `guidance_scale` | float | None | CFG guidance scale |
|
||||
| `seed` | int | None | Random seed for reproducibility |
|
||||
|
||||
## Notes
|
||||
|
||||
- S2V requires both a reference image (`image_reference`) and an audio reference
|
||||
(`audio_reference`). The generated video will show a person matching the reference
|
||||
image with lip movements synchronized to the audio.
|
||||
- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be:
|
||||
- An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`)
|
||||
- A base64 data URL (e.g., `data:audio/mp3;base64,...`)
|
||||
- `--model-class-name WanS2VPipeline` is required on the server to select the
|
||||
S2V pipeline (distinct from the T2V/I2V pipelines).
|
||||
- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps.
|
||||
- Audio is muxed into the output MP4 automatically.
|
||||
- Always pass `fps=16` to ensure correct video/audio alignment.
|
||||
@@ -0,0 +1,117 @@
|
||||
# vLLM-Omni Failure Mode Injection Scenarios and Expected Behavior Matrix
|
||||
|
||||
## Background and Goals
|
||||
|
||||
The online serving pipeline of vLLM-Omni can be abstracted as: `Serve` (service entry process) -> `Engine` (model orchestration unit) -> `Worker` (execution process) -> `GPU` (compute resource).
|
||||
|
||||
- `Serve` provides external APIs, accepts requests, and manages lifecycle.
|
||||
- `Engine` handles request scheduling and inference orchestration; the number and organization of engines can vary by model type.
|
||||
- `Worker` performs concrete model computation, usually bound to underlying GPU resources.
|
||||
- `GPU` is the resource layer that ultimately carries inference compute and memory usage.
|
||||
|
||||
<div markdown style="text-align: center">
|
||||
|
||||
```mermaid
|
||||
%%{init: {'flowchart': {'useMaxWidth': false}}}%%
|
||||
flowchart TD
|
||||
R(request)--> A(serve)
|
||||
A --> B(engine)
|
||||
A --> C(engine)
|
||||
A --> D(engine)
|
||||
B --> E(worker)
|
||||
C --> F(worker)
|
||||
D --> I(worker)
|
||||
E --> G(GPU)
|
||||
F --> H(GPU)
|
||||
I --> L(GPU)
|
||||
```
|
||||
|
||||
</div>
|
||||
|
||||
The diagram above shows an abstract relationship. Omni models and Diffusion models share the same hierarchy; differences mainly come from engine decomposition, worker count, and scheduling strategy.
|
||||
|
||||
In real customer environments, the system will inevitably encounter accidental operations or environmental disturbances, such as unintended process termination, container/node jitter, or short-term resource contention. From the system perspective, many of these events map to two fault categories: "process receives abnormal signals" or "GPU memory is squeezed." For example, pressing `Ctrl+C` essentially sends a `SIGINT` signal to the `serve` process; running `kill -15 <pid>` maps to graceful termination via `SIGTERM`; running `kill -9 <pid>` maps to forced termination via `SIGKILL`; when `docker kill` is executed or a Pod is deleted, the container runtime usually sends `SIGKILL`.
|
||||
|
||||
This document currently covers two core failure modes, and the system is expected to show predictable and observable behavior under each mode:
|
||||
|
||||
- **Failure Mode 1: Process receives abnormal signals (`SIGINT` / `SIGTERM` / `SIGKILL`)**
|
||||
Focus on three aspects: whether processes exit as expected and complete cleanup, whether requests fail fast or connections are interrupted, and whether GPU resources are eventually fully released with no residue.
|
||||
- **Failure Mode 2: OOM (occupy all free GPU memory via an extra process)**
|
||||
Focus on two aspects: whether service health degrades to the expected state (for example, `503`), and whether requests fail within an acceptable time without hanging.
|
||||
|
||||
More failure modes (for example, network jitter and network interruption) will be added in future iterations to improve end-to-end reliability validation coverage.
|
||||
|
||||
> Note: independent fault injection on the `engine` component is not covered in the current version and will be added gradually in future versions.
|
||||
|
||||
## Fault Injection Scenario Matrix
|
||||
|
||||
| Scenario | Fault Type | System Behavior | Current Status |
|
||||
|------|----------------------------|----------|----------|
|
||||
| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | |
|
||||
| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | |
|
||||
| No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725) <br>[#43060](https://github.com/vllm-project/vllm/issues/43060) |
|
||||
| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | |
|
||||
| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | |
|
||||
| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | |
|
||||
| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | |
|
||||
| Under load | Send `SIGKILL` to Worker | In-flight requests are hard interrupted (5xx/connection drop); main service becomes unavailable | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | |
|
||||
| Under load | Send `SIGKILL` to serve main process | serve is hard-killed and current connections are interrupted; in-flight requests fail; after cleanup there are no residual processes and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGTERM` to serve main process | serve stops accepting new requests and executes shutdown flow; in-flight requests fail; no residue remains and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | `Ctrl+C`-style serve shutdown; in-flight requests fail (5xx/connection interruption); service unavailable; after exit there is no residue and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | |
|
||||
| Under load | Send `SIGTERM` to all related processes | All processes exit gracefully; in-flight requests fail; service unavailable; GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| OOM | Occupy all free GPU memory via an extra process | After OOM injection process starts, GPU memory is continuously saturated; service enters unavailable/degraded state and health check drops to 503; different request types (chat/speech, etc.) fail fast within a fixed time and return 500 (no hanging) | [#4285](https://github.com/vllm-project/vllm-omni/issues/4285) |
|
||||
|
||||
## Source of Conclusions
|
||||
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`.
|
||||
|
||||
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
|
||||
|
||||
The following full log is from a real run where a `SIGTERM` signal was sent to the serve root process.
|
||||
|
||||
```text
|
||||
Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
Server ready on 127.0.0.1:60675 (OmniServer startup took 206.023s)
|
||||
OmniServer started successfully
|
||||
[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[reliability][process-kill] current_server_proc pid=557118 name=python3 cmdline=/workspace/.venv/bin/python3 -c from multiprocessing.resource_tracker import main;main(57)
|
||||
[reliability][process-kill] current_server_proc pid=557119 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
[reliability][process-kill] current_server_proc pid=557122 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
[reliability][process-kill] current_server_proc pid=558201 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
|
||||
[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [async_omni_engine.py:2397] [AsyncOmniEngine] Shutting down Orchestrator
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:351] [Orchestrator] Received shutdown signal
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:1359] [Orchestrator] Shutting down all 3 client(s)
|
||||
|
||||
[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
|
||||
[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1265][0m Shutdown complete
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Shutting down
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Waiting for application shutdown.
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Application shutdown complete.
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Finished server process [556855]
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
|
||||
[rank0]:[W519 12:24:53.427894833 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:56 [stage_pool.py:731] [StagePool] Stage 0 replica 0 shut down
|
||||
[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
|
||||
[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1265][0m Shutdown complete
|
||||
[rank0]:[W519 12:24:57.569425065 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:25:00 [stage_pool.py:731] [StagePool] Stage 1 replica 0 shut down
|
||||
[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
|
||||
[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1265][0m Shutdown complete
|
||||
[0;36m(APIServer pid=556855)[0;0m WARNING 05-19 12:25:01 [async_omni_engine.py:2406] [AsyncOmniEngine] Orchestrator thread did not exit in time
|
||||
[rank0]:[W519 12:25:01.867242016 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
|
||||
```
|
||||
|
||||
### How this maps to the matrix
|
||||
|
||||
- This log corresponds to the `Send SIGTERM to serve main process` scenario in the matrix.
|
||||
- The `root-kill ... signal=SIGTERM` line confirms the fault injection action is applied as expected.
|
||||
- The sequence from `Shutting down Orchestrator` to stage-level `Shutdown complete` and `Stage ... shut down` matches the expected graceful shutdown path and child-process cleanup behavior described in the matrix.
|
||||
|
||||
## Related Documentation
|
||||
|
||||
- For practical CI debugging procedures and triage guidance, see [CI Failures](../contributing/ci/failures.md).
|
||||
@@ -143,6 +143,20 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
required=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=None,
|
||||
metavar="W",
|
||||
help="Output image width in pixels. Default: None (pipeline's default).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=None,
|
||||
metavar="H",
|
||||
help="Output image height in pixels. Default: None (pipeline's default).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
@@ -236,8 +250,8 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--resolution",
|
||||
type=int,
|
||||
default=640,
|
||||
help="Bucket in (640, 1024) to determine the condition and output resolution",
|
||||
default=None,
|
||||
help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
@@ -384,6 +398,15 @@ def parse_args() -> argparse.Namespace:
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
if args.resolution and (args.width or args.height):
|
||||
raise ValueError("--resolution and --width/--height cannot be specified together")
|
||||
if args.width is not None and args.width <= 0:
|
||||
raise ValueError("--width must be a positive integer")
|
||||
if args.height is not None and args.height <= 0:
|
||||
raise ValueError("--height must be a positive integer")
|
||||
if not args.width and not args.height and not args.resolution:
|
||||
args.resolution = 640
|
||||
|
||||
# Validate input images exist and load them
|
||||
input_images = []
|
||||
for image_path in args.image:
|
||||
@@ -464,6 +487,7 @@ def main():
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
|
||||
)
|
||||
print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
@@ -473,22 +497,26 @@ def main():
|
||||
omni.start_profile()
|
||||
|
||||
# Generate edited image
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_2,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_outputs_per_prompt,
|
||||
layers=args.layers,
|
||||
)
|
||||
if args.resolution:
|
||||
sampling_params.resolution = args.resolution
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"multi_modal_data": {"image": input_image},
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_2,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_outputs_per_prompt,
|
||||
layers=args.layers,
|
||||
resolution=args.resolution,
|
||||
),
|
||||
sampling_params,
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -0,0 +1,375 @@
|
||||
# Step-Audio2 Offline Inference Examples
|
||||
|
||||
This directory contains examples for running offline inference with Step-Audio2 using vLLM-Omni.
|
||||
|
||||
## Model Overview
|
||||
|
||||
Step-Audio2 is a two-stage audio model:
|
||||
|
||||
- **Stage 0 (Thinker)**: Audio understanding → Text + Audio tokens
|
||||
- Input: Audio (16kHz)
|
||||
- Output: Text transcription + Audio tokens for synthesis
|
||||
|
||||
- **Stage 1 (Token2Wav)**: Audio synthesis
|
||||
- Input: Audio tokens + Speaker prompt wav
|
||||
- Output: Synthesized audio waveform (24kHz)
|
||||
|
||||
## Hardware Requirements
|
||||
|
||||
| Mode | GPU Configuration | VRAM Required |
|
||||
|------|-------------------|---------------|
|
||||
| ASR (S2T) | 1x GPU | ~20-25GB |
|
||||
| TTS/S2ST (single GPU) | 1x GPU | ~40-50GB |
|
||||
| TTS/S2ST (multi GPU) | 2x GPU | GPU0: ~28GB, GPU1: ~22GB |
|
||||
|
||||
**Tested on:**
|
||||
- 1x NVIDIA H100 80GB (single-card S2ST)
|
||||
- 2x NVIDIA A10 40GB (multi-card S2ST)
|
||||
|
||||
**Notes:**
|
||||
- Single GPU mode requires high VRAM due to both stages sharing memory
|
||||
- Multi GPU mode separates Stage 0 (Thinker) and Stage 1 (Token2Wav) across GPUs
|
||||
- VRAM usage can be adjusted via `gpu_memory_utilization` in stage config
|
||||
|
||||
## Performance Benchmark
|
||||
|
||||
### vLLM-Omni vs Official Step-Audio2
|
||||
|
||||
Single request latency comparison between vLLM-Omni and official Step-Audio2 implementation.
|
||||
|
||||
| Task | Tokens | vllm-omni | Step-Audio2 | Speedup |
|
||||
|------|--------|-----------|-------------|---------|
|
||||
| S2ST | ~85 | 5.45s | 7.36s | **1.35x** |
|
||||
| S2ST | ~160 | 6.67s | 13.92s | **2.09x** |
|
||||
| S2ST | ~315 | 9.42s | 31.50s | **3.34x** |
|
||||
| TTS | ~1024 | 16.65s | ~87s | **~5.2x** |
|
||||
|
||||
**Key observations:**
|
||||
- Speedup increases with sequence length due to vLLM's efficient KV cache management
|
||||
- TTS (pure generation) shows the largest speedup (~5x)
|
||||
- S2ST benefits from optimized multi-stage pipeline
|
||||
|
||||
**Benchmark environment:**
|
||||
- GPU: NVIDIA H100 80GB (single card)
|
||||
- Model: Step-Audio2-mini
|
||||
- Warmup: 1 run, Measured: 3 runs (averaged)
|
||||
|
||||
### Async Chunk Streaming Performance
|
||||
|
||||
Comparison between sequential (non-async) and async chunk modes via `/v1/audio/speech` TTS endpoint.
|
||||
|
||||
| Mode | Mean TTFP | Mean E2E | Mean RTF | Audio Throughput |
|
||||
|------|-----------|----------|----------|-----------------|
|
||||
| Sequential | 4316ms | 4316ms | 0.938 | 1.07x realtime |
|
||||
| **Async Chunk** | **1437ms** | 4362ms | 0.949 | 1.06x realtime |
|
||||
| **Improvement** | **-67% (3x faster)** | ~same | ~same | ~same |
|
||||
|
||||
**Key observations:**
|
||||
- Async chunk reduces **time-to-first-audio (TTFP) by 67%** (4.3s → 1.4s)
|
||||
- E2E latency remains comparable — async chunk overlaps Thinker decode with Token2Wav synthesis
|
||||
- RTF < 1 in both modes (real-time capable)
|
||||
- Sequential mode: TTFP ≈ E2E (must wait for all audio tokens before synthesis starts)
|
||||
- Async chunk mode: Token2Wav starts after first 28 tokens (chunk_size=25 + lookahead=3)
|
||||
|
||||
**Benchmark environment:**
|
||||
- GPU: 4x NVIDIA RTX 3090 24GB (TP=2 for Thinker, 1 GPU for Token2Wav)
|
||||
- Model: Step-Audio2-mini
|
||||
- Endpoint: `/v1/audio/speech` (10 prompts, concurrency=1)
|
||||
- Measured via `bench_tts_serve.py`
|
||||
|
||||
## Installation
|
||||
|
||||
Make sure you have installed vLLM-Omni and all required dependencies:
|
||||
|
||||
```bash
|
||||
# Install vLLM-Omni
|
||||
pip install vllm-omni
|
||||
|
||||
# Install Step-Audio2 (REQUIRED for Token2Wav stage)
|
||||
pip install step-audio2
|
||||
```
|
||||
|
||||
## Model Setup
|
||||
|
||||
### Option 1: Auto-download from HuggingFace (Recommended)
|
||||
|
||||
The script will **automatically download** the model on first run:
|
||||
|
||||
```bash
|
||||
# Just run without specifying --model, it will auto-download stepfun-ai/Step-Audio2-mini
|
||||
python end2end.py --query-type audio_to_text
|
||||
|
||||
# Or explicitly specify the HuggingFace model
|
||||
python end2end.py --query-type audio_to_text --model stepfun-ai/Step-Audio2-mini
|
||||
```
|
||||
|
||||
Models will be cached in `~/.cache/huggingface/hub/` for future use.
|
||||
|
||||
**Available models**:
|
||||
- `stepfun-ai/Step-Audio2-mini` (smaller, faster)
|
||||
- `stepfun-ai/Step-Audio2-7B` (larger, better quality)
|
||||
|
||||
### Option 2: Manual Download (for offline use)
|
||||
|
||||
Download and use locally:
|
||||
|
||||
```bash
|
||||
# Download from HuggingFace
|
||||
huggingface-cli download stepfun-ai/Step-Audio2-mini --local-dir ./models/Step-Audio2-mini
|
||||
|
||||
# Then use the local path
|
||||
python end2end.py --query-type audio_to_text --model ./models/Step-Audio2-mini
|
||||
```
|
||||
|
||||
Ensure the model directory contains:
|
||||
```
|
||||
Step-Audio2-mini/
|
||||
├── config.json
|
||||
├── model.safetensors (or pytorch_model.bin)
|
||||
├── tokenizer.json
|
||||
├── tokenizer_config.json
|
||||
└── token2wav/ # Token2Wav models (REQUIRED)
|
||||
├── speech_tokenizer_v2_25hz.onnx # Audio tokenizer
|
||||
├── campplus.onnx # Speaker encoder
|
||||
├── flow.yaml # Flow model config
|
||||
├── flow.pt # Flow model weights
|
||||
└── hift.pt # HiFT vocoder weights
|
||||
```
|
||||
|
||||
## Usage Examples
|
||||
|
||||
### 1. Audio to Text (ASR - Speech Recognition)
|
||||
|
||||
Transcribe audio to text:
|
||||
|
||||
```bash
|
||||
# Quick start - Using default model and test audio
|
||||
python end2end.py --query-type audio_to_text
|
||||
|
||||
# Using your own audio file (model will auto-download)
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--audio-path /path/to/input.wav
|
||||
|
||||
# With specific model
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--audio-path /path/to/input.wav \
|
||||
--model stepfun-ai/Step-Audio2-7B
|
||||
|
||||
# With custom question
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--audio-path input.wav \
|
||||
--question "What is the speaker saying?"
|
||||
```
|
||||
|
||||
**Output**: Text transcription saved to `output_step_audio2/00000_text.txt`
|
||||
|
||||
### 2. Text to Audio (TTS - Speech Synthesis)
|
||||
|
||||
Convert text to speech:
|
||||
|
||||
```bash
|
||||
# Basic TTS (model auto-downloads)
|
||||
python end2end.py --query-type text_to_audio \
|
||||
--text "Hello, this is a test of Step Audio 2 synthesis."
|
||||
|
||||
# With specific model
|
||||
python end2end.py --query-type text_to_audio \
|
||||
--text "Hello, this is a test." \
|
||||
--model stepfun-ai/Step-Audio2-7B
|
||||
```
|
||||
|
||||
**Note**: Speaker voice is controlled by the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable or the default prompt wav bundled with the model.
|
||||
|
||||
**Output**:
|
||||
- Text: `output_step_audio2/00000_text.txt`
|
||||
- Audio: `output_step_audio2/00000_output.wav` (24kHz)
|
||||
|
||||
### 3. Audio to Audio (Voice Conversion)
|
||||
|
||||
Process input audio and generate output audio:
|
||||
|
||||
```bash
|
||||
# Basic voice conversion (model auto-downloads)
|
||||
python end2end.py --query-type audio_to_audio \
|
||||
--audio-path /path/to/source_audio.wav
|
||||
|
||||
# With specific model
|
||||
python end2end.py --query-type audio_to_audio \
|
||||
--audio-path source.wav \
|
||||
--model stepfun-ai/Step-Audio2-7B
|
||||
```
|
||||
|
||||
This mode:
|
||||
1. Understands the content in `--audio-path` (source)
|
||||
2. Generates audio output with the default voice
|
||||
|
||||
**Note**: To use a custom speaker voice, set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable.
|
||||
|
||||
### Advanced Options
|
||||
|
||||
```bash
|
||||
# Use custom stage configuration
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--stage-configs-path /path/to/custom_config.yaml
|
||||
|
||||
# Multiple prompts (for batch testing)
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--audio-path input.wav \
|
||||
--num-prompts 5
|
||||
|
||||
# Custom output directory
|
||||
python end2end.py --query-type text_to_audio \
|
||||
--text "Test synthesis" \
|
||||
--output-dir ./my_outputs
|
||||
|
||||
# Enable detailed logging
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--audio-path input.wav \
|
||||
--enable-stats
|
||||
|
||||
# Adjust generation parameters
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--audio-path input.wav \
|
||||
--max-tokens 2048
|
||||
|
||||
# Use custom speaker voice via environment variable
|
||||
STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav python end2end.py \
|
||||
--query-type text_to_audio \
|
||||
--text "Hello world"
|
||||
|
||||
# Use Ray backend for distributed processing
|
||||
python end2end.py --query-type text_to_audio \
|
||||
--text "Hello world" \
|
||||
--worker-backend ray \
|
||||
--ray-address "auto"
|
||||
```
|
||||
|
||||
## Configuration
|
||||
|
||||
### Stage Configuration
|
||||
|
||||
The default configuration (`step_audio_2.yaml`) uses:
|
||||
|
||||
- **Stage 0 (Thinker)**: GPU 0, 80% memory
|
||||
- **Stage 1 (Token2Wav)**: GPU 1, 30% memory
|
||||
|
||||
For **single GPU** setup, edit the config to use `devices: "0"` for both stages.
|
||||
|
||||
### Sampling Parameters
|
||||
|
||||
- **Thinker (Stage 0)**:
|
||||
- Temperature: 0.7 (balanced creativity)
|
||||
- Top-p: 0.9
|
||||
- Max tokens: 1024 (configurable)
|
||||
|
||||
- **Token2Wav (Stage 1)**:
|
||||
- Temperature: 0.0 (deterministic)
|
||||
- Operates in generation mode (not sampling)
|
||||
|
||||
## Common Issues
|
||||
|
||||
### 1. ImportError: No module named 's3tokenizer'
|
||||
|
||||
**Solution**: Install Step-Audio2 package:
|
||||
```bash
|
||||
pip install step-audio2
|
||||
```
|
||||
|
||||
### 2. FileNotFoundError: prompt_wav file not found
|
||||
|
||||
**Solution**: Set the `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable to a valid audio file:
|
||||
```bash
|
||||
export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav
|
||||
python end2end.py --query-type text_to_audio --text "Hello"
|
||||
```
|
||||
Or ensure the default prompt wav (`default_female.wav`) exists in your model directory.
|
||||
|
||||
### 3. FileNotFoundError: token2wav models not found
|
||||
|
||||
**Solution**: Ensure your model directory has the complete `token2wav/` subdirectory with all ONNX and PyTorch models.
|
||||
|
||||
### 4. CUDA Out of Memory
|
||||
|
||||
**Solutions**:
|
||||
- Use single GPU mode (set both stages to `devices: "0"`)
|
||||
- Reduce `gpu_memory_utilization` in config
|
||||
- Reduce `max_num_batched_tokens`
|
||||
- Process fewer prompts at once
|
||||
|
||||
### 5. Model not found in registry
|
||||
|
||||
**Solution**: Ensure you're using vLLM-Omni's entry point with `--omni` flag or install vllm-omni properly:
|
||||
```bash
|
||||
pip install vllm-omni
|
||||
```
|
||||
|
||||
## Output Files
|
||||
|
||||
The script generates files in the output directory (default: `output_step_audio2/`):
|
||||
|
||||
```
|
||||
output_step_audio2/
|
||||
├── 00000_text.txt # Text output from Thinker stage
|
||||
├── 00000_output.wav # Audio output from Token2Wav stage (24kHz)
|
||||
├── 00001_text.txt # (if multiple prompts)
|
||||
└── 00001_output.wav
|
||||
```
|
||||
|
||||
## Performance Tips
|
||||
|
||||
1. **First run is slow**: Stage initialization takes 20-60 seconds
|
||||
2. **Single GPU**: Set both stages to `devices: "0"` in config
|
||||
3. **Multiple prompts**: Use `--num-prompts N` for batch testing
|
||||
4. **Ray backend**: For multi-node or advanced scheduling
|
||||
5. **Logging**: Use `--enable-stats` to debug performance issues
|
||||
|
||||
## Speaker Voice Configuration
|
||||
|
||||
The Token2Wav stage requires a speaker prompt wav for voice conditioning. It is automatically resolved in this order:
|
||||
|
||||
1. `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable (if set)
|
||||
2. `{model_dir}/assets/default_female.wav`
|
||||
3. `{model_dir}/default_female.wav`
|
||||
|
||||
If none are found, set the environment variable explicitly:
|
||||
```bash
|
||||
export STEP_AUDIO2_DEFAULT_PROMPT_WAV=/path/to/speaker.wav
|
||||
```
|
||||
|
||||
**Guidelines for custom speaker prompt:**
|
||||
|
||||
- **Duration**: 3-10 seconds recommended
|
||||
- **Quality**: Clean audio, minimal background noise
|
||||
- **Format**: WAV, MP3, FLAC (will be resampled internally)
|
||||
- **Content**: Clear speech, representative of target voice
|
||||
|
||||
## Example Workflow
|
||||
|
||||
Complete example from audio to final output:
|
||||
|
||||
```bash
|
||||
# 1. ASR: Transcribe audio
|
||||
python end2end.py --query-type audio_to_text \
|
||||
--audio-path interview.wav \
|
||||
--model ./models/Step-Audio2-7B \
|
||||
--output-dir ./outputs
|
||||
|
||||
# 2. Check the transcription
|
||||
cat ./outputs/00000_text.txt
|
||||
|
||||
# 3. TTS: Synthesize new speech (with custom voice)
|
||||
STEP_AUDIO2_DEFAULT_PROMPT_WAV=./speaker_samples/female_voice.wav \
|
||||
python end2end.py --query-type text_to_audio \
|
||||
--text "The quick brown fox jumps over the lazy dog" \
|
||||
--model ./models/Step-Audio2-7B \
|
||||
--output-dir ./outputs
|
||||
|
||||
# 4. Listen to the result
|
||||
# Audio saved to: ./outputs/00000_output.wav
|
||||
```
|
||||
|
||||
## References
|
||||
|
||||
- [Step-Audio2 Paper](https://arxiv.org/abs/2507.16632)
|
||||
- [vLLM-Omni Documentation](https://vllm-omni.readthedocs.io)
|
||||
- [Model on HuggingFace](https://huggingface.co/stepfun-ai/Step-Audio-2-mini)
|
||||
@@ -0,0 +1,527 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
This example shows how to use vLLM-Omni for running offline inference
|
||||
with Step-Audio2 model for audio understanding and generation.
|
||||
|
||||
Step-Audio2 is a two-stage model:
|
||||
- Stage 0 (Thinker): Audio understanding → Text + Audio tokens
|
||||
- Stage 1 (Token2Wav): Audio tokens → Waveform (24kHz)
|
||||
|
||||
Usage examples:
|
||||
# Audio to text (ASR)
|
||||
python end2end.py --query-type audio_to_text --audio-path input.wav --model-path "$MODEL_PATH"
|
||||
|
||||
# Text to audio (TTS with default voice)
|
||||
python end2end.py --query-type text_to_audio --model-path "$MODEL_PATH"
|
||||
|
||||
# Audio to audio (Voice conversion with default voice)
|
||||
python end2end.py --query-type audio_to_audio --audio-path input.wav --model-path "$MODEL_PATH"
|
||||
"""
|
||||
|
||||
import os
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import NamedTuple
|
||||
|
||||
import librosa # noqa: TID251
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.sampling_params import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
SEED = 42
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
inputs: dict
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
def get_audio_to_text_query(
|
||||
audio_path: str | None = None,
|
||||
question: str | None = None,
|
||||
sampling_rate: int = 16000,
|
||||
) -> QueryResult:
|
||||
"""
|
||||
Audio Speech Recognition (ASR) - Convert audio to text.
|
||||
|
||||
Args:
|
||||
audio_path: Path to audio file
|
||||
question: Question about the audio (optional)
|
||||
sampling_rate: Target sampling rate (16kHz for Step-Audio2)
|
||||
|
||||
Returns:
|
||||
QueryResult with prompt and audio data
|
||||
"""
|
||||
if question is None:
|
||||
question = "请记录下你所听到的语音内容。"
|
||||
|
||||
prompt = (
|
||||
f"<|im_start|>system\n{question}<|im_end|>\n<|im_start|>user\n<audio_patch><|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
# Load audio
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
|
||||
audio_data = (audio_signal.astype(np.float32), sr)
|
||||
else:
|
||||
# Use default test audio
|
||||
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
# Resample to 16kHz if needed
|
||||
if audio_data[1] != sampling_rate:
|
||||
audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate)
|
||||
audio_data = (audio_signal.astype(np.float32), sampling_rate)
|
||||
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_data,
|
||||
},
|
||||
},
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
|
||||
|
||||
def get_text_to_audio_query(
|
||||
text: str | None = None,
|
||||
sampling_rate: int = 16000,
|
||||
) -> QueryResult:
|
||||
"""
|
||||
Text-to-Speech (TTS) - Convert text to audio.
|
||||
|
||||
Args:
|
||||
text: Text to synthesize
|
||||
sampling_rate: Target sampling rate (16kHz for Step-Audio2)
|
||||
|
||||
Returns:
|
||||
QueryResult with text prompt and speaker wav path
|
||||
|
||||
Note:
|
||||
For audio generation, the prompt must end with "<tts_start>"
|
||||
to signal the model to generate audio tokens instead of just text.
|
||||
"""
|
||||
if text is None:
|
||||
text = "Hello, this is a test of Step Audio 2 text to speech synthesis."
|
||||
|
||||
prompt = (
|
||||
"<|im_start|>system\n"
|
||||
"请逐字朗读用户提供的文本,不要改写或补充。<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
f"{text}<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
"<tts_start>"
|
||||
)
|
||||
|
||||
inputs = {
|
||||
"prompt": prompt,
|
||||
}
|
||||
|
||||
return QueryResult(
|
||||
inputs=inputs,
|
||||
limit_mm_per_prompt={},
|
||||
)
|
||||
|
||||
|
||||
def get_audio_to_audio_query(
|
||||
audio_path: str | None = None,
|
||||
question: str | None = None,
|
||||
sampling_rate: int = 16000,
|
||||
) -> QueryResult:
|
||||
"""
|
||||
Audio-to-Audio - Voice conversion/cloning.
|
||||
|
||||
Args:
|
||||
audio_path: Path to source audio file
|
||||
question: Question/instruction about the audio
|
||||
sampling_rate: Target sampling rate (16kHz for Step-Audio2)
|
||||
|
||||
Returns:
|
||||
QueryResult with audio input
|
||||
|
||||
Note:
|
||||
This mode processes input audio and generates output audio.
|
||||
Speaker voice is controlled by STEP_AUDIO2_DEFAULT_PROMPT_WAV env var.
|
||||
"""
|
||||
if question is None:
|
||||
question = "请仔细聆听这段语音,然后复述其内容。"
|
||||
|
||||
prompt = (
|
||||
"<|im_start|>system\n"
|
||||
f"{question}<|im_end|>\n"
|
||||
"<|im_start|>user\n"
|
||||
"<audio_patch><|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
"<tts_start>"
|
||||
)
|
||||
|
||||
# Load source audio
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
audio_signal, sr = librosa.load(audio_path, sr=sampling_rate)
|
||||
audio_data = (audio_signal.astype(np.float32), sr)
|
||||
|
||||
else:
|
||||
audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
if audio_data[1] != sampling_rate:
|
||||
audio_signal = librosa.resample(audio_data[0], orig_sr=audio_data[1], target_sr=sampling_rate)
|
||||
audio_data = (audio_signal.astype(np.float32), sampling_rate)
|
||||
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_data,
|
||||
},
|
||||
},
|
||||
limit_mm_per_prompt={"audio": 1},
|
||||
)
|
||||
|
||||
|
||||
query_map = {
|
||||
"audio_to_text": get_audio_to_text_query,
|
||||
"text_to_audio": get_text_to_audio_query,
|
||||
"audio_to_audio": get_audio_to_audio_query,
|
||||
}
|
||||
|
||||
|
||||
def main(args):
|
||||
# Use local model path or HuggingFace model name
|
||||
if not args.model:
|
||||
sys.exit(1)
|
||||
|
||||
model_name = args.model
|
||||
|
||||
# Check if it's a local path
|
||||
if os.path.isdir(model_name):
|
||||
# Verify essential files exist
|
||||
config_path = os.path.join(model_name, "config.json")
|
||||
if not os.path.exists(config_path):
|
||||
sys.exit(1)
|
||||
|
||||
# Get query configuration
|
||||
query_func = query_map[args.query_type]
|
||||
|
||||
# Build query based on type
|
||||
if args.query_type == "audio_to_text":
|
||||
query_result = query_func(
|
||||
audio_path=args.audio_path,
|
||||
question=args.question,
|
||||
sampling_rate=args.sampling_rate,
|
||||
)
|
||||
elif args.query_type == "text_to_audio":
|
||||
query_result = query_func(
|
||||
text=args.text,
|
||||
sampling_rate=args.sampling_rate,
|
||||
)
|
||||
elif args.query_type == "audio_to_audio":
|
||||
query_result = query_func(
|
||||
audio_path=args.audio_path,
|
||||
question=args.question,
|
||||
sampling_rate=args.sampling_rate,
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown query type: {args.query_type}")
|
||||
|
||||
# Initialize vLLM-Omni with Step-Audio2
|
||||
# Resolve stage config path
|
||||
if args.stage_configs_path:
|
||||
stage_config_path = args.stage_configs_path
|
||||
else:
|
||||
configs_dir = Path(__file__).parent.parent.parent.parent / "vllm_omni/model_executor/stage_configs"
|
||||
if args.query_type == "audio_to_text":
|
||||
# ASR only needs Thinker (Stage 0), no Token2Wav
|
||||
stage_config_path = str(configs_dir / "step_audio_2_asr.yaml")
|
||||
if not os.path.exists(stage_config_path):
|
||||
stage_config_path = str(configs_dir / "step_audio_2.yaml")
|
||||
else:
|
||||
# TTS/S2ST need both stages
|
||||
stage_config_path = str(configs_dir / "step_audio_2.yaml")
|
||||
|
||||
omni_llm = Omni(
|
||||
model=model_name,
|
||||
log_stats=args.enable_stats,
|
||||
log_file=("step_audio2_pipeline.log" if args.enable_stats else None),
|
||||
init_sleep_seconds=args.init_sleep_seconds,
|
||||
batch_timeout=args.batch_timeout,
|
||||
init_timeout=args.init_timeout,
|
||||
shm_threshold_bytes=args.shm_threshold_bytes,
|
||||
stage_configs_path=stage_config_path,
|
||||
worker_backend=args.worker_backend,
|
||||
ray_address=args.ray_address,
|
||||
)
|
||||
|
||||
# Configure sampling parameters for each stage
|
||||
# Stage 0 (Thinker): Generate text and audio tokens
|
||||
# Note: For audio generation tasks, we need to allow the model to generate audio tokens
|
||||
# Audio tokens are in range [151696, 158257], so we should not stop too early
|
||||
# Qwen EOS token is typically 151645 (<|im_end|>), but we want to allow audio generation
|
||||
# So we don't set stop_token_ids for audio tasks, but we do set it for text-only tasks
|
||||
stop_token_ids = None
|
||||
if args.query_type == "audio_to_text":
|
||||
# For ASR, stop at EOS token
|
||||
stop_token_ids = [151645] # Qwen EOS token <|im_end|>
|
||||
|
||||
# Adjust sampling parameters based on query type
|
||||
# For audio generation tasks, we need higher repetition_penalty to avoid text loops
|
||||
# and allow the model to transition to audio token generation
|
||||
repetition_penalty = 1.1 if args.query_type in ["text_to_audio", "audio_to_audio"] else 1.05
|
||||
|
||||
thinker_sampling_params = SamplingParams(
|
||||
temperature=0.7,
|
||||
top_p=0.9,
|
||||
top_k=-1,
|
||||
max_tokens=args.max_tokens,
|
||||
seed=SEED,
|
||||
detokenize=True,
|
||||
repetition_penalty=repetition_penalty,
|
||||
stop_token_ids=stop_token_ids,
|
||||
)
|
||||
|
||||
# Stage 1 (Token2Wav): Convert audio tokens to waveform
|
||||
# Note: This stage doesn't actually sample, it's deterministic generation
|
||||
token2wav_sampling_params = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=1,
|
||||
seed=SEED,
|
||||
detokenize=False,
|
||||
)
|
||||
|
||||
if args.query_type == "audio_to_text":
|
||||
# ASR: single stage, only Thinker
|
||||
sampling_params_list = [thinker_sampling_params]
|
||||
else:
|
||||
# TTS/S2ST: two stages
|
||||
sampling_params_list = [thinker_sampling_params, token2wav_sampling_params]
|
||||
|
||||
# Prepare prompts
|
||||
if args.num_prompts > 1:
|
||||
prompts = [query_result.inputs for _ in range(args.num_prompts)]
|
||||
else:
|
||||
prompts = [query_result.inputs]
|
||||
|
||||
# Run inference
|
||||
omni_outputs = omni_llm.generate(prompts, sampling_params_list)
|
||||
|
||||
# Create output directory
|
||||
output_dir = args.output_dir
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
|
||||
def _get_request_index(req_id: str, fallback: int) -> int:
|
||||
"""Parse request index from vLLM request_id.
|
||||
|
||||
vLLM may return IDs like "0_<uuid>". Use the leading integer if present,
|
||||
otherwise fall back to the loop index.
|
||||
"""
|
||||
try:
|
||||
return int(req_id)
|
||||
except (TypeError, ValueError):
|
||||
if isinstance(req_id, str):
|
||||
head = req_id.split("_", 1)[0]
|
||||
if head.isdigit():
|
||||
return int(head)
|
||||
return fallback
|
||||
|
||||
# Process outputs from each stage
|
||||
for stage_idx, stage_outputs in enumerate(omni_outputs):
|
||||
# request_output may be a single object or a list
|
||||
raw_output = stage_outputs.request_output
|
||||
outputs_list = raw_output if isinstance(raw_output, list) else [raw_output]
|
||||
|
||||
if stage_outputs.final_output_type == "text":
|
||||
# Stage 0 (Thinker) text output
|
||||
for i, output in enumerate(outputs_list):
|
||||
request_id = _get_request_index(output.request_id, i)
|
||||
text_output = output.outputs[0].text
|
||||
|
||||
# Save to file
|
||||
out_txt = os.path.join(output_dir, f"{request_id:05d}_text.txt")
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
f.write(f"Prompt:\n{prompts[request_id]['prompt']}\n\n")
|
||||
f.write(f"Output:\n{text_output}\n")
|
||||
print(f"Text output saved to: {out_txt}")
|
||||
print(f"Text: {text_output[:200]}")
|
||||
|
||||
elif stage_outputs.final_output_type == "audio":
|
||||
# Stage 1 (Token2Wav) audio output
|
||||
for i, output in enumerate(outputs_list):
|
||||
request_id = _get_request_index(output.request_id, i)
|
||||
|
||||
# Get audio tensor (24kHz)
|
||||
mm_out = {}
|
||||
if hasattr(output, "multimodal_output") and output.multimodal_output:
|
||||
mm_out = output.multimodal_output
|
||||
elif output.outputs and hasattr(output.outputs[0], "multimodal_output"):
|
||||
mm_out = output.outputs[0].multimodal_output or {}
|
||||
audio_tensor = None
|
||||
for key in ("audio", "wav", "waveform", "audio_pcm", "pcm", "model_outputs"):
|
||||
if key in mm_out and mm_out[key] is not None:
|
||||
audio_tensor = mm_out[key]
|
||||
break
|
||||
|
||||
if audio_tensor is not None:
|
||||
# Save audio file
|
||||
output_wav = os.path.join(output_dir, f"{request_id:05d}_output.wav")
|
||||
if isinstance(audio_tensor, list):
|
||||
chunks = []
|
||||
for chunk in audio_tensor:
|
||||
if chunk is None:
|
||||
continue
|
||||
if isinstance(chunk, torch.Tensor):
|
||||
chunk_np = chunk.detach().cpu().numpy()
|
||||
else:
|
||||
chunk_np = np.asarray(chunk)
|
||||
if chunk_np.size > 0:
|
||||
chunks.append(chunk_np)
|
||||
if not chunks:
|
||||
continue
|
||||
audio_numpy = np.concatenate(chunks, axis=-1)
|
||||
elif isinstance(audio_tensor, torch.Tensor):
|
||||
audio_numpy = audio_tensor.detach().cpu().numpy()
|
||||
else:
|
||||
audio_numpy = np.asarray(audio_tensor)
|
||||
sf.write(output_wav, audio_numpy, samplerate=24000)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="Demo on using vLLM-Omni for offline inference with Step-Audio2")
|
||||
|
||||
# Model and config
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
"-m",
|
||||
type=str,
|
||||
required=True,
|
||||
help=(
|
||||
"Model path (required). Can be:\n"
|
||||
" - Local path: /path/to/step-audio-2\n"
|
||||
" - HuggingFace ID: stepfun-ai/Step-Audio2-mini"
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--stage-configs-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to stage config YAML file (default: use built-in step_audio_2.yaml)",
|
||||
)
|
||||
|
||||
# Query configuration
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="audio_to_text",
|
||||
choices=query_map.keys(),
|
||||
help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio (voice conversion)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
"-a",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to input audio file (for audio_to_text and audio_to_audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
"-t",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Text to synthesize (for text_to_audio mode)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--question",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Question/instruction about the audio",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sampling-rate",
|
||||
type=int,
|
||||
default=16000,
|
||||
help="Audio sampling rate for input (default: 16000 Hz)",
|
||||
)
|
||||
|
||||
# Generation parameters
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="Maximum tokens to generate in Thinker stage (default: 1024)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--num-prompts",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of times to repeat the prompt (for testing)",
|
||||
)
|
||||
|
||||
# System configuration
|
||||
parser.add_argument(
|
||||
"--enable-stats",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable detailed statistics logging",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-sleep-seconds",
|
||||
type=int,
|
||||
default=20,
|
||||
help="Sleep seconds after starting each stage (default: 20)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--batch-timeout",
|
||||
type=int,
|
||||
default=5,
|
||||
help="Batch timeout in seconds (default: 5)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--init-timeout",
|
||||
type=int,
|
||||
default=300,
|
||||
help="Initialization timeout in seconds (default: 300)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--shm-threshold-bytes",
|
||||
type=int,
|
||||
default=65536,
|
||||
help="Shared memory threshold in bytes (default: 65536)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--worker-backend",
|
||||
type=str,
|
||||
default="multi_process",
|
||||
choices=["multi_process", "ray"],
|
||||
help="Worker backend (default: multi_process)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ray-address",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Ray cluster address (for ray backend)",
|
||||
)
|
||||
|
||||
# Output
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
"-o",
|
||||
type=str,
|
||||
default="output_step_audio2",
|
||||
help="Output directory for generated files (default: output_step_audio2)",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
main(args)
|
||||
@@ -17,6 +17,7 @@ list of supported architectures across all modalities, see
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
@@ -159,6 +160,46 @@ Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
---
|
||||
|
||||
## Ming-omni-tts
|
||||
|
||||
Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case style \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case zero_shot \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Streaming
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case basic \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--streaming \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `style`, `ip`, `bgm`, and `tta` do not require reference audio.
|
||||
- Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`.
|
||||
- `podcast` uses multiple references via `--ref-audio-paths`.
|
||||
- Full case details live in [`ming_tts/README.md`](ming_tts/README.md).
|
||||
|
||||
---
|
||||
|
||||
## Ming-flash-omni-TTS
|
||||
|
||||
Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio.
|
||||
|
||||
@@ -0,0 +1,248 @@
|
||||
# Ming-omni-tts Offline Inference
|
||||
|
||||
`end2end.py` runs Ming dense 0.5B end to end with vLLM-Omni. It uses the in-repo Ming prompt assembly helper directly, so the example request shape matches the real integration instead of a simplified wrapper.
|
||||
|
||||
## Files
|
||||
|
||||
| File | Purpose |
|
||||
|---|---|
|
||||
| `end2end.py` | Driver: CLI, case loading, prompt construction, orchestration (~150 lines) |
|
||||
| `cases.yaml` | All 11 built-in case definitions (prompt, text, instruction, ref-audio flags, flow controls) |
|
||||
| `runner.py` | Engine management and audio output (streaming + blocking paths) |
|
||||
|
||||
## Model Overview
|
||||
|
||||
Ming dense 0.5B is exposed here as a two-stage offline pipeline:
|
||||
|
||||
- **Stage 0**: Qwen2-based AR generation with Ming prompt formatting and inline flow controls
|
||||
- **Stage 1**: audio VAE decode to mono 44.1 kHz waveform
|
||||
|
||||
`config_ming_tts.py` adapts the checkpoint's HuggingFace config fields
|
||||
(LLM, DiT, aggregator, AudioVAE, token ids). `vllm_omni/deploy/ming_tts.yaml`
|
||||
selects the vLLM-Omni pipeline and stage runtime topology, including
|
||||
connectors, async chunking, memory limits, and sampling defaults.
|
||||
|
||||
The example supports both:
|
||||
|
||||
- **Blocking eager** via `vllm_omni/deploy/ming_tts.yaml`
|
||||
- **Async chunk eager** via `vllm_omni/deploy/ming_tts.yaml` (default `async_chunk: true`)
|
||||
|
||||
## Setup
|
||||
|
||||
Install vLLM-Omni with the platform requirements for your accelerator:
|
||||
|
||||
```bash
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
The Ming offline example does not require a separate upstream Ming package.
|
||||
Reference-audio cases use the repo dependencies for audio loading,
|
||||
resampling, and CampPlus speaker extraction, including `soundfile`,
|
||||
`torchaudio`, and `onnxruntime-rocm` in the documented ROCm environment.
|
||||
|
||||
The tested ROCm environment is summarized in the
|
||||
[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
|
||||
|
||||
## Supported Cases
|
||||
|
||||
These cases cover the upstream dense 0.5B cookbook surface that maps cleanly onto the current vLLM-Omni example:
|
||||
|
||||
- `style`: zero-speaker style-conditioned speech
|
||||
- `ip`: zero-speaker IP voice generation
|
||||
- `bgm`: music-only generation
|
||||
- `tta`: text-to-audio event generation with FlowLoss controls
|
||||
- `emotion`: reference-audio speech with emotion control
|
||||
- `basic`: reference-audio speech with speed / pitch / volume control
|
||||
- `dialect`: reference-audio speech with dialect control
|
||||
- `zero_shot`: reference-audio cloning with explicit transcript
|
||||
- `podcast`: multi-reference dialogue generation with automatic speaker embedding extraction
|
||||
- `speech_bgm`: speech with background music conditioning
|
||||
- `speech_sound`: speech with environmental sound conditioning
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the zero-speaker style example:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case style \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Run zero-shot cloning with a transcript:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case zero_shot \
|
||||
--ref-audio /path/to/10002287-00000094.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Run emotion-controlled speech:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case emotion \
|
||||
--ref-audio /path/to/emotion_prompt.wav \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Run podcast generation with two reference clips:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case podcast \
|
||||
--ref-audio-paths /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
The script automatically extracts one 192-d speaker embedding per reference WAV using the Ming model's `campplus.onnx`.
|
||||
|
||||
If you already have precomputed multi-speaker embeddings, you can override extraction with:
|
||||
|
||||
```bash
|
||||
--speaker-embedding /path/to/podcast_speaker_embeddings.json
|
||||
```
|
||||
|
||||
where the JSON is a list of speaker embeddings, one 192-d vector per speaker.
|
||||
|
||||
Run text-to-audio event generation:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case tta \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Use async_chunk streaming:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case basic \
|
||||
--ref-audio /path/to/10002287-00000095.wav \
|
||||
--streaming \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently
|
||||
supports one prompt per process invocation; use blocking mode for
|
||||
`--num-prompts > 1`.
|
||||
|
||||
Collect runtime stats and a manifest:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case style \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager \
|
||||
--enable-stats \
|
||||
--stats-log-file output_audio/ming_style_pipeline.log \
|
||||
--metadata-json output_audio/ming_style_manifest.json
|
||||
```
|
||||
|
||||
## Reference Fixtures
|
||||
|
||||
The upstream Ming cookbook uses these public audio fixtures from `inclusionAI/Ming-omni-tts/data/wavs`:
|
||||
|
||||
- `10002287-00000094.wav` for zero-shot cloning
|
||||
- `10002287-00000095.wav` for `basic`
|
||||
- `emotion_prompt.wav` for `emotion`
|
||||
- `yue_prompt.wav` for `dialect`
|
||||
- `00000309-00000300.wav` for `speech_bgm` and `speech_sound`
|
||||
- `CTS-CN-F2F-2019-11-11-423-012-A.wav` and `CTS-CN-F2F-2019-11-11-423-012-B.wav` for `podcast`
|
||||
|
||||
## Validation Matrix
|
||||
|
||||
The repo-facing example is intended to cover the same dense TTS workflows used
|
||||
by the local Ming validation script:
|
||||
|
||||
| Case | Blocking `deploy/ming_tts.yaml` | Async chunk `deploy/ming_tts.yaml` | Extra inputs |
|
||||
|---|---:|---:|---|
|
||||
| `style` | Yes | Optional smoke test | none |
|
||||
| `ip` | Yes | Optional smoke test | none |
|
||||
| `bgm` | Yes | Optional smoke test | none |
|
||||
| `tta` | Yes | Optional smoke test | none |
|
||||
| `emotion` | Yes | Yes | `--ref-audio emotion_prompt.wav` |
|
||||
| `basic` | Yes | Yes | `--ref-audio 10002287-00000095.wav` |
|
||||
| `dialect` | Yes | Yes | `--ref-audio yue_prompt.wav` |
|
||||
| `zero_shot` | Yes | Yes | `--ref-audio 10002287-00000094.wav --ref-text ...` |
|
||||
| `podcast` | Yes | Yes | two `--ref-audio-paths` |
|
||||
| `speech_bgm` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
|
||||
| `speech_sound` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
|
||||
|
||||
## Validated Outputs
|
||||
|
||||
The following measurements are retained from an earlier L4 CUDA validation;
|
||||
they are not ROCm benchmark results. Default async_chunk matched blocking
|
||||
output frame counts and Stage-1 patch counts for every case:
|
||||
|
||||
| Case | Blocking frames / patches / sec | Async chunk frames / patches / sec |
|
||||
|---|---:|---:|
|
||||
| `style` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
|
||||
| `ip` | 183456 / 13 / 4.16 | 183456 / 13 / 4.16 |
|
||||
| `bgm` | 1326528 / 94 / 30.08 | 1326528 / 94 / 30.08 |
|
||||
| `tta` | 465696 / 33 / 10.56 | 465696 / 33 / 10.56 |
|
||||
| `emotion` | 324576 / 23 / 7.36 | 324576 / 23 / 7.36 |
|
||||
| `basic` | 211680 / 15 / 4.80 | 211680 / 15 / 4.80 |
|
||||
| `dialect` | 239904 / 17 / 5.44 | 239904 / 17 / 5.44 |
|
||||
| `zero_shot` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
|
||||
| `podcast` | 437472 / 31 / 9.92 | 437472 / 31 / 9.92 |
|
||||
| `speech_bgm` | 296352 / 21 / 6.72 | 296352 / 21 / 6.72 |
|
||||
| `speech_sound` | 352800 / 25 / 8.00 | 352800 / 25 / 8.00 |
|
||||
|
||||
## Key Arguments
|
||||
|
||||
| Argument | Description |
|
||||
|---|---|
|
||||
| `--model` | Hugging Face repo or local Ming checkpoint path |
|
||||
| `--deploy-config` | Deploy config YAML. Use `vllm_omni/deploy/ming_tts.yaml` |
|
||||
| `--case` | Built-in demo case |
|
||||
| `--ref-audio` | Single reference wav path for cloning-style cases |
|
||||
| `--ref-audio-paths` | Multiple reference wav paths, used by `podcast` |
|
||||
| `--ref-text` | Reference transcript. Required for `zero_shot` |
|
||||
| `--instructions` | Free-form Ming instruction string |
|
||||
| `--instruction-json` | Structured Ming instruction JSON |
|
||||
| `--speaker-embedding` | JSON file containing a 192-d speaker embedding |
|
||||
| `--extract-speaker-embeddings` | Force CampPlus speaker extraction from the provided reference audio paths |
|
||||
| `--max-decode-steps` | Override `ming_max_decode_steps` |
|
||||
| `--num-prompts` | Repeat the same case N times. Outputs are indexed when `N > 1` |
|
||||
| `--streaming` | Use `AsyncOmni` and async_chunk transport |
|
||||
| `--enforce-eager` | Recommended for Ming dense; non-eager is out of scope |
|
||||
| `--enable-stats` / `--log-stats` | Enable vLLM-Omni per-request stats logging |
|
||||
| `--stats-log-file` | Optional path for the stats log |
|
||||
| `--metadata-json` | Optional path for the run manifest JSON |
|
||||
| `--stage-init-timeout` | Per-stage initialization timeout in seconds |
|
||||
| `--init-timeout` | Total initialization timeout in seconds |
|
||||
| `--batch-timeout` | Batch timeout in seconds |
|
||||
| `--worker-backend` | `multi_process` or `ray` |
|
||||
| `--ray-address` | Ray cluster address when using `--worker-backend ray` |
|
||||
|
||||
## Output
|
||||
|
||||
- The script writes one mono 44.1 kHz WAV file per run
|
||||
- Default output directory: `output_audio/`
|
||||
- Default filename: `ming_<case>.wav`
|
||||
- When `--num-prompts > 1`, outputs are indexed as `ming_<case>_00000.wav`, `..._00001.wav`, etc.
|
||||
- When stats are enabled, the script can also write:
|
||||
- a stats log file such as `ming_style_pipeline.log`
|
||||
- a manifest JSON with per-output metadata, stage durations, peak memory info,
|
||||
and streaming client latency metrics when `--streaming` is used
|
||||
|
||||
## Notes
|
||||
|
||||
- `style` and `ip` are zero-speaker paths and do not require a reference clip
|
||||
- `emotion`, `basic`, `dialect`, `speech_bgm`, and `speech_sound` require one reference clip
|
||||
- `zero_shot` requires both `--ref-audio` and `--ref-text`
|
||||
- `podcast` requires at least two reference clips via `--ref-audio-paths`
|
||||
- `podcast` automatically extracts one speaker embedding per reference clip
|
||||
- `--speaker-embedding` may contain either one 192-d vector or a list of 192-d vectors
|
||||
- `--enforce-eager` was used for the validated runs
|
||||
- The earlier L4 validation used SDPA for the Ming audio VAE instead of
|
||||
FlashAttention2, which is the preferred default when available.
|
||||
@@ -0,0 +1,110 @@
|
||||
style:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?"
|
||||
instruction:
|
||||
风格: >-
|
||||
这是一种ASMR耳语,属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语,声音气音成分重。音量极低,紧贴麦克风,语速极慢,旨在制造触发听者颅内快感的声学刺激。
|
||||
use_zero_spk_emb: true
|
||||
max_decode_steps: 200
|
||||
|
||||
ip:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "这款产品的名字,叫变态坑爹牛肉丸。"
|
||||
instruction:
|
||||
IP: "灵小甄"
|
||||
use_zero_spk_emb: true
|
||||
max_decode_steps: 200
|
||||
|
||||
bgm:
|
||||
prompt: "Please generate music based on the following description.\n"
|
||||
text: " Genre: 电子舞曲. Mood: 自信 / 坚定. Instrument: 架子鼓. Theme: 节日. Duration: 30s."
|
||||
instruction: null
|
||||
use_zero_spk_emb: false
|
||||
max_decode_steps: 400
|
||||
|
||||
tta:
|
||||
prompt: "Please generate audio events based on given text.\n"
|
||||
text: "Thunder and a gentle rain"
|
||||
instruction: null
|
||||
use_zero_spk_emb: false
|
||||
max_decode_steps: 200
|
||||
cfg: 4.5
|
||||
sigma: 0.3
|
||||
temperature: 2.5
|
||||
|
||||
emotion:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我竟然抢到了陈奕迅的演唱会门票!太棒了!终于可以现场听一听他的歌声了!"
|
||||
instruction:
|
||||
情感: "高兴"
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
basic:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "简单地说,这相当于惠普把消费领域市场拱手相让了。"
|
||||
instruction:
|
||||
语速: "快速"
|
||||
基频: "中"
|
||||
音量: "中"
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
dialect:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我觉得社会企业同个人都有责任"
|
||||
instruction:
|
||||
方言: "广粤话"
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
zero_shot:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。"
|
||||
instruction: null
|
||||
requires_ref_audio: true
|
||||
requires_ref_text: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
podcast:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: " speaker_1:你可以说一下,就大概说一下,可能虽然我也不知道,我看过那部电影没有。\n speaker_2:就是那个叫什么,变相一节课的嘛。\n speaker_1:嗯。\n speaker_2:一部搞笑的电影。\n speaker_1:一部搞笑的。\n"
|
||||
instruction: null
|
||||
prompt_text: " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试,做个,当服务员还要去笔试了\n speaker_2:对啊,这真的很奇怪,就是 单纯的因,单纯自己工资不高,只是因为可能人家那个店比较出名一点,就对你苛刻要求\n"
|
||||
requires_ref_audio_count: 2
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
speech_bgm:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
|
||||
instruction:
|
||||
BGM:
|
||||
Genre: "当代古典音乐."
|
||||
Mood: "温暖 / 友善."
|
||||
Instrument: "电吉他"
|
||||
Theme: "节日."
|
||||
SNR: 10.0
|
||||
ENV: null
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
speech_sound:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
|
||||
instruction:
|
||||
BGM:
|
||||
ENV: "Birds chirping"
|
||||
SNR: 10.0
|
||||
Genre: null
|
||||
Mood: null
|
||||
Instrument: null
|
||||
Theme: null
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
@@ -0,0 +1,200 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference demo for Ming-omni-tts via vLLM Omni."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
import torchaudio
|
||||
import yaml
|
||||
from transformers import AutoTokenizer
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
KEY_CFG,
|
||||
KEY_MAX_DECODE_STEPS,
|
||||
KEY_SIGMA,
|
||||
KEY_TEMPERATURE,
|
||||
SAMPLE_RATE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
|
||||
from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
|
||||
|
||||
try:
|
||||
from .runner import (
|
||||
build_manifest,
|
||||
build_sampling_params,
|
||||
resolve_metadata_json,
|
||||
resolve_stats_log_file,
|
||||
run_generation,
|
||||
)
|
||||
except ImportError:
|
||||
from runner import (
|
||||
build_manifest,
|
||||
build_sampling_params,
|
||||
resolve_metadata_json,
|
||||
resolve_stats_log_file,
|
||||
run_generation,
|
||||
)
|
||||
|
||||
_DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni/deploy/ming_tts.yaml")
|
||||
_CASES_FILE = Path(__file__).with_name("cases.yaml")
|
||||
|
||||
CASE_DEFAULTS = yaml.safe_load(_CASES_FILE.read_text(encoding="utf-8")) or {}
|
||||
if not CASE_DEFAULTS:
|
||||
raise RuntimeError(f"Empty or missing case definitions in {_CASES_FILE}")
|
||||
|
||||
|
||||
def _build_parser() -> FlexibleArgumentParser:
|
||||
p = FlexibleArgumentParser(description="Offline Ming-omni-tts example")
|
||||
p.add_argument("--model", default=_DEFAULT_MODEL, help="Model name or local path")
|
||||
p.add_argument("--deploy-config", default=None, help="Deploy config YAML; auto-selected when omitted")
|
||||
p.add_argument("--case", choices=sorted(CASE_DEFAULTS), default="style", help="Built-in demo case")
|
||||
p.add_argument("--text", default=None, help="Override case text")
|
||||
p.add_argument("--prompt", default=None, help="Override the system prompt prefix")
|
||||
p.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
|
||||
p.add_argument(
|
||||
"--instruction-json", default=None, help='Structured Ming instruction JSON, e.g. \'{"方言":"广粤话"}\''
|
||||
)
|
||||
p.add_argument("--ref-audio", default=None, help="Single reference audio path for cloning-style cases")
|
||||
p.add_argument("--ref-audio-paths", nargs="+", default=None, help="Multiple reference audio paths (podcast)")
|
||||
p.add_argument("--ref-text", default=None, help="Reference transcript; required for zero_shot")
|
||||
p.add_argument("--speaker-embedding", default=None, help="Path to a JSON speaker embedding file")
|
||||
p.add_argument(
|
||||
"--extract-speaker-embeddings",
|
||||
action="store_true",
|
||||
help="Extract speaker embeddings from ref audio via campplus.onnx",
|
||||
)
|
||||
p.add_argument("--max-decode-steps", type=int, default=None, help="Override ming_max_decode_steps")
|
||||
p.add_argument("--output-dir", default="output_audio", help="Directory for output wav files")
|
||||
p.add_argument("--output-name", default=None, help="Output wav filename")
|
||||
p.add_argument("--num-prompts", type=int, default=1, help="Repeat the same prompt N times")
|
||||
p.add_argument("--streaming", action="store_true", help="Use AsyncOmni with async_chunk streaming")
|
||||
p.add_argument("--trust-remote-code", action="store_true")
|
||||
p.add_argument("--enforce-eager", action="store_true")
|
||||
p.add_argument(
|
||||
"--log-stats", "--enable-stats", dest="log_stats", action="store_true", help="Enable Omni stats logging"
|
||||
)
|
||||
p.add_argument("--stats-log-file", default=None, help="Path for the Omni stats log file")
|
||||
p.add_argument("--metadata-json", default=None, help="Path for the run manifest JSON")
|
||||
p.add_argument("--stage-init-timeout", type=int, default=300, help="Per-stage init timeout (s)")
|
||||
p.add_argument("--init-timeout", type=int, default=600, help="Total init timeout (s)")
|
||||
p.add_argument("--batch-timeout", type=int, default=5, help="Batch timeout (s)")
|
||||
p.add_argument("--shm-threshold-bytes", type=int, default=65536)
|
||||
p.add_argument("--worker-backend", default="multi_process", choices=["multi_process", "ray"])
|
||||
p.add_argument("--ray-address", default=None, help="Ray cluster address (--worker-backend ray)")
|
||||
return p
|
||||
|
||||
|
||||
def _finalize_args(args) -> None:
|
||||
if args.instructions is not None and args.instruction_json is not None:
|
||||
raise RuntimeError("Use either --instructions or --instruction-json, not both")
|
||||
if args.num_prompts < 1:
|
||||
raise RuntimeError("--num-prompts must be at least 1")
|
||||
if args.streaming and args.num_prompts != 1:
|
||||
raise RuntimeError("--streaming currently supports exactly one prompt")
|
||||
if args.deploy_config is None:
|
||||
args.deploy_config = _DEFAULT_DEPLOY_CONFIG
|
||||
|
||||
|
||||
def _load_waveform(path: str) -> torch.Tensor:
|
||||
samples, sr = sf.read(path, dtype="float32")
|
||||
wav = torch.as_tensor(samples, dtype=torch.float32)
|
||||
if wav.ndim == 2:
|
||||
wav = wav.mean(dim=1)
|
||||
wav = wav.reshape(1, -1)
|
||||
if int(sr) != SAMPLE_RATE:
|
||||
wav = torchaudio.functional.resample(wav, int(sr), SAMPLE_RATE)
|
||||
return wav
|
||||
|
||||
|
||||
def _load_speaker_embedding(path: str) -> torch.Tensor:
|
||||
return torch.as_tensor(json.loads(Path(path).read_text(encoding="utf-8")), dtype=torch.float32)
|
||||
|
||||
|
||||
def _ref_audio_paths(args) -> list[str]:
|
||||
if args.ref_audio is not None and args.ref_audio_paths is not None:
|
||||
raise RuntimeError("Use either --ref-audio or --ref-audio-paths, not both")
|
||||
if args.ref_audio_paths is not None:
|
||||
return list(args.ref_audio_paths)
|
||||
return [args.ref_audio] if args.ref_audio else []
|
||||
|
||||
|
||||
def _resolve_reference_inputs(args, case: dict, paths: list[str]):
|
||||
required = int(case.get("requires_ref_audio_count", 0))
|
||||
if required > 0 and len(paths) < required:
|
||||
raise RuntimeError(f"Case '{args.case}' needs {required} ref audio paths via --ref-audio-paths")
|
||||
if required <= 0 and case.get("requires_ref_audio") and not paths:
|
||||
raise RuntimeError(f"--ref-audio required for case '{args.case}'")
|
||||
if not paths:
|
||||
return None
|
||||
wavs = [_load_waveform(p) for p in paths]
|
||||
return wavs[0] if len(wavs) == 1 else wavs
|
||||
|
||||
|
||||
def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
|
||||
if args.speaker_embedding:
|
||||
return _load_speaker_embedding(args.speaker_embedding)
|
||||
if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
|
||||
return None
|
||||
embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
|
||||
if not embs:
|
||||
raise RuntimeError("Speaker extraction produced no embeddings")
|
||||
return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)
|
||||
|
||||
|
||||
def _build_prompt_payload(tokenizer, args):
|
||||
case = CASE_DEFAULTS[args.case]
|
||||
paths = _ref_audio_paths(args)
|
||||
prompt_text = args.ref_text if args.ref_text is not None else case.get("prompt_text")
|
||||
if case.get("requires_ref_text") and not prompt_text:
|
||||
raise RuntimeError(f"--ref-text required for case '{args.case}'")
|
||||
prompt_waveform = _resolve_reference_inputs(args, case, paths) if prompt_text is not None else None
|
||||
speaker_embedding = _resolve_speaker_embedding(args, case, paths)
|
||||
use_zero_spk_emb = bool(case.get("use_zero_spk_emb")) and prompt_waveform is None and speaker_embedding is None
|
||||
runtime_controls = {KEY_MAX_DECODE_STEPS: args.max_decode_steps or case["max_decode_steps"]}
|
||||
for key, field in [(KEY_CFG, "cfg"), (KEY_SIGMA, "sigma"), (KEY_TEMPERATURE, "temperature")]:
|
||||
if field in case:
|
||||
runtime_controls[key] = case[field]
|
||||
instruction = (
|
||||
json.loads(args.instruction_json) if args.instruction_json else (args.instructions or case.get("instruction"))
|
||||
)
|
||||
return build_ming_dense_prompt(
|
||||
tokenizer,
|
||||
prompt=args.prompt or case["prompt"],
|
||||
text=args.text or case["text"],
|
||||
runtime_controls=runtime_controls,
|
||||
instruction=instruction,
|
||||
prompt_text=prompt_text,
|
||||
prompt_waveform=prompt_waveform,
|
||||
speaker_embedding=speaker_embedding,
|
||||
use_zero_spk_emb=use_zero_spk_emb,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
args = _build_parser().parse_args()
|
||||
_finalize_args(args)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False)
|
||||
prompt_payload = _build_prompt_payload(tokenizer, args)
|
||||
case = CASE_DEFAULTS[args.case]
|
||||
sampling_params_list = build_sampling_params(args.max_decode_steps or case["max_decode_steps"])
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stats_log_file = resolve_stats_log_file(args)
|
||||
summaries = run_generation(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
|
||||
metadata_json = resolve_metadata_json(args)
|
||||
manifest = build_manifest(args, prompt_payload, stats_log_file, summaries)
|
||||
if metadata_json is not None:
|
||||
Path(metadata_json).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(f"Saved run manifest to {metadata_json}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
main()
|
||||
@@ -0,0 +1,277 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
KEY_SPEAKER_EMBEDDING,
|
||||
SAMPLE_RATE,
|
||||
TEXT_EOS_TOKEN_ID,
|
||||
)
|
||||
|
||||
|
||||
def coerce_audio_tensor(audio, *, async_chunk: bool) -> torch.Tensor:
|
||||
if isinstance(audio, list):
|
||||
parts = []
|
||||
for item in audio:
|
||||
tensor = torch.as_tensor(item, dtype=torch.float32).reshape(-1)
|
||||
if tensor.numel() > 0:
|
||||
parts.append(tensor)
|
||||
if not parts:
|
||||
return torch.zeros((0,), dtype=torch.float32)
|
||||
return torch.cat(parts, dim=0)
|
||||
|
||||
return torch.as_tensor(audio, dtype=torch.float32).reshape(-1)
|
||||
|
||||
|
||||
def resolve_sr(sr) -> int:
|
||||
if isinstance(sr, list):
|
||||
sr = sr[-1]
|
||||
if hasattr(sr, "item"):
|
||||
return int(sr.item())
|
||||
return int(sr)
|
||||
|
||||
|
||||
def extract_sample_rate(multimodal_output: dict) -> int:
|
||||
sr = multimodal_output.get("sr")
|
||||
if sr is None:
|
||||
raise RuntimeError("Expected multimodal_output['sr']")
|
||||
return resolve_sr(sr)
|
||||
|
||||
|
||||
def write_wav(path: str, audio: torch.Tensor, sample_rate: int) -> None:
|
||||
audio = audio.clamp(-1.0, 1.0)
|
||||
pcm16 = (audio * 32767.0).round().to(torch.int16).cpu().numpy()
|
||||
with wave.open(path, "wb") as wav_file:
|
||||
wav_file.setnchannels(1)
|
||||
wav_file.setsampwidth(2)
|
||||
wav_file.setframerate(int(sample_rate))
|
||||
wav_file.writeframes(pcm16.tobytes())
|
||||
|
||||
|
||||
def request_index(request_id: str | None, fallback: int) -> int:
|
||||
try:
|
||||
return int(request_id)
|
||||
except (TypeError, ValueError):
|
||||
if isinstance(request_id, str):
|
||||
head = request_id.split("_", 1)[0]
|
||||
if head.isdigit():
|
||||
return int(head)
|
||||
return fallback
|
||||
|
||||
|
||||
def audio_summary(audio: torch.Tensor, sample_rate: int) -> dict:
|
||||
waveform = audio.detach().cpu().reshape(-1).to(torch.float32)
|
||||
return {
|
||||
"sample_rate": int(sample_rate),
|
||||
"num_samples": int(waveform.numel()),
|
||||
"duration_seconds": float(waveform.numel()) / float(sample_rate),
|
||||
"max_abs_amplitude": float(waveform.abs().max().item()) if waveform.numel() > 0 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def resolve_output_name(output_name: str | None, case: str, index: int, total: int) -> str:
|
||||
if total == 1:
|
||||
return output_name or f"ming_{case}.wav"
|
||||
base = Path(output_name or f"ming_{case}.wav")
|
||||
return f"{base.stem}_{index:05d}{base.suffix or '.wav'}"
|
||||
|
||||
|
||||
def resolve_stats_log_file(args) -> str | None:
|
||||
if not args.log_stats:
|
||||
return None
|
||||
if args.stats_log_file:
|
||||
return args.stats_log_file
|
||||
base = Path(args.output_name or f"ming_{args.case}.wav").stem
|
||||
return str(Path(args.output_dir) / f"{base}_pipeline.log")
|
||||
|
||||
|
||||
def resolve_metadata_json(args) -> str | None:
|
||||
if args.metadata_json:
|
||||
return args.metadata_json
|
||||
if args.log_stats:
|
||||
base = Path(args.output_name or f"ming_{args.case}.wav").stem
|
||||
return str(Path(args.output_dir) / f"{base}_manifest.json")
|
||||
return None
|
||||
|
||||
|
||||
def build_manifest(args, prompt_payload, stats_log_file: str | None, outputs: list[dict]) -> dict:
|
||||
additional_information = {}
|
||||
if isinstance(prompt_payload, dict):
|
||||
additional_information = dict(prompt_payload.get("additional_information", {}))
|
||||
return {
|
||||
"model": args.model,
|
||||
"case": args.case,
|
||||
"streaming": bool(args.streaming),
|
||||
"deploy_config": args.deploy_config,
|
||||
"enforce_eager": bool(args.enforce_eager),
|
||||
"num_prompts": int(args.num_prompts),
|
||||
"log_stats": bool(args.log_stats),
|
||||
"stats_log_file": stats_log_file,
|
||||
"prompt_text": additional_information.get("prompt_text"),
|
||||
"instruction": additional_information.get("instruction"),
|
||||
"speaker_embedding_shape": (
|
||||
list(additional_information[KEY_SPEAKER_EMBEDDING].shape)
|
||||
if KEY_SPEAKER_EMBEDDING in additional_information
|
||||
and hasattr(additional_information[KEY_SPEAKER_EMBEDDING], "shape")
|
||||
else None
|
||||
),
|
||||
"outputs": outputs,
|
||||
"generated_at_unix": time.time(),
|
||||
}
|
||||
|
||||
|
||||
def build_engine_kwargs(args, stats_log_file: str | None) -> dict:
|
||||
kwargs = {
|
||||
"model": args.model,
|
||||
"deploy_config": args.deploy_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"trust_remote_code": args.trust_remote_code,
|
||||
"log_stats": args.log_stats,
|
||||
"stage_init_timeout": args.stage_init_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
"batch_timeout": args.batch_timeout,
|
||||
"shm_threshold_bytes": args.shm_threshold_bytes,
|
||||
"worker_backend": args.worker_backend,
|
||||
}
|
||||
if stats_log_file is not None:
|
||||
kwargs["log_file"] = stats_log_file
|
||||
if args.ray_address is not None:
|
||||
kwargs["ray_address"] = args.ray_address
|
||||
return kwargs
|
||||
|
||||
|
||||
def build_sampling_params(max_decode_steps: int) -> list[SamplingParams]:
|
||||
return [
|
||||
SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=max_decode_steps + 1,
|
||||
stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
|
||||
),
|
||||
SamplingParams(temperature=0.0, max_tokens=1),
|
||||
]
|
||||
|
||||
|
||||
async def run_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
|
||||
engine = AsyncOmni(**build_engine_kwargs(args, stats_log_file))
|
||||
try:
|
||||
all_audio_chunks = []
|
||||
accumulated_samples = 0
|
||||
chunk_idx = 0
|
||||
start_time = time.time()
|
||||
chunk_times = []
|
||||
ttfp_seconds = None
|
||||
final_stage_output = None
|
||||
async for stage_output in engine.generate(
|
||||
prompt=prompt_payload,
|
||||
request_id=str(uuid.uuid4()),
|
||||
sampling_params_list=sampling_params_list,
|
||||
):
|
||||
final_stage_output = stage_output
|
||||
multimodal_output = stage_output.multimodal_output or {}
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
continue
|
||||
|
||||
finished = stage_output.finished
|
||||
if isinstance(audio, torch.Tensor):
|
||||
if finished:
|
||||
audio_chunk = audio[accumulated_samples:].float().detach().cpu()
|
||||
else:
|
||||
audio_chunk = audio.float().detach().cpu()
|
||||
elif isinstance(audio, list):
|
||||
audio_chunk = torch.as_tensor(audio[chunk_idx], dtype=torch.float32).reshape(-1).cpu()
|
||||
else:
|
||||
audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
|
||||
|
||||
accumulated_samples += int(audio_chunk.numel())
|
||||
chunk_idx += 1
|
||||
if audio_chunk.numel() > 0:
|
||||
now = time.time()
|
||||
if ttfp_seconds is None:
|
||||
ttfp_seconds = now - start_time
|
||||
chunk_times.append(now)
|
||||
all_audio_chunks.append(audio_chunk)
|
||||
|
||||
if not all_audio_chunks:
|
||||
raise RuntimeError("Streaming Ming example produced no audio chunks")
|
||||
|
||||
waveform = torch.cat(all_audio_chunks, dim=0)
|
||||
output_name = resolve_output_name(args.output_name, args.case, 0, 1)
|
||||
output_path = str(output_dir / output_name)
|
||||
write_wav(output_path, waveform, SAMPLE_RATE)
|
||||
summary = {
|
||||
"request_id": getattr(final_stage_output, "request_id", None),
|
||||
"stage_id": getattr(final_stage_output, "stage_id", None),
|
||||
"output_path": output_path,
|
||||
"stage_durations": getattr(final_stage_output, "stage_durations", {}),
|
||||
"peak_memory_mb": getattr(final_stage_output, "peak_memory_mb", 0.0),
|
||||
"ttfp_seconds": ttfp_seconds,
|
||||
"mean_inter_chunk_seconds": (
|
||||
sum(t1 - t0 for t0, t1 in zip(chunk_times, chunk_times[1:])) / (len(chunk_times) - 1)
|
||||
if len(chunk_times) > 1
|
||||
else None
|
||||
),
|
||||
}
|
||||
summary.update(audio_summary(waveform, SAMPLE_RATE))
|
||||
print(f"Saved streaming output to {output_path}")
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
return [summary]
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
def run_non_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
|
||||
engine = Omni(**build_engine_kwargs(args, stats_log_file))
|
||||
try:
|
||||
outputs = engine.generate(
|
||||
prompts=[prompt_payload for _ in range(args.num_prompts)],
|
||||
sampling_params_list=sampling_params_list,
|
||||
py_generator=False,
|
||||
)
|
||||
summaries = []
|
||||
for fallback_index, output in enumerate(outputs):
|
||||
if output.final_output_type != "audio":
|
||||
continue
|
||||
multimodal_output = output.multimodal_output or {}
|
||||
waveform = coerce_audio_tensor(multimodal_output.get("audio"), async_chunk=False)
|
||||
sample_rate = extract_sample_rate(multimodal_output)
|
||||
output_name = resolve_output_name(
|
||||
args.output_name,
|
||||
args.case,
|
||||
request_index(output.request_id, fallback_index),
|
||||
args.num_prompts,
|
||||
)
|
||||
output_path = str(output_dir / output_name)
|
||||
write_wav(output_path, waveform, sample_rate)
|
||||
summary = {
|
||||
"request_id": output.request_id,
|
||||
"stage_id": output.stage_id,
|
||||
"output_path": output_path,
|
||||
"stage_durations": output.stage_durations,
|
||||
"peak_memory_mb": output.peak_memory_mb,
|
||||
}
|
||||
summary.update(audio_summary(waveform, sample_rate))
|
||||
summaries.append(summary)
|
||||
print(f"Saved output to {output_path}")
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
if not summaries:
|
||||
raise RuntimeError("Non-streaming Ming example produced no audio outputs")
|
||||
return summaries
|
||||
finally:
|
||||
engine.close()
|
||||
|
||||
|
||||
def run_generation(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
|
||||
if args.streaming:
|
||||
return asyncio.run(run_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file))
|
||||
return run_non_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
|
||||
@@ -212,7 +212,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
|
||||
output_audio_dur = 0.0
|
||||
|
||||
for batch_idx, o in enumerate(outputs):
|
||||
audio_tensor = torch.cat(o.multimodal_output["audio"])
|
||||
audio_tensor = torch.cat((o.multimodal_output["audio"],))
|
||||
audio_array = audio_tensor.tolist()
|
||||
output_audio_dur += float(len(audio_array)) / 24000
|
||||
if args.write_audio:
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
# Speech-To-Video
|
||||
|
||||
This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Model ID |
|
||||
|-------|----------|
|
||||
| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` |
|
||||
|
||||
## Start Server
|
||||
|
||||
### Basic Start
|
||||
|
||||
```bash
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size 2 \
|
||||
--flow-shift 3.0 \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--cache-backend cache_dit \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
### Start with Script
|
||||
|
||||
```bash
|
||||
bash run_server.sh
|
||||
```
|
||||
|
||||
The script allows overriding:
|
||||
- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`)
|
||||
- `PORT` (default: `8091`)
|
||||
- `FLOW_SHIFT` (default: `3.0`)
|
||||
- `TP` (default: `2`)
|
||||
- `CACHE_BACKEND` (default: `cache_dit`)
|
||||
|
||||
### Start with 4 GPUs
|
||||
|
||||
```bash
|
||||
TP=4 bash run_server.sh
|
||||
```
|
||||
|
||||
## Sync API (Recommended for Testing)
|
||||
|
||||
`POST /v1/videos/sync` blocks until generation completes and returns the raw
|
||||
video bytes (`video/mp4`) directly in the response body.
|
||||
|
||||
```bash
|
||||
bash run_curl_speech_to_video.sh
|
||||
```
|
||||
|
||||
## Async Job API
|
||||
|
||||
`POST /v1/videos` creates an asynchronous job and returns immediately.
|
||||
|
||||
```bash
|
||||
no_proxy=127.0.0.1 \
|
||||
create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A person singing" \
|
||||
-F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
|
||||
-F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
|
||||
-F "width=832" -F "height=480" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.5" \
|
||||
-F "fps=16")
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
|
||||
# Poll until complete
|
||||
while true; do
|
||||
status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
# Download
|
||||
curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4
|
||||
```
|
||||
|
||||
## Request Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `prompt` | str | - | Text description of the desired video |
|
||||
| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `width` | int | None | Video width in pixels |
|
||||
| `height` | int | None | Video height in pixels |
|
||||
| `fps` | int | None | Output video frame rate |
|
||||
| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) |
|
||||
| `guidance_scale` | float | None | CFG guidance scale |
|
||||
| `seed` | int | None | Random seed for reproducibility |
|
||||
|
||||
## Notes
|
||||
|
||||
- S2V requires both a reference image (`image_reference`) and an audio reference
|
||||
(`audio_reference`). The generated video will show a person matching the reference
|
||||
image with lip movements synchronized to the audio.
|
||||
- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be:
|
||||
- An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`)
|
||||
- A base64 data URL (e.g., `data:audio/mp3;base64,...`)
|
||||
- `--model-class-name WanS2VPipeline` is required on the server to select the
|
||||
S2V pipeline (distinct from the T2V/I2V pipelines).
|
||||
- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps.
|
||||
- Audio is muxed into the output MP4 automatically.
|
||||
- Always pass `fps=16` to ensure correct video/audio alignment.
|
||||
@@ -0,0 +1,44 @@
|
||||
#!/bin/bash
|
||||
# Wan2.2 S2V (speech-to-video) curl example using the sync video API.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
BASE_URL="${BASE_URL:-http://localhost:8091}"
|
||||
OUTPUT_PATH="${OUTPUT_PATH:-s2v_480p_serve.mp4}"
|
||||
IMAGE_URL="${IMAGE_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png}"
|
||||
AUDIO_URL="${AUDIO_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3}"
|
||||
PROMPT="${PROMPT:-A person singing}"
|
||||
WIDTH="${WIDTH:-832}"
|
||||
HEIGHT="${HEIGHT:-480}"
|
||||
NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-40}"
|
||||
GUIDANCE_SCALE="${GUIDANCE_SCALE:-4.5}"
|
||||
FPS="${FPS:-16}"
|
||||
|
||||
echo "Sending S2V request..."
|
||||
echo " Image URL: $IMAGE_URL"
|
||||
echo " Audio URL: $AUDIO_URL"
|
||||
echo " Prompt: $PROMPT"
|
||||
echo " Resolution: ${WIDTH}x${HEIGHT}"
|
||||
echo " Steps: $NUM_INFERENCE_STEPS"
|
||||
echo " FPS: $FPS"
|
||||
|
||||
IMAGE_REF_JSON="{\"image_url\": \"${IMAGE_URL}\"}"
|
||||
AUDIO_REF_JSON="{\"audio_url\": \"${AUDIO_URL}\"}"
|
||||
|
||||
no_proxy=127.0.0.1 \
|
||||
curl -X POST "${BASE_URL}/v1/videos/sync" \
|
||||
-F "prompt=${PROMPT}" \
|
||||
-F "image_reference=${IMAGE_REF_JSON}" \
|
||||
-F "audio_reference=${AUDIO_REF_JSON}" \
|
||||
-F "width=${WIDTH}" -F "height=${HEIGHT}" \
|
||||
-F "num_inference_steps=${NUM_INFERENCE_STEPS}" \
|
||||
-F "guidance_scale=${GUIDANCE_SCALE}" \
|
||||
-F "fps=${FPS}" \
|
||||
--output "${OUTPUT_PATH}"
|
||||
|
||||
if [ -f "$OUTPUT_PATH" ] && [ -s "$OUTPUT_PATH" ]; then
|
||||
echo "Saved video to ${OUTPUT_PATH} ($(du -h "$OUTPUT_PATH" | cut -f1))"
|
||||
else
|
||||
echo "ERROR: Output file is empty or missing"
|
||||
exit 1
|
||||
fi
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
#!/bin/bash
|
||||
# Wan2.2 S2V (speech-to-video) online serving startup script
|
||||
|
||||
MODEL="${MODEL:-Wan-AI/Wan2.2-S2V-14B}"
|
||||
PORT="${PORT:-8091}"
|
||||
FLOW_SHIFT="${FLOW_SHIFT:-3.0}"
|
||||
TP="${TP:-2}"
|
||||
CACHE_BACKEND="${CACHE_BACKEND:-cache_dit}"
|
||||
|
||||
echo "Starting Wan2.2 S2V server..."
|
||||
echo "Model: $MODEL"
|
||||
echo "Port: $PORT"
|
||||
echo "Flow shift: $FLOW_SHIFT"
|
||||
echo "Tensor parallel size: $TP"
|
||||
echo "Cache backend: $CACHE_BACKEND"
|
||||
|
||||
CACHE_BACKEND_FLAG=""
|
||||
if [ "$CACHE_BACKEND" != "none" ]; then
|
||||
CACHE_BACKEND_FLAG="--cache-backend $CACHE_BACKEND"
|
||||
fi
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
vllm serve "$MODEL" --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size "$TP" \
|
||||
--flow-shift "$FLOW_SHIFT" \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--port "$PORT" \
|
||||
$CACHE_BACKEND_FLAG
|
||||
@@ -0,0 +1,193 @@
|
||||
# Step-Audio2 Online Serving
|
||||
|
||||
This directory contains examples for running Step-Audio2 with vLLM-Omni's online serving API.
|
||||
|
||||
## Installation
|
||||
|
||||
Please refer to [README.md](../../../README.md)
|
||||
|
||||
## Launch the Server
|
||||
|
||||
```bash
|
||||
# Async chunk mode (recommended — lower first-packet latency for TTS)
|
||||
vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2_async_chunk.yaml \
|
||||
--trust-remote-code --enforce-eager
|
||||
```
|
||||
|
||||
Sequential mode:
|
||||
```bash
|
||||
vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092 \
|
||||
--stage-configs-path vllm_omni/model_executor/stage_configs/step_audio_2.yaml \
|
||||
--trust-remote-code --enforce-eager
|
||||
```
|
||||
|
||||
With local model:
|
||||
```bash
|
||||
vllm serve /path/to/Step-Audio-2-mini --omni --port 8092 \
|
||||
--trust-remote-code --enforce-eager
|
||||
```
|
||||
|
||||
## Send Requests
|
||||
|
||||
### TTS via `/v1/audio/speech` (Recommended)
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/step_audio2
|
||||
|
||||
# Python client
|
||||
python openai_speech_client.py --text "你好世界"
|
||||
|
||||
# With custom system prompt
|
||||
python openai_speech_client.py --text "Hello, how are you?" \
|
||||
--instructions "You are a friendly assistant."
|
||||
|
||||
# Save to specific file
|
||||
python openai_speech_client.py --text "你好世界" -o output.wav
|
||||
```
|
||||
|
||||
Or via curl:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8092/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model":"stepfun-ai/Step-Audio2-mini","input":"你好世界","voice":"default"}' \
|
||||
--output output.wav
|
||||
```
|
||||
|
||||
This endpoint bypasses the chat template and directly triggers TTS mode. It supports async chunk streaming for low first-packet latency.
|
||||
|
||||
**Note**: Speaker voice is controlled by `STEP_AUDIO2_DEFAULT_PROMPT_WAV` env var on the server side.
|
||||
|
||||
### Chat Completions (ASR / S2ST)
|
||||
|
||||
```bash
|
||||
# Audio to Text (ASR)
|
||||
python openai_chat_completion_client.py --query-type audio_to_text
|
||||
|
||||
# Audio to Audio (S2ST)
|
||||
python openai_chat_completion_client.py --query-type audio_to_audio --audio-path /path/to/input.wav
|
||||
```
|
||||
|
||||
| Argument | Description |
|
||||
|----------|-------------|
|
||||
| `--query-type`, `-q` | Query type: `audio_to_text`, `text_to_audio`, `audio_to_audio` |
|
||||
| `--audio-path`, `-a` | Path to input audio file (local or URL) |
|
||||
| `--text`, `-t` | Text to synthesize (for TTS mode) |
|
||||
| `--prompt`, `-p` | Custom prompt/question |
|
||||
| `--output-dir`, `-o` | Output directory for audio files (default: `output_online`) |
|
||||
| `--api-base` | API base URL (default: `http://localhost:8092/v1`) |
|
||||
|
||||
### Curl (Chat Completions)
|
||||
|
||||
```bash
|
||||
# Audio to Text
|
||||
bash run_curl.sh audio_to_text
|
||||
|
||||
# Text to Audio
|
||||
bash run_curl.sh text_to_audio
|
||||
|
||||
# Audio to Audio
|
||||
bash run_curl.sh audio_to_audio
|
||||
```
|
||||
|
||||
## Query Types
|
||||
|
||||
### 1. Audio to Text (ASR)
|
||||
|
||||
Transcribe audio to text.
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client.py \
|
||||
--query-type audio_to_text \
|
||||
--audio-path /path/to/speech.wav \
|
||||
--prompt "Transcribe this audio."
|
||||
```
|
||||
|
||||
### 2. Text to Audio (TTS)
|
||||
|
||||
Convert text to speech:
|
||||
|
||||
```bash
|
||||
# Via speech endpoint (recommended, returns WAV directly)
|
||||
python openai_speech_client.py --text "Hello, welcome to Step-Audio2."
|
||||
|
||||
# Via chat completions
|
||||
python openai_chat_completion_client.py \
|
||||
--query-type text_to_audio \
|
||||
--text "Hello, welcome to Step-Audio2."
|
||||
```
|
||||
|
||||
### 3. Audio to Audio (S2ST)
|
||||
|
||||
Process input audio and generate text transcription + audio output.
|
||||
|
||||
```bash
|
||||
python openai_chat_completion_client.py \
|
||||
--query-type audio_to_audio \
|
||||
--audio-path /path/to/source.wav
|
||||
```
|
||||
|
||||
## Output
|
||||
|
||||
- **Text output**: Printed to console
|
||||
- **Audio output**: Saved to `output_online/audio_0.wav` (24kHz WAV)
|
||||
|
||||
## API Format
|
||||
|
||||
Step-Audio2 uses the OpenAI-compatible chat completions API:
|
||||
|
||||
```json
|
||||
{
|
||||
"model": "stepfun-ai/Step-Audio2-mini",
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": "Transcribe the audio."}]
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": "..."}},
|
||||
{"type": "text", "text": "Please transcribe."}
|
||||
]
|
||||
}
|
||||
],
|
||||
"sampling_params_list": [
|
||||
{"temperature": 0.7, "max_tokens": 1024},
|
||||
{"temperature": 0.0, "max_tokens": 1}
|
||||
]
|
||||
}
|
||||
```
|
||||
|
||||
## Performance
|
||||
|
||||
### Async Chunk vs Sequential
|
||||
|
||||
Benchmark via `/v1/audio/speech` (4x RTX 3090, 10 prompts, concurrency=1):
|
||||
|
||||
| Mode | Mean TTFP | Mean E2E | Mean RTF |
|
||||
|------|-----------|----------|----------|
|
||||
| Sequential | 4316ms | 4316ms | 0.938 |
|
||||
| **Async Chunk** | **1437ms** | 4362ms | 0.949 |
|
||||
|
||||
Async chunk reduces TTFP by **67%** by streaming audio token chunks from Thinker to Token2Wav as they are generated. RTF < 1 in both modes (real-time capable).
|
||||
|
||||
## Troubleshooting
|
||||
|
||||
### Server not responding
|
||||
- Check if the server is running: `curl http://localhost:8092/v1/models`
|
||||
- Verify the port number matches
|
||||
|
||||
### FileNotFoundError: prompt_wav file not found
|
||||
- Ensure `default_female.wav` exists at `{model_dir}/assets/default_female.wav`
|
||||
- Or set `STEP_AUDIO2_DEFAULT_PROMPT_WAV` environment variable when launching the server
|
||||
|
||||
### Audio not generated
|
||||
- For TTS, use the `/v1/audio/speech` endpoint (recommended) or `openai_speech_client.py`
|
||||
- For chat completions TTS, ensure the prompt ends with `<tts_start>`
|
||||
- Check server logs for errors
|
||||
|
||||
### Out of memory
|
||||
- Reduce `gpu_memory_utilization` in stage configs
|
||||
- Use a smaller batch size
|
||||
@@ -0,0 +1,342 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
OpenAI-compatible client for Step-Audio2 online serving.
|
||||
|
||||
Step-Audio2 is a two-stage model:
|
||||
- Stage 0 (Thinker): Audio understanding -> Text + Audio tokens
|
||||
- Stage 1 (Token2Wav): Audio tokens -> Waveform (24kHz)
|
||||
|
||||
Usage examples:
|
||||
# Start server first:
|
||||
vllm serve stepfun-ai/Step-Audio2-mini --omni --port 8092
|
||||
|
||||
# Audio to text (ASR)
|
||||
python openai_chat_completion_client.py --query-type audio_to_text
|
||||
|
||||
# Text to audio (TTS)
|
||||
python openai_chat_completion_client.py --query-type text_to_audio --text "Hello world"
|
||||
|
||||
# Audio to audio (Voice conversion)
|
||||
python openai_chat_completion_client.py --query-type audio_to_audio --audio-path input.wav
|
||||
"""
|
||||
|
||||
import base64
|
||||
import logging
|
||||
import os
|
||||
|
||||
import requests
|
||||
from openai import OpenAI
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
logging.basicConfig(level=logging.INFO, format="%(levelname)s: %(message)s")
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
SEED = 42
|
||||
|
||||
|
||||
def encode_base64_content_from_file(file_path: str) -> str:
|
||||
"""Encode a local file to base64 format."""
|
||||
with open(file_path, "rb") as f:
|
||||
return base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
|
||||
def encode_base64_content_from_url(url: str) -> str:
|
||||
"""Encode content from URL to base64 format."""
|
||||
with requests.get(url) as response:
|
||||
response.raise_for_status()
|
||||
return base64.b64encode(response.content).decode("utf-8")
|
||||
|
||||
|
||||
def get_audio_url(audio_path: str | None = None) -> str:
|
||||
"""Convert audio path to URL format for the API.
|
||||
|
||||
Args:
|
||||
audio_path: Local file path or URL. If None, uses default test audio.
|
||||
|
||||
Returns:
|
||||
Audio URL (either original URL or base64 data URL)
|
||||
"""
|
||||
if not audio_path:
|
||||
return AudioAsset("mary_had_lamb").url
|
||||
|
||||
if audio_path.startswith(("http://", "https://")):
|
||||
return audio_path
|
||||
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
|
||||
# Detect MIME type
|
||||
audio_path_lower = audio_path.lower()
|
||||
if audio_path_lower.endswith(".wav"):
|
||||
mime_type = "audio/wav"
|
||||
elif audio_path_lower.endswith((".mp3", ".mpeg")):
|
||||
mime_type = "audio/mpeg"
|
||||
elif audio_path_lower.endswith(".ogg"):
|
||||
mime_type = "audio/ogg"
|
||||
elif audio_path_lower.endswith(".flac"):
|
||||
mime_type = "audio/flac"
|
||||
else:
|
||||
mime_type = "audio/wav"
|
||||
|
||||
audio_base64 = encode_base64_content_from_file(audio_path)
|
||||
return f"data:{mime_type};base64,{audio_base64}"
|
||||
|
||||
|
||||
def get_audio_to_text_query(
|
||||
audio_path: str | None = None,
|
||||
custom_prompt: str | None = None,
|
||||
) -> dict:
|
||||
"""Build query for Audio-to-Text (ASR) mode."""
|
||||
question = custom_prompt or "Please transcribe the audio content."
|
||||
audio_url = get_audio_url(audio_path)
|
||||
|
||||
return {
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": audio_url}},
|
||||
{"type": "text", "text": question},
|
||||
],
|
||||
}
|
||||
|
||||
|
||||
def get_text_to_audio_query(
|
||||
text: str | None = None,
|
||||
custom_prompt: str | None = None,
|
||||
) -> dict:
|
||||
"""Build query for Text-to-Audio (TTS) mode.
|
||||
|
||||
Returns a list of messages: user message + assistant message with <tts_start>.
|
||||
The assistant message triggers TTS mode and must be used with
|
||||
continue_final_message=True so the chat template does not append <|im_end|>.
|
||||
"""
|
||||
text_to_speak = text or "Hello, this is a test of Step Audio 2 text to speech."
|
||||
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [{"type": "text", "text": text_to_speak}],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "<tts_start>",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def get_audio_to_audio_query(
|
||||
audio_path: str | None = None,
|
||||
custom_prompt: str | None = None,
|
||||
) -> dict:
|
||||
"""Build query for Audio-to-Audio (S2ST) mode.
|
||||
|
||||
Returns a list of messages: user message (with audio) + assistant message
|
||||
with <tts_start>. Must be used with continue_final_message=True.
|
||||
"""
|
||||
question = custom_prompt or "Please listen to this audio and repeat its content."
|
||||
audio_url = get_audio_url(audio_path)
|
||||
|
||||
return [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": audio_url}},
|
||||
{"type": "text", "text": question},
|
||||
],
|
||||
},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": "<tts_start>",
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def get_system_prompt(query_type: str) -> dict:
|
||||
"""Get system prompt based on query type."""
|
||||
if query_type == "audio_to_text":
|
||||
system_text = "You are a speech recognition assistant. Transcribe the audio accurately."
|
||||
elif query_type == "text_to_audio":
|
||||
system_text = "You are a text-to-speech assistant. Read the text aloud exactly as provided."
|
||||
else: # audio_to_audio
|
||||
system_text = "You are an audio processing assistant. Listen and repeat the audio content."
|
||||
|
||||
return {
|
||||
"role": "system",
|
||||
"content": [{"type": "text", "text": system_text}],
|
||||
}
|
||||
|
||||
|
||||
QUERY_MAP = {
|
||||
"audio_to_text": get_audio_to_text_query,
|
||||
"text_to_audio": get_text_to_audio_query,
|
||||
"audio_to_audio": get_audio_to_audio_query,
|
||||
}
|
||||
|
||||
|
||||
def run_step_audio2(args) -> None:
|
||||
"""Run Step-Audio2 inference via OpenAI-compatible API."""
|
||||
|
||||
# Initialize OpenAI client
|
||||
client = OpenAI(
|
||||
api_key="EMPTY",
|
||||
base_url=args.api_base,
|
||||
)
|
||||
|
||||
model_name = args.model
|
||||
|
||||
# Sampling parameters for each stage
|
||||
thinker_sampling_params = {
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"top_k": -1,
|
||||
"max_tokens": args.max_tokens,
|
||||
"seed": SEED,
|
||||
"detokenize": True,
|
||||
"repetition_penalty": 1.1 if args.query_type != "audio_to_text" else 1.05,
|
||||
}
|
||||
|
||||
# For ASR, stop at EOS token
|
||||
if args.query_type == "audio_to_text":
|
||||
thinker_sampling_params["stop_token_ids"] = [151645]
|
||||
|
||||
token2wav_sampling_params = {
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 1,
|
||||
"seed": SEED,
|
||||
"detokenize": False,
|
||||
}
|
||||
|
||||
sampling_params_list = [
|
||||
thinker_sampling_params,
|
||||
token2wav_sampling_params,
|
||||
]
|
||||
|
||||
# Build query based on type
|
||||
query_func = QUERY_MAP[args.query_type]
|
||||
if args.query_type == "audio_to_text":
|
||||
user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt)
|
||||
elif args.query_type == "text_to_audio":
|
||||
user_message = query_func(text=args.text, custom_prompt=args.prompt)
|
||||
else: # audio_to_audio
|
||||
user_message = query_func(audio_path=args.audio_path, custom_prompt=args.prompt)
|
||||
|
||||
system_message = get_system_prompt(args.query_type)
|
||||
|
||||
logger.info(f"Query type: {args.query_type}")
|
||||
logger.info(f"Model: {model_name}")
|
||||
logger.info("Sending request to server...")
|
||||
|
||||
# Build messages list.
|
||||
# TTS and S2ST return a list of messages (user + assistant with <tts_start>);
|
||||
# ASR returns a single user message dict.
|
||||
if isinstance(user_message, list):
|
||||
messages = [system_message] + user_message
|
||||
else:
|
||||
messages = [system_message, user_message]
|
||||
|
||||
# For TTS/S2ST, use continue_final_message so <tts_start> in the
|
||||
# assistant turn is not followed by <|im_end|>.
|
||||
extra_body: dict = {"sampling_params_list": sampling_params_list}
|
||||
if args.query_type in ("text_to_audio", "audio_to_audio"):
|
||||
extra_body["continue_final_message"] = True
|
||||
extra_body["add_generation_prompt"] = False
|
||||
|
||||
# Send request
|
||||
chat_completion = client.chat.completions.create(
|
||||
messages=messages,
|
||||
model=model_name,
|
||||
extra_body=extra_body,
|
||||
)
|
||||
|
||||
# Process response
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
audio_count = 0
|
||||
for choice in chat_completion.choices:
|
||||
if choice.message.audio:
|
||||
# Audio response
|
||||
audio_data = base64.b64decode(choice.message.audio.data)
|
||||
audio_file = os.path.join(args.output_dir, f"audio_{audio_count}.wav")
|
||||
with open(audio_file, "wb") as f:
|
||||
f.write(audio_data)
|
||||
logger.info(f"Audio saved to: {audio_file}")
|
||||
audio_count += 1
|
||||
|
||||
if choice.message.content:
|
||||
# Text response
|
||||
logger.info(f"Text output: {choice.message.content}")
|
||||
|
||||
if audio_count == 0 and args.query_type != "audio_to_text":
|
||||
logger.warning("No audio output received. Check server logs for details.")
|
||||
|
||||
logger.info("Done!")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="OpenAI-compatible client for Step-Audio2 online serving")
|
||||
|
||||
parser.add_argument(
|
||||
"--query-type",
|
||||
"-q",
|
||||
type=str,
|
||||
default="audio_to_text",
|
||||
choices=QUERY_MAP.keys(),
|
||||
help="Query type: audio_to_text (ASR), text_to_audio (TTS), audio_to_audio",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--audio-path",
|
||||
"-a",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to input audio file (local path or URL)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
"-t",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Text to synthesize (for text_to_audio mode)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
"-p",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Custom prompt/question",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
"-m",
|
||||
type=str,
|
||||
default="stepfun-ai/Step-Audio2-mini",
|
||||
help="Model name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-base",
|
||||
type=str,
|
||||
default="http://localhost:8092/v1",
|
||||
help="API base URL",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-tokens",
|
||||
type=int,
|
||||
default=1024,
|
||||
help="Maximum tokens for Thinker stage",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
"-o",
|
||||
type=str,
|
||||
default="output_online",
|
||||
help="Output directory for audio files",
|
||||
)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_step_audio2(args)
|
||||
@@ -0,0 +1,130 @@
|
||||
"""OpenAI-compatible client for Step-Audio2 TTS via /v1/audio/speech endpoint.
|
||||
|
||||
Examples:
|
||||
# Basic TTS
|
||||
python openai_speech_client.py --text "你好世界"
|
||||
|
||||
# With custom system prompt
|
||||
python openai_speech_client.py --text "Hello, how are you?" \
|
||||
--instructions "You are a friendly assistant."
|
||||
|
||||
# Save to specific file
|
||||
python openai_speech_client.py --text "你好世界" -o output.wav
|
||||
"""
|
||||
|
||||
import argparse
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8092"
|
||||
DEFAULT_API_KEY = "EMPTY"
|
||||
|
||||
|
||||
def run_tts_generation(args) -> None:
|
||||
"""Run TTS generation via /v1/audio/speech API."""
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"voice": args.voice,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
|
||||
if args.instructions:
|
||||
payload["instructions"] = args.instructions
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Text: {args.text}")
|
||||
print(f"Voice: {args.voice}")
|
||||
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 "tts_output.wav"
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(response.content)
|
||||
print(f"Audio saved to: {output_path}")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Step-Audio2 TTS client via /v1/audio/speech",
|
||||
formatter_class=argparse.RawDescriptionHelpFormatter,
|
||||
epilog=__doc__,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-base",
|
||||
type=str,
|
||||
default=DEFAULT_API_BASE,
|
||||
help=f"API base URL (default: {DEFAULT_API_BASE})",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--api-key",
|
||||
type=str,
|
||||
default=DEFAULT_API_KEY,
|
||||
help="API key (default: EMPTY)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
"-m",
|
||||
type=str,
|
||||
default="stepfun-ai/Step-Audio2-mini",
|
||||
help="Model name/path",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text",
|
||||
type=str,
|
||||
required=True,
|
||||
help="Text to synthesize",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--voice",
|
||||
type=str,
|
||||
default="default",
|
||||
help="Voice name (default: default)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--instructions",
|
||||
type=str,
|
||||
default=None,
|
||||
help="System prompt for the Thinker stage",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--response-format",
|
||||
type=str,
|
||||
default="wav",
|
||||
choices=["wav", "pcm"],
|
||||
help="Audio output format (default: wav)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Output audio file path (default: tts_output.wav)",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
args = parse_args()
|
||||
run_tts_generation(args)
|
||||
+149
@@ -0,0 +1,149 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Step-Audio2 curl client for online serving
|
||||
# Usage: bash run_curl.sh [audio_to_text|text_to_audio|audio_to_audio]
|
||||
|
||||
QUERY_TYPE="${1:-audio_to_text}"
|
||||
API_BASE="${API_BASE:-http://localhost:8092}"
|
||||
|
||||
# Validate query type
|
||||
if [[ ! "$QUERY_TYPE" =~ ^(audio_to_text|text_to_audio|audio_to_audio)$ ]]; then
|
||||
echo "Error: Invalid query type '$QUERY_TYPE'"
|
||||
echo "Usage: $0 [audio_to_text|text_to_audio|audio_to_audio]"
|
||||
echo " audio_to_text: Speech recognition (ASR)"
|
||||
echo " text_to_audio: Text-to-speech (TTS)"
|
||||
echo " audio_to_audio: Voice conversion"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
SEED=42
|
||||
|
||||
# Default test audio URL
|
||||
MARY_HAD_LAMB_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
|
||||
|
||||
# Sampling parameters for Thinker stage
|
||||
thinker_sampling_params='{
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"top_k": -1,
|
||||
"max_tokens": 1024,
|
||||
"seed": 42,
|
||||
"detokenize": true,
|
||||
"repetition_penalty": 1.05
|
||||
}'
|
||||
|
||||
# Sampling parameters for Token2Wav stage
|
||||
token2wav_sampling_params='{
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 1,
|
||||
"seed": 42,
|
||||
"detokenize": false
|
||||
}'
|
||||
|
||||
# Build request based on query type
|
||||
case "$QUERY_TYPE" in
|
||||
audio_to_text)
|
||||
system_content='[{"type": "text", "text": "You are a speech recognition assistant. Transcribe the audio accurately."}]'
|
||||
user_content='[
|
||||
{"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}},
|
||||
{"type": "text", "text": "Please transcribe this audio."}
|
||||
]'
|
||||
# Add stop token for ASR
|
||||
thinker_sampling_params='{
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"top_k": -1,
|
||||
"max_tokens": 1024,
|
||||
"seed": 42,
|
||||
"detokenize": true,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [151645]
|
||||
}'
|
||||
;;
|
||||
text_to_audio)
|
||||
system_content='[{"type": "text", "text": "You are a text-to-speech assistant. Read the text aloud exactly as provided."}]'
|
||||
user_content='[
|
||||
{"type": "text", "text": "Hello, this is a test of Step Audio 2 text to speech synthesis.<tts_start>"}
|
||||
]'
|
||||
thinker_sampling_params='{
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"top_k": -1,
|
||||
"max_tokens": 1024,
|
||||
"seed": 42,
|
||||
"detokenize": true,
|
||||
"repetition_penalty": 1.1
|
||||
}'
|
||||
;;
|
||||
audio_to_audio)
|
||||
system_content='[{"type": "text", "text": "You are an audio processing assistant. Listen and repeat the audio content."}]'
|
||||
user_content='[
|
||||
{"type": "audio_url", "audio_url": {"url": "'"$MARY_HAD_LAMB_URL"'"}},
|
||||
{"type": "text", "text": "Please listen to this audio and repeat its content.<tts_start>"}
|
||||
]'
|
||||
thinker_sampling_params='{
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
"top_k": -1,
|
||||
"max_tokens": 1024,
|
||||
"seed": 42,
|
||||
"detokenize": true,
|
||||
"repetition_penalty": 1.1
|
||||
}'
|
||||
;;
|
||||
esac
|
||||
|
||||
sampling_params_list='[
|
||||
'"$thinker_sampling_params"',
|
||||
'"$token2wav_sampling_params"'
|
||||
]'
|
||||
|
||||
echo "Query type: $QUERY_TYPE"
|
||||
echo "API base: $API_BASE"
|
||||
echo "Sending request..."
|
||||
echo ""
|
||||
|
||||
output=$(curl -sS -X POST "${API_BASE}/v1/chat/completions" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @- <<EOF
|
||||
{
|
||||
"model": "stepfun-ai/Step-Audio2-mini",
|
||||
"sampling_params_list": $sampling_params_list,
|
||||
"messages": [
|
||||
{
|
||||
"role": "system",
|
||||
"content": $system_content
|
||||
},
|
||||
{
|
||||
"role": "user",
|
||||
"content": $user_content
|
||||
}
|
||||
]
|
||||
}
|
||||
EOF
|
||||
)
|
||||
|
||||
# Extract and display text content
|
||||
text_content=$(echo "$output" | jq -r '.choices[0].message.content // empty')
|
||||
if [[ -n "$text_content" ]]; then
|
||||
echo "Text output: $text_content"
|
||||
fi
|
||||
|
||||
# Check for audio content
|
||||
audio_data=$(echo "$output" | jq -r '.choices[0].message.audio.data // empty')
|
||||
if [[ -n "$audio_data" ]]; then
|
||||
echo "Audio output received (base64 encoded)"
|
||||
echo "To save audio, use the Python client or decode the base64 data"
|
||||
fi
|
||||
|
||||
# Check for errors
|
||||
error=$(echo "$output" | jq -r '.error // empty')
|
||||
if [[ -n "$error" ]]; then
|
||||
echo "Error: $error"
|
||||
fi
|
||||
|
||||
echo ""
|
||||
echo "Done!"
|
||||
@@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
|
||||
@@ -202,6 +203,51 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
|
||||
|
||||
---
|
||||
|
||||
## Ming-omni-tts
|
||||
|
||||
Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
bash examples/online_serving/text_to_speech/ming_tts/run_server.sh
|
||||
```
|
||||
Equivalent manual command:
|
||||
```bash
|
||||
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--host 0.0.0.0 --port 8091 \
|
||||
--enforce-eager --omni
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "你好,这是 Ming 在线语音合成测试。"
|
||||
```
|
||||
|
||||
Structured dialect control:
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' \
|
||||
--ref-audio /path/to/yue_prompt.wav
|
||||
```
|
||||
|
||||
Zero-shot cloning:
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
|
||||
--ref-audio /path/to/10002287-00000094.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。"
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook.
|
||||
- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples.
|
||||
- Full request details live in [`ming_tts/README.md`](ming_tts/README.md).
|
||||
|
||||
---
|
||||
|
||||
## Ming-flash-omni-TTS
|
||||
|
||||
Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`.
|
||||
|
||||
@@ -0,0 +1,113 @@
|
||||
# Ming-omni-tts Online Serving
|
||||
|
||||
Serve the dense `inclusionAI/Ming-omni-tts-0.5B` two-stage TTS model through
|
||||
the OpenAI-compatible `/v1/audio/speech` endpoint.
|
||||
|
||||
## Start Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Or:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/ming_tts
|
||||
./run_server.sh
|
||||
```
|
||||
|
||||
The tested ROCm environment is summarized in the
|
||||
[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
|
||||
|
||||
## Send Requests
|
||||
|
||||
The Python client targets `http://localhost:8091/v1` with `api_key=EMPTY`; it
|
||||
does not call OpenAI's hosted API.
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "你好,这是 Ming 在线语音合成测试。" \
|
||||
--max-new-tokens 200
|
||||
```
|
||||
|
||||
Style or dialect controls can be plain text or Ming JSON. The upstream
|
||||
dialect example also uses `yue_prompt.wav` for speaker conditioning:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' \
|
||||
--ref-audio /path/to/yue_prompt.wav \
|
||||
--max-new-tokens 200
|
||||
```
|
||||
|
||||
When `--ref-audio` is supplied without `--ref-text`, the server extracts the
|
||||
Ming speaker embedding, matching upstream `use_spk_emb=True`, without using the
|
||||
audio as a zero-shot prompt.
|
||||
|
||||
Reference-audio cloning:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
|
||||
--ref-audio /path/to/10002287-00000094.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--max-new-tokens 200
|
||||
```
|
||||
|
||||
Podcast-style multi-speaker prompt:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text " speaker_1:你可以说一下,就大概说一下,可能虽然我也不知道,我看过那部电影没有。
|
||||
speaker_2:就是那个叫什么,变相一节课的嘛。
|
||||
speaker_1:嗯。
|
||||
speaker_2:一部搞笑的电影。
|
||||
speaker_1:一部搞笑的。" \
|
||||
--ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav \
|
||||
--ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
|
||||
--ref-text " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试,做个,当服务员还要去笔试了
|
||||
speaker_2:对啊,这真的很奇怪,就是 单纯的因,单纯自己工资不高,只是因为可能人家那个店比较出名一点,就对你苛刻要求"
|
||||
```
|
||||
|
||||
Streaming PCM:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "你好,这是流式输出测试。" \
|
||||
--stream \
|
||||
--output ming_output.pcm
|
||||
```
|
||||
|
||||
`run_curl.sh` keeps small smoke checks:
|
||||
|
||||
```bash
|
||||
./run_curl.sh basic
|
||||
REF_AUDIO=/path/to/reference.wav REF_TEXT="在此奉劝大家别乱打美白针。" ./run_curl.sh zero_shot
|
||||
./run_curl.sh stream
|
||||
```
|
||||
|
||||
## Request Fields
|
||||
|
||||
| Field | Ming meaning |
|
||||
|-------|--------------|
|
||||
| `input` | target text |
|
||||
| `instructions` | plain style text, or JSON object for structured Ming controls |
|
||||
| `voice` | Ming IP voice label unless it resolves to an uploaded speaker |
|
||||
| `language` | Ming `方言` control |
|
||||
| `ref_audio` | speaker reference; with `ref_text`, also supplies the prompt waveform |
|
||||
| `ref_text` | transcript enabling zero-shot or podcast prompt-latent conditioning |
|
||||
| `speaker_embedding` | 192-d Ming speaker embedding |
|
||||
| `max_new_tokens` | Ming `max_decode_steps` |
|
||||
|
||||
## Notes
|
||||
|
||||
- `ref_audio` accepts local paths through the client, remote URLs, `file://`,
|
||||
or `data:` URLs.
|
||||
- Non-streaming responses return WAV bytes; streaming responses return PCM.
|
||||
- Music-only `bgm` generation is offline-only until the API exposes Ming
|
||||
prompt-mode selection.
|
||||
@@ -0,0 +1,225 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech.
|
||||
|
||||
Examples:
|
||||
python openai_speech_client.py --text "你好,世界"
|
||||
python openai_speech_client.py --text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' --ref-audio yue_prompt.wav --max-new-tokens 200
|
||||
python openai_speech_client.py \
|
||||
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
|
||||
--ref-audio 10002287-00000094.wav --ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--max-new-tokens 200
|
||||
python openai_speech_client.py --text "speaker_1:你好。 speaker_2:你好。" \
|
||||
--ref-audio speaker_1.wav --ref-audio speaker_2.wav --ref-text "speaker_1:你好。 speaker_2:你好。"
|
||||
python openai_speech_client.py --text "你好,这是流式输出测试。" \
|
||||
--stream --output ming_output.pcm
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8091"
|
||||
DEFAULT_API_KEY = "EMPTY"
|
||||
DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
|
||||
EXPECTED_SPEAKER_EMBEDDING_DIM = 192
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_path: str) -> str:
|
||||
"""Encode a local audio file to a base64 data URL."""
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
ext = audio_path.lower().rsplit(".", 1)[-1]
|
||||
mime_map = {
|
||||
"wav": "audio/wav",
|
||||
"mp3": "audio/mpeg",
|
||||
"flac": "audio/flac",
|
||||
"ogg": "audio/ogg",
|
||||
"aac": "audio/aac",
|
||||
}
|
||||
mime_type = mime_map.get(ext, "audio/wav")
|
||||
with open(audio_path, "rb") as f:
|
||||
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
return f"data:{mime_type};base64,{audio_b64}"
|
||||
|
||||
|
||||
def load_speaker_embedding(path: str) -> list[float]:
|
||||
"""Load and validate a 192-d Ming speaker embedding JSON file."""
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
if not isinstance(data, list):
|
||||
raise ValueError("speaker_embedding file must contain a JSON list")
|
||||
if len(data) != EXPECTED_SPEAKER_EMBEDDING_DIM:
|
||||
raise ValueError(
|
||||
f"Ming dense speaker_embedding must have {EXPECTED_SPEAKER_EMBEDDING_DIM} values, got {len(data)}"
|
||||
)
|
||||
|
||||
values = []
|
||||
for index, value in enumerate(data):
|
||||
try:
|
||||
values.append(float(value))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"speaker_embedding[{index}] must be a number, got {value!r}") from exc
|
||||
return values
|
||||
|
||||
|
||||
def build_instruction_payload(args) -> str | None:
|
||||
"""Return a string payload for the API `instructions` field."""
|
||||
if args.instructions and args.instruction_json:
|
||||
raise ValueError("Use either --instructions or --instruction-json, not both")
|
||||
if args.instruction_json:
|
||||
parsed = json.loads(args.instruction_json)
|
||||
return json.dumps(parsed, ensure_ascii=False)
|
||||
return args.instructions
|
||||
|
||||
|
||||
def validate_args(args) -> None:
|
||||
"""Fail fast on invalid combinations before hitting the server."""
|
||||
if args.ref_text and not args.ref_audio:
|
||||
raise ValueError("--ref-audio is required when --ref-text is provided")
|
||||
if args.speaker_embedding and args.ref_audio and len(args.ref_audio) > 1:
|
||||
raise ValueError("--speaker-embedding cannot be combined with multiple --ref-audio values")
|
||||
|
||||
|
||||
def run_tts(args) -> None:
|
||||
"""Generate speech via the OpenAI-compatible /v1/audio/speech API."""
|
||||
validate_args(args)
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
|
||||
if args.voice:
|
||||
payload["voice"] = args.voice
|
||||
if args.task_type:
|
||||
payload["task_type"] = args.task_type
|
||||
if args.dialect:
|
||||
payload["language"] = args.dialect
|
||||
|
||||
instructions = build_instruction_payload(args)
|
||||
if instructions:
|
||||
payload["instructions"] = instructions
|
||||
|
||||
if args.ref_audio:
|
||||
ref_audio = []
|
||||
for audio in args.ref_audio:
|
||||
if audio.startswith(("http://", "https://", "data:")):
|
||||
ref_audio.append(audio)
|
||||
else:
|
||||
ref_audio.append(encode_audio_to_base64(audio))
|
||||
payload["ref_audio"] = ref_audio[0] if len(ref_audio) == 1 else ref_audio
|
||||
if args.ref_text:
|
||||
payload["ref_text"] = args.ref_text
|
||||
if args.speaker_embedding:
|
||||
payload["speaker_embedding"] = load_speaker_embedding(args.speaker_embedding)
|
||||
if args.max_new_tokens:
|
||||
payload["max_new_tokens"] = args.max_new_tokens
|
||||
if args.stream:
|
||||
payload["stream"] = True
|
||||
payload["response_format"] = "pcm"
|
||||
|
||||
api_url = f"{args.api_base}/v1/audio/speech"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {args.api_key}",
|
||||
}
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Text: {args.text}")
|
||||
print(f"Payload keys: {sorted(payload)}")
|
||||
|
||||
if args.stream:
|
||||
output_path = args.output or "ming_output.pcm"
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
with client.stream("POST", api_url, json=payload, headers=headers) as response:
|
||||
if response.status_code != 200:
|
||||
print(f"Error: {response.status_code}")
|
||||
print(response.read().decode())
|
||||
return
|
||||
with open(output_path, "wb") as f:
|
||||
for chunk in response.iter_bytes():
|
||||
f.write(chunk)
|
||||
print(f"Streamed PCM audio to: {output_path}")
|
||||
return
|
||||
|
||||
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 "ming_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="OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech")
|
||||
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=DEFAULT_MODEL, help="Model name or path")
|
||||
parser.add_argument("--text", required=True, help="Text to synthesize")
|
||||
parser.add_argument(
|
||||
"--task-type",
|
||||
default=None,
|
||||
choices=["CustomVoice", "VoiceDesign", "Base"],
|
||||
help="Optional compatibility task type. Ming accepts the same field but primarily uses prompt metadata.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--voice",
|
||||
default=None,
|
||||
help="Maps to Ming `IP` when using built-in character voices, or to an uploaded voice sample name",
|
||||
)
|
||||
parser.add_argument("--dialect", default=None, help="Maps to Ming `方言`")
|
||||
parser.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
|
||||
parser.add_argument(
|
||||
"--instruction-json",
|
||||
default=None,
|
||||
help='Structured Ming instruction JSON, for example \'{"情感":"高兴"}\'',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
action="append",
|
||||
default=None,
|
||||
help="Reference audio path, URL, or data URL. Repeat for podcast-style multi-speaker prompts.",
|
||||
)
|
||||
parser.add_argument("--ref-text", default=None, help="Reference transcript for cloning")
|
||||
parser.add_argument(
|
||||
"--speaker-embedding", default=None, help="Path to a JSON file containing a 192-d speaker embedding"
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=None, help="Override ming_max_decode_steps")
|
||||
parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output")
|
||||
parser.add_argument(
|
||||
"--response-format",
|
||||
default="wav",
|
||||
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
|
||||
help="Audio format when not streaming",
|
||||
)
|
||||
parser.add_argument("--output", "-o", default=None, help="Output file path")
|
||||
args = parser.parse_args()
|
||||
try:
|
||||
run_tts(args)
|
||||
except (FileNotFoundError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise SystemExit(f"Error: {exc}") from exc
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,76 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
MODE="${1:-basic}"
|
||||
HOST="${HOST:-localhost}"
|
||||
PORT="${PORT:-8091}"
|
||||
MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
|
||||
API_URL="http://${HOST}:${PORT}/v1/audio/speech"
|
||||
TEXT="${TEXT:-你好,这是 Ming 在线语音合成测试。}"
|
||||
OUTPUT="${OUTPUT:-ming_output.wav}"
|
||||
STREAM_OUTPUT="${STREAM_OUTPUT:-ming_output.pcm}"
|
||||
REF_AUDIO="${REF_AUDIO:-}"
|
||||
REF_TEXT="${REF_TEXT:-}"
|
||||
|
||||
post_json() {
|
||||
local payload="$1"
|
||||
local output_path="$2"
|
||||
curl -X POST "$API_URL" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer EMPTY" \
|
||||
-d "$payload" \
|
||||
--output "$output_path"
|
||||
}
|
||||
|
||||
case "$MODE" in
|
||||
basic)
|
||||
post_json "{
|
||||
\"model\": \"${MODEL}\",
|
||||
\"input\": \"${TEXT}\",
|
||||
\"response_format\": \"wav\"
|
||||
}" "$OUTPUT"
|
||||
;;
|
||||
zero_shot)
|
||||
if [ -z "$REF_AUDIO" ] || [ -z "$REF_TEXT" ]; then
|
||||
echo "zero_shot requires REF_AUDIO and REF_TEXT" >&2
|
||||
exit 1
|
||||
fi
|
||||
python - <<'PY' > /tmp/ming_zero_shot_payload.json
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
path = Path(os.environ["REF_AUDIO"])
|
||||
mime_type = mimetypes.guess_type(path.name)[0] or "audio/wav"
|
||||
payload = {
|
||||
"model": os.environ["MODEL"],
|
||||
"input": os.environ["TEXT"],
|
||||
"ref_audio": f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('utf-8')}",
|
||||
"ref_text": os.environ["REF_TEXT"],
|
||||
"response_format": "wav",
|
||||
}
|
||||
print(json.dumps(payload, ensure_ascii=False))
|
||||
PY
|
||||
curl -X POST "$API_URL" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer EMPTY" \
|
||||
--data-binary @/tmp/ming_zero_shot_payload.json \
|
||||
--output "$OUTPUT"
|
||||
rm -f /tmp/ming_zero_shot_payload.json
|
||||
;;
|
||||
stream)
|
||||
post_json "{
|
||||
\"model\": \"${MODEL}\",
|
||||
\"input\": \"${TEXT}\",
|
||||
\"stream\": true,
|
||||
\"response_format\": \"pcm\"
|
||||
}" "$STREAM_OUTPUT"
|
||||
;;
|
||||
*)
|
||||
echo "Unknown mode: $MODE" >&2
|
||||
echo "Supported sanity checks: basic, zero_shot, stream" >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for Ming-omni-tts.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh
|
||||
# PORT=8000 ./run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
ROOT="$(cd "$DIR/../../../.." && pwd)"
|
||||
|
||||
MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
|
||||
PORT="${PORT:-8091}"
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT/vllm_omni/deploy/ming_tts.yaml}"
|
||||
|
||||
echo "Starting Ming-omni-tts server with model: $MODEL"
|
||||
echo "Deploy config: $DEPLOY_CONFIG"
|
||||
|
||||
vllm-omni serve "$MODEL" \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--enforce-eager \
|
||||
--omni
|
||||
@@ -408,7 +408,10 @@ def create_app(api_base: str):
|
||||
await resp.aclose()
|
||||
await client.aclose()
|
||||
|
||||
return StreamingResponse(relay(), media_type="application/octet-stream")
|
||||
# Use audio/pcm to bypass browsers (e.g. Chromium) MIME type sniffing, which
|
||||
# buffers application/octet-stream responses before delivering chunks to
|
||||
# fetch().body.getReader(), inflating TTFP to total_ms.
|
||||
return StreamingResponse(relay(), media_type="audio/pcm")
|
||||
|
||||
# ── Gradio UI ────────────────────────────────────────────────
|
||||
voices = fetch_voices(api_base)
|
||||
|
||||
+7
-1
@@ -66,7 +66,8 @@ dev = [
|
||||
"voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
|
||||
"torchdiffeq>=0.2.5", # for covo_audio Test
|
||||
"FlagEmbedding", # for hunyuan_image3 tests
|
||||
"kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
||||
"kernels<0.15.0", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
||||
"s3tokenizer>=0.3.0" # for tts test
|
||||
]
|
||||
|
||||
demo = [
|
||||
@@ -280,6 +281,7 @@ extend-ignore-identifiers-re = [
|
||||
".*[Oo]no_[Aa]nna.*",
|
||||
".*cann.*",
|
||||
".*[Nn][Dd]3.*",
|
||||
".*updator.*",
|
||||
"byt5.*",
|
||||
]
|
||||
extend-ignore-re = [
|
||||
@@ -292,5 +294,9 @@ fullset = "fullset"
|
||||
Vai = "Vai"
|
||||
tockens = "tockens"
|
||||
CANN = "CANN"
|
||||
# Step-Audio reasoning parser streams `</think>` tokens split across chunks
|
||||
# (e.g. "</" + "thi" + "nk>"). The fragment "thi" appears literally in its
|
||||
# parser logic and tests and is not a typo of "the"/"this".
|
||||
thi = "thi"
|
||||
ND = "ND"
|
||||
byt5 = "byt5"
|
||||
@@ -52,6 +52,7 @@ vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
|
||||
|
||||
The default deploy config `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` is
|
||||
loaded automatically by model registry (HF `model_type=higgs_multimodal_qwen3`).
|
||||
It matches the high-throughput profile.
|
||||
|
||||
**Offline batch inference:**
|
||||
|
||||
@@ -103,8 +104,28 @@ python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client
|
||||
- Voice cloning: `ref_audio` accepts WAV/FLAC/MP3; `ref_text` is optional but improves fidelity.
|
||||
- Deploy config: `vllm_omni/deploy/higgs_multimodal_qwen3.yaml` (auto-discovered from `model_type`).
|
||||
- `max_num_seqs=16` for both stages.
|
||||
- `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
|
||||
- Stage 0 and Stage 1 default to the same device (`0`) for single-GPU serving.
|
||||
- Stage 0 intentionally keeps `enforce_eager=true`. This preserves the Higgs-specific local MLP CUDA graph path, which is the current high-throughput default.
|
||||
- Stage 1 remains `enforce_eager=true` for the codec decoder.
|
||||
- Deploy profiles:
|
||||
- High throughput: `vllm_omni/deploy/higgs_multimodal_qwen3_high_throughput.yaml`.
|
||||
Use this for medium/high concurrency. The auto-discovered
|
||||
`higgs_multimodal_qwen3.yaml` is kept as a compatibility/default alias for
|
||||
this profile.
|
||||
- Low latency: `vllm_omni/deploy/higgs_multimodal_qwen3_low_latency.yaml`.
|
||||
Use this for low-concurrency serving (for example c1-c4). It sets Stage 0
|
||||
`enforce_eager=false` and enables vLLM `FULL_DECODE_ONLY` CUDA graph through
|
||||
YAML `compilation_config`; no environment variable is required.
|
||||
- Profile details: `vllm_omni/deploy/README_higgs_audio_v3.md`.
|
||||
- Performance note:
|
||||
- Do not switch the auto-discovered default Stage 0 profile to vLLM
|
||||
`FULL_DECODE_ONLY` CUDA graph without an end-to-end throughput and
|
||||
audio-quality revalidation. On the H20 SeedTTS c16/full-dataset benchmark,
|
||||
the eager Stage 0 path with Higgs local MLP CUDA graph reproduced
|
||||
~35 audio_s/s (`1088/1088` OK per run, three runs around 134-136s). A
|
||||
separate H20 low-concurrency smoke after the FULL_DECODE audio-feedback
|
||||
capture fix showed lower c1/c4 latency, so the FULL_DECODE profile is kept
|
||||
as an explicit low-latency option rather than the throughput default.
|
||||
- Known limitations:
|
||||
- CUDA graph for stage 0 helps batch=1 (-13% RTF) but not batch>1 due to model-owned sampler running outside graph.
|
||||
- Stage 1 (code2wav) must use `enforce_eager=true` (`@torch.inference_mode` incompatible with graph capture).
|
||||
- Async streaming (chunk-based) not yet implemented; sync-only pipeline.
|
||||
- Stage 0 full-decode CUDA graph is experimental; sampler, delay-state updates, staging, and request postprocess remain outside the graph.
|
||||
@@ -33,6 +33,7 @@ recipes/
|
||||
| [`fishaudio/Fish-Speech-S2-Pro.md`](./fishaudio/Fish-Speech-S2-Pro.md) | Online serving for TTS | 1x A800 80GB |
|
||||
| [`Helios/Helios.md`](./Helios/Helios.md) | Text-to-video, image-to-video, and video-to-video generation | 1x NVIDIA H20 |
|
||||
| [`inclusionAI/Ming-flash-omni-2.0.md`](./inclusionAI/Ming-flash-omni-2.0.md) | Online serving for multimodal chat + standalone TTS | 4x H100 / 1x H100 80GB |
|
||||
| [`inclusionAI/Ming-omni-tts-0.5B.md`](./inclusionAI/Ming-omni-tts-0.5B.md) | Offline + online dense Ming TTS/audio generation | 1x AMD `gfx942` (ROCm 7.2) |
|
||||
| [`LTX/LTX-2.md`](./LTX/LTX-2.md) | Text-to-video and image-to-video serving | 1x H200 141GB |
|
||||
| [`LTX/LTX-2.3.md`](./LTX/LTX-2.3.md) | Text-to-video with audio generation (22B) | 1x GPU (96GB VRAM) |
|
||||
| [`mistralai/Voxtral-TTS.md`](./mistralai/Voxtral-TTS.md) | Online serving for TTS | 1x RTX 4090 24GB |
|
||||
|
||||
@@ -171,6 +171,34 @@ python examples/offline_inference/speech_to_video/speech_to_video.py \
|
||||
```
|
||||
|
||||
|
||||
### Online Serving
|
||||
|
||||
#### Server
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size 2 \
|
||||
--flow-shift 3.0 \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
#### Client
|
||||
|
||||
```bash
|
||||
no_proxy=127.0.0.1 \
|
||||
curl -X POST http://127.0.0.1:8091/v1/videos/sync \
|
||||
-F "prompt=A person singing" \
|
||||
-F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
|
||||
-F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
|
||||
-F "width=832" -F "height=480" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.5" \
|
||||
-F "fps=16" \
|
||||
--output s2v_480p_serve.mp4
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The output video should show a talking/singing person matching the reference
|
||||
|
||||
@@ -0,0 +1,82 @@
|
||||
# Ming-omni-tts 0.5B
|
||||
|
||||
> Offline and online dense Ming text-to-speech on ROCm
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: inclusionAI
|
||||
- Model: `inclusionAI/Ming-omni-tts-0.5B`
|
||||
- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
|
||||
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
|
||||
- Maintainer: Community
|
||||
|
||||
## References
|
||||
|
||||
- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
|
||||
- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts)
|
||||
- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
|
||||
- [Online example](../../examples/online_serving/text_to_speech/ming_tts/)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## ROCm
|
||||
|
||||
### 1x AMD `gfx942`
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04.5 LTS, x86_64
|
||||
- Python: 3.12.13
|
||||
- ROCm / HIP: 7.2.53211
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
|
||||
- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
|
||||
|
||||
#### Command
|
||||
|
||||
From the vLLM-Omni repository root:
|
||||
|
||||
```bash
|
||||
docker run --rm \
|
||||
--group-add=video \
|
||||
--ipc=host \
|
||||
--cap-add=SYS_PTRACE \
|
||||
--security-opt seccomp=unconfined \
|
||||
--device /dev/kfd \
|
||||
--device /dev/dri \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
-v "$PWD":/app/vllm-omni \
|
||||
-w /app/vllm-omni \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.22.0 \
|
||||
--model inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' \
|
||||
--ref-audio /path/to/yue_prompt.wav \
|
||||
--max-new-tokens 200 \
|
||||
--output dialect.wav
|
||||
```
|
||||
|
||||
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
|
||||
for the dialect case.
|
||||
|
||||
## Notes
|
||||
|
||||
- The official ROCm image includes the platform dependencies.
|
||||
- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md)
|
||||
for interactive and source-build workflows.
|
||||
- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
|
||||
can be run from an interactive container.
|
||||
- The tested environment uses `--enforce-eager`.
|
||||
- Output is mono 44.1 kHz audio.
|
||||
@@ -22,4 +22,3 @@ prettytable>=3.8.0
|
||||
aenum==3.1.16
|
||||
pyzmq>=25.0.0
|
||||
janus>=1.0.0
|
||||
pydub
|
||||
@@ -1,3 +1,3 @@
|
||||
-r common.txt
|
||||
onnxruntime-cann>=1.23.2
|
||||
torchaudio==2.9.0
|
||||
torchaudio==2.10.0
|
||||
@@ -1,2 +1,3 @@
|
||||
-r common.txt
|
||||
onnxruntime>=1.23.2
|
||||
auto-round-lib
|
||||
@@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreEventType
|
||||
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
|
||||
# isort: on
|
||||
|
||||
@@ -142,3 +143,300 @@ class TestReplaceSessionWithStreamingUpdate:
|
||||
|
||||
assert session.events
|
||||
assert session.events[-1].type == EngineCoreEventType.QUEUED
|
||||
|
||||
|
||||
class _RealignSchedulerStub(OmniSchedulerMixin):
|
||||
"""Minimal scheduler surface for ``_realign_request_status_to_queues``.
|
||||
|
||||
Pins ``self.requests`` / ``self.running`` / ``self.waiting`` to the
|
||||
same shapes the upstream ``Scheduler.finish_requests`` reads, so the
|
||||
helper sees realistic state without spinning a real scheduler.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
requests: dict,
|
||||
running: list,
|
||||
waiting: list,
|
||||
) -> None:
|
||||
self.requests = requests
|
||||
self.running = running
|
||||
self.waiting = waiting
|
||||
|
||||
|
||||
class TestRealignRequestStatusToQueues:
|
||||
"""Regression for the residual hang described in
|
||||
https://github.com/vllm-project/vllm-omni/pull/3774 -- chunk-transfer
|
||||
adapter's ``requests_origin_status`` table goes stale on the
|
||||
``waiting → running`` admit transition, and an abort that lands
|
||||
before the next deque round-trip stomps stale ``WAITING`` onto a
|
||||
request that actually lives in ``self.running``. After
|
||||
``max_num_seqs`` such aborts every ``input_batch`` slot is leaked and
|
||||
new requests hang at ``chunks=0``.
|
||||
"""
|
||||
|
||||
def test_running_with_stale_waiting_status_is_realigned_to_running(
|
||||
self,
|
||||
) -> None:
|
||||
"""The exact race the hang reproduces: a request lives in
|
||||
``self.running`` but its ``status`` is ``WAITING``. After
|
||||
realign, status must be ``RUNNING`` so upstream
|
||||
``Scheduler.finish_requests`` removes it from ``self.running``.
|
||||
"""
|
||||
req = _make_request(request_id="req-stale")
|
||||
req.status = RequestStatus.WAITING # stale -- actually in running
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[req],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.RUNNING
|
||||
|
||||
def test_waiting_with_stale_running_status_is_realigned_to_waiting(
|
||||
self,
|
||||
) -> None:
|
||||
"""Symmetric case: request lives in ``self.waiting`` but status
|
||||
is ``RUNNING``. Upstream's else branch should run, so realign to
|
||||
``WAITING``.
|
||||
"""
|
||||
req = _make_request(request_id="req-stale-2")
|
||||
req.status = RequestStatus.RUNNING # stale -- actually in waiting
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[],
|
||||
waiting=[req],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.WAITING
|
||||
|
||||
def test_already_aligned_status_is_left_unchanged(self) -> None:
|
||||
"""Healthy case: status matches actual queue. No status mutation
|
||||
means no spurious side effects.
|
||||
"""
|
||||
req_running = _make_request(request_id="req-r")
|
||||
req_running.status = RequestStatus.RUNNING
|
||||
req_waiting = _make_request(request_id="req-w")
|
||||
req_waiting.status = RequestStatus.WAITING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
req_running.request_id: req_running,
|
||||
req_waiting.request_id: req_waiting,
|
||||
},
|
||||
running=[req_running],
|
||||
waiting=[req_waiting],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req_running.request_id, req_waiting.request_id])
|
||||
|
||||
assert req_running.status == RequestStatus.RUNNING
|
||||
assert req_waiting.status == RequestStatus.WAITING
|
||||
|
||||
def test_unknown_request_id_is_skipped_silently(self) -> None:
|
||||
"""Aligning ids that aren't in ``self.requests`` (already freed
|
||||
upstream, or never existed) must be a no-op.
|
||||
"""
|
||||
sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
|
||||
sched._realign_request_status_to_queues(["never-existed"]) # no raise
|
||||
|
||||
def test_request_in_neither_queue_is_left_unchanged(self) -> None:
|
||||
"""If a tracked request is in neither queue (e.g. parked in a
|
||||
chunk-transfer deque), realign must not invent a status. The
|
||||
adapter / deque purge owns that surface; the realign helper is
|
||||
only for the admit-transition staleness.
|
||||
"""
|
||||
req = _make_request(request_id="req-parked")
|
||||
req.status = RequestStatus.WAITING_FOR_CHUNK
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
|
||||
def test_request_ids_str_is_treated_as_single_id(self) -> None:
|
||||
"""Match upstream ``Scheduler.finish_requests`` resolution: a
|
||||
bare string is treated as one id, not iterated as characters.
|
||||
"""
|
||||
req = _make_request(request_id="req-s")
|
||||
req.status = RequestStatus.WAITING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[req],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues(req.request_id)
|
||||
|
||||
assert req.status == RequestStatus.RUNNING
|
||||
|
||||
def test_request_ids_none_aligns_every_known_request(self) -> None:
|
||||
"""``request_ids=None`` matches upstream's "all requests" path.
|
||||
Realign must walk every entry in ``self.requests`` and fix any
|
||||
stale status it finds.
|
||||
"""
|
||||
stale = _make_request(request_id="req-stale-none")
|
||||
stale.status = RequestStatus.WAITING # but actually in running
|
||||
|
||||
clean = _make_request(request_id="req-clean-none")
|
||||
clean.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
stale.request_id: stale,
|
||||
clean.request_id: clean,
|
||||
},
|
||||
running=[stale, clean],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues(None)
|
||||
|
||||
assert stale.status == RequestStatus.RUNNING
|
||||
assert clean.status == RequestStatus.RUNNING
|
||||
|
||||
def test_finished_request_is_skipped(self) -> None:
|
||||
"""Already-finished requests must not be touched -- they may
|
||||
have legitimate finished statuses (FINISHED_STOPPED etc.) that
|
||||
a status flip would corrupt.
|
||||
"""
|
||||
req = _make_request(request_id="req-finished")
|
||||
req.status = RequestStatus.FINISHED_STOPPED
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[req], # not realistic, but we want to prove the guard fires
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.FINISHED_STOPPED
|
||||
|
||||
|
||||
class TestPurgeFinishedFromRunning:
|
||||
"""Regression for the residual ``self.running`` slot leak surface
|
||||
paired with ``_realign_request_status_to_queues``: even after
|
||||
realign + ``super().finish_requests`` runs, corner cases can leave
|
||||
already-finished or untracked entries in ``self.running`` -- e.g.
|
||||
a connector cleanup that pops from ``self.requests`` without
|
||||
unwinding ``self.running``, or a ``status`` mid-transition when
|
||||
finish ran. The defensive post-finish purge sweeps those residues
|
||||
so the worker's ``input_batch`` slot never pins a freed request.
|
||||
|
||||
See https://github.com/vllm-project/vllm-omni/pull/3774 discussion
|
||||
and the residual-hang reproduction in that PR.
|
||||
"""
|
||||
|
||||
def test_finished_request_is_purged_from_running(self) -> None:
|
||||
"""``is_finished()`` request lingering in ``self.running`` must
|
||||
be swept so its ``input_batch`` slot is freed."""
|
||||
finished = _make_request(request_id="req-finished")
|
||||
finished.status = RequestStatus.FINISHED_STOPPED # is_finished() True
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={finished.request_id: finished},
|
||||
running=[finished],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == []
|
||||
|
||||
def test_untracked_request_is_purged_from_running(self) -> None:
|
||||
"""Request lingering in ``self.running`` but no longer present
|
||||
in ``self.requests`` (already freed by upstream / connector)
|
||||
must be swept."""
|
||||
untracked = _make_request(request_id="req-untracked")
|
||||
untracked.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={}, # already deleted from self.requests
|
||||
running=[untracked],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == []
|
||||
|
||||
def test_healthy_running_is_left_unchanged(self) -> None:
|
||||
"""Live tracked running requests must not be touched -- the
|
||||
purge is defensive, not aggressive."""
|
||||
alive = _make_request(request_id="req-alive")
|
||||
alive.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={alive.request_id: alive},
|
||||
running=[alive],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == [alive]
|
||||
|
||||
def test_empty_running_is_noop(self) -> None:
|
||||
"""Empty ``self.running`` must short-circuit cleanly."""
|
||||
sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == []
|
||||
|
||||
def test_mixed_alive_and_dead_keeps_only_alive_in_order(self) -> None:
|
||||
"""Mixed ``self.running`` -- some alive, some finished, some
|
||||
untracked. Sweep keeps only the alive ones and preserves their
|
||||
relative order."""
|
||||
alive_a = _make_request(request_id="req-alive-a")
|
||||
alive_a.status = RequestStatus.RUNNING
|
||||
finished_b = _make_request(request_id="req-finished-b")
|
||||
finished_b.status = RequestStatus.FINISHED_STOPPED
|
||||
untracked_c = _make_request(request_id="req-untracked-c")
|
||||
untracked_c.status = RequestStatus.RUNNING
|
||||
alive_d = _make_request(request_id="req-alive-d")
|
||||
alive_d.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
alive_a.request_id: alive_a,
|
||||
finished_b.request_id: finished_b, # tracked but is_finished
|
||||
# untracked_c is NOT in self.requests
|
||||
alive_d.request_id: alive_d,
|
||||
},
|
||||
running=[alive_a, finished_b, untracked_c, alive_d],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == [alive_a, alive_d]
|
||||
|
||||
def test_in_place_mutation_keeps_helper_local_list_identity(self) -> None:
|
||||
"""The slice-assign form (``self.running[:] = ...``) avoids an
|
||||
extra rebind inside this helper. Note that across a full
|
||||
``finish_requests`` call upstream ``Scheduler.finish_requests``
|
||||
already rebinds ``self.running``, so global identity is not a
|
||||
contract -- this test only pins the local helper behavior so a
|
||||
future refactor away from in-place mutation is intentional.
|
||||
"""
|
||||
finished = _make_request(request_id="req-finished")
|
||||
finished.status = RequestStatus.FINISHED_STOPPED
|
||||
alive = _make_request(request_id="req-alive")
|
||||
alive.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
finished.request_id: finished,
|
||||
alive.request_id: alive,
|
||||
},
|
||||
running=[finished, alive],
|
||||
waiting=[],
|
||||
)
|
||||
held_ref = sched.running
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert held_ref is sched.running
|
||||
assert held_ref == [alive]
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Integration tests for the post-finish defensive purge of
|
||||
``self.running`` placed in
|
||||
``OmniARScheduler.finish_requests`` /
|
||||
``OmniGenerationScheduler.finish_requests``.
|
||||
|
||||
Pins *placement* — that ``_purge_finished_from_running`` runs AFTER
|
||||
``super().finish_requests`` so it catches entries upstream's removal
|
||||
missed (e.g. status mid-transition, connector cleanup popping
|
||||
``self.requests`` without unwinding ``self.running``). Direct helper
|
||||
unit tests live in ``test_omni_scheduler_mixin.py``; these tests
|
||||
exercise the full ``finish_requests`` path end-to-end.
|
||||
|
||||
Regression for the residual self.running slot leak surface paired
|
||||
with ``_realign_request_status_to_queues`` (vllm-omni#3774).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
|
||||
import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _StubRequest:
|
||||
"""Minimal Request stub with the surface the helpers exercise."""
|
||||
|
||||
def __init__(self, request_id: str, status: RequestStatus) -> None:
|
||||
self.request_id = request_id
|
||||
self.status = status
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return RequestStatus.is_finished(self.status)
|
||||
|
||||
|
||||
_SCHEDULER_PARAMS = [
|
||||
pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
|
||||
pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
|
||||
]
|
||||
|
||||
|
||||
def _make_scheduler(scheduler_cls, *, requests, running, waiting):
|
||||
"""Bare scheduler with just the surface ``finish_requests`` reads."""
|
||||
scheduler = scheduler_cls.__new__(scheduler_cls)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.requests = requests
|
||||
scheduler.running = running
|
||||
scheduler.waiting = waiting
|
||||
return scheduler
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_purges_finished_request_from_running(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""A request whose ``status`` is a finished value but still lives
|
||||
in ``self.running`` after ``super().finish_requests`` returns must
|
||||
be swept by the post-finish purge. Mocks ``super()`` to simulate
|
||||
the corner case where upstream's status-driven removal missed the
|
||||
entry.
|
||||
"""
|
||||
finished = _StubRequest("req-finished", RequestStatus.FINISHED_STOPPED)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={finished.request_id: finished},
|
||||
running=[finished],
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
# super() returning the abort tuple but leaving self.running untouched
|
||||
return [(finished.request_id, 0)]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
result = scheduler_cls.finish_requests(
|
||||
scheduler,
|
||||
[finished.request_id],
|
||||
RequestStatus.FINISHED_STOPPED,
|
||||
)
|
||||
|
||||
assert scheduler.running == []
|
||||
assert result == [(finished.request_id, 0)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_purges_untracked_request_from_running(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""A request that has been deleted from ``self.requests`` (e.g.
|
||||
after upstream's ``_free_request`` ran but a status mismatch
|
||||
caused the ``self.running`` removal branch to no-op) must be
|
||||
swept by the post-finish purge.
|
||||
"""
|
||||
untracked = _StubRequest("req-untracked", RequestStatus.RUNNING)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={}, # already deleted from self.requests
|
||||
running=[untracked],
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
return [(untracked.request_id, 0)]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
result = scheduler_cls.finish_requests(
|
||||
scheduler,
|
||||
[untracked.request_id],
|
||||
RequestStatus.FINISHED_STOPPED,
|
||||
)
|
||||
|
||||
assert scheduler.running == []
|
||||
assert result == [(untracked.request_id, 0)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_leaves_healthy_running_intact(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""A live ``RUNNING`` request that is not in the finish set must
|
||||
not be swept by the defensive purge -- this pins that the predicate
|
||||
is conservative and won't aggressively drop healthy entries.
|
||||
"""
|
||||
alive = _StubRequest("req-alive", RequestStatus.RUNNING)
|
||||
leaving_id = "req-leaving"
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={alive.request_id: alive},
|
||||
running=[alive], # leaving already removed by super() (typical happy path)
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
return [(leaving_id, 0)]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
scheduler_cls.finish_requests(
|
||||
scheduler,
|
||||
[leaving_id],
|
||||
RequestStatus.FINISHED_STOPPED,
|
||||
)
|
||||
|
||||
assert scheduler.running == [alive]
|
||||
@@ -46,6 +46,9 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
|
||||
scheduler = scheduler_cls.__new__(scheduler_cls)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = coordinator
|
||||
scheduler.requests = {}
|
||||
scheduler.running = []
|
||||
scheduler.waiting = []
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
assert request_ids == ["req-a", "req-b"]
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_higgs_audio_v3",
|
||||
"server_params": {
|
||||
"model": "bosonai/higgs-audio-v3-tts-4b",
|
||||
"stage_overrides": {
|
||||
"0": {
|
||||
"gpu_memory_utilization": 0.3
|
||||
},
|
||||
"1": {
|
||||
"gpu_memory_utilization": 0.2
|
||||
}
|
||||
},
|
||||
"extra_cli_args": [
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/higgs_multimodal_qwen3.yaml",
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "latency",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [
|
||||
20
|
||||
],
|
||||
"max_concurrency": [
|
||||
1
|
||||
],
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
3000
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.5
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "throughput",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [
|
||||
500
|
||||
],
|
||||
"max_concurrency": [
|
||||
16
|
||||
],
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
5000
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
1.2
|
||||
],
|
||||
"audio_throughput": [
|
||||
30
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -339,5 +339,66 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_voxtral_tts",
|
||||
"server_params": {
|
||||
"model": "mistralai/Voxtral-4B-TTS-2603"
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "latency",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [20],
|
||||
"max_concurrency": [1],
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {"voice": "casual_male"},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [150],
|
||||
"median_audio_rtf": [0.15]
|
||||
}
|
||||
},
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "throughput",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [80],
|
||||
"max_concurrency": [8],
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {"voice": "casual_male"},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [1500],
|
||||
"median_audio_rtf": [0.30],
|
||||
"audio_throughput": [30.0]
|
||||
}
|
||||
},
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "stress",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [100],
|
||||
"request_rate": [2.0],
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {"voice": "casual_male"},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [3000],
|
||||
"median_audio_rtf": [0.25],
|
||||
"audio_throughput": [4.0]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -198,19 +198,78 @@ def test_video_chat_stream_invalid_requests(
|
||||
# WS /v1/realtime
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
|
||||
_REALTIME_WS_MODEL_MISMATCH = object()
|
||||
_REALTIME_WS_INVALID_AUDIO_APPEND = object()
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize(
|
||||
"send_frames_spec, ws_error_code, err_message",
|
||||
[
|
||||
pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
|
||||
pytest.param(
|
||||
json.dumps({"type": "session.update"}),
|
||||
"invalid_event",
|
||||
("model", "Missing required field"),
|
||||
id="session_update_missing_model",
|
||||
),
|
||||
pytest.param(
|
||||
_REALTIME_WS_MODEL_MISMATCH,
|
||||
"model_not_found",
|
||||
("model", "does not exist"),
|
||||
id="session_update_model_mismatch",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "input_audio_buffer.commit", "final": False}),
|
||||
"model_not_validated",
|
||||
("session.update", "validate the model"),
|
||||
id="commit_without_model_validation",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "unknown.event"}),
|
||||
"unknown_event",
|
||||
"Unknown event type",
|
||||
id="unknown_event_type",
|
||||
),
|
||||
pytest.param(
|
||||
_REALTIME_WS_INVALID_AUDIO_APPEND,
|
||||
"invalid_audio",
|
||||
"Invalid audio data",
|
||||
id="invalid_audio_append",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||
def test_realtime_rejected_when_async_chunk_enabled(
|
||||
def test_realtime_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
send_frames_spec: Any,
|
||||
ws_error_code: str,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
|
||||
send_frames = json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"model": "this-model-is-not-served-on-this-instance",
|
||||
}
|
||||
)
|
||||
elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
|
||||
send_frames = [
|
||||
json.dumps({"type": "session.update", "model": omni_server.model}),
|
||||
json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
|
||||
]
|
||||
else:
|
||||
send_frames = send_frames_spec
|
||||
openai_client.send_realtime_ws_request(
|
||||
{
|
||||
"send_frames": send_frames,
|
||||
"timeout": 120,
|
||||
"ws_max_size": None,
|
||||
"ws_skip_types": ["session.created"],
|
||||
"ws_json_type": "error",
|
||||
"ws_error_code": "unsupported",
|
||||
"err_message": "async_chunk",
|
||||
"ws_error_code": ws_error_code,
|
||||
"err_message": err_message,
|
||||
}
|
||||
)
|
||||
@@ -319,6 +319,7 @@ QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLO
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
|
||||
@@ -387,6 +388,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
|
||||
@@ -429,6 +431,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
|
||||
@@ -771,7 +774,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="OOM injection causes the serving instance to exit")
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
|
||||
|
||||
@@ -131,6 +131,38 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
|
||||
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
||||
|
||||
|
||||
def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
||||
|
||||
class RecordingGather(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
def forward(self, hidden_gen: torch.Tensor) -> torch.Tensor:
|
||||
self.calls += 1
|
||||
return hidden_gen
|
||||
|
||||
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
||||
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
||||
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
||||
)
|
||||
gather = RecordingGather()
|
||||
model.gen_sp_gather = gather
|
||||
|
||||
output = model(
|
||||
hidden_states=torch.zeros(1, 2, 1, 2, 2),
|
||||
timestep=torch.tensor([1.0]),
|
||||
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
||||
text_mask=torch.ones(1, 2, dtype=torch.long),
|
||||
video_shape=(1, 2, 2),
|
||||
fps=24.0,
|
||||
)
|
||||
|
||||
assert gather.calls == 1
|
||||
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
||||
|
||||
|
||||
def test_sound_and_action_modules_follow_config() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
||||
|
||||
|
||||
@@ -0,0 +1,728 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for multi-resolution support in HunyuanImage3.
|
||||
|
||||
Covers three areas:
|
||||
|
||||
1. ``ResolutionGroup._calc_by_step`` — training-bucket generation,
|
||||
verifying that each resolution has correct width×height and that
|
||||
indices match expected ratio ordering.
|
||||
|
||||
2. ``HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS`` — extra resolution matching
|
||||
via ``Resolution.append`` and new-entry insertion in
|
||||
``ResolutionGroup.__init__``.
|
||||
|
||||
3. ``get_cached_resolution_group`` and ``ar2diffusion`` — caching
|
||||
behaviour and correct ``ResolutionGroup`` construction for the
|
||||
AR→DiT stage transition.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
|
||||
HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS,
|
||||
Resolution,
|
||||
ResolutionGroup,
|
||||
get_cached_resolution_group,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_reso(desc: str | tuple[int, int]) -> Resolution:
|
||||
"""Build a ``Resolution`` from a ``"HxW"`` string or ``(h, w)`` tuple."""
|
||||
if isinstance(desc, str):
|
||||
h, w = desc.split("x")
|
||||
return Resolution(int(h), int(w))
|
||||
return Resolution(desc[0], desc[1])
|
||||
|
||||
|
||||
def _reso_set(rg: ResolutionGroup) -> set[tuple[int, int]]:
|
||||
"""Return all ``(h, w)`` tuples reachable through a ``ResolutionGroup``.
|
||||
|
||||
For entries that carry ``extra_res`` this includes every alternative
|
||||
resolution in that set.
|
||||
"""
|
||||
out: set[tuple[int, int]] = set()
|
||||
for r in rg.data:
|
||||
out.add((r.h, r.w))
|
||||
out |= r.extra_res
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. _calc_by_step — training-bucket resolution parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BASE = 1024
|
||||
_STEP = _BASE // 16 # 64
|
||||
_ALIGN = 16
|
||||
_MIN_SIDE = _BASE // 2 # 512
|
||||
_MAX_SIDE = _BASE * 2 # 2048
|
||||
|
||||
|
||||
def _build_plain_group(
|
||||
base_size: int = _BASE,
|
||||
step: int | None = None,
|
||||
align: int = _ALIGN,
|
||||
) -> ResolutionGroup:
|
||||
"""Return a ``ResolutionGroup`` with *no* extra resolutions."""
|
||||
return ResolutionGroup(base_size=base_size, step=step, align=align, extra_resolutions=None)
|
||||
|
||||
|
||||
class TestCalcByStep:
|
||||
"""Verify ``_calc_by_step`` generates the expected training buckets."""
|
||||
|
||||
# -- count & range -------------------------------------------------------
|
||||
|
||||
def test_count_for_base_1024(self):
|
||||
"""33 buckets: 1 centre + 2×16 from the stepping loops."""
|
||||
rg = _build_plain_group()
|
||||
# 1 centre + 16 portrait + 16 landscape
|
||||
assert len(rg.data) == 33, f"expected 33 training buckets, got {len(rg.data)}"
|
||||
|
||||
def test_all_sizes_within_bounds(self):
|
||||
rg = _build_plain_group()
|
||||
for r in rg.data:
|
||||
assert _MIN_SIDE <= r.h <= _MAX_SIDE, f"height {r.h} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
|
||||
assert _MIN_SIDE <= r.w <= _MAX_SIDE, f"width {r.w} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
|
||||
|
||||
def test_all_sizes_aligned(self):
|
||||
rg = _build_plain_group()
|
||||
for r in rg.data:
|
||||
assert r.h % _ALIGN == 0, f"height {r.h} not {_ALIGN}-aligned"
|
||||
assert r.w % _ALIGN == 0, f"width {r.w} not {_ALIGN}-aligned"
|
||||
|
||||
# -- ratio ordering ------------------------------------------------------
|
||||
|
||||
def test_sorted_by_ratio_ascending(self):
|
||||
rg = _build_plain_group()
|
||||
ratios = [r.ratio for r in rg.data]
|
||||
for i in range(1, len(ratios)):
|
||||
assert ratios[i - 1] <= ratios[i], f"ratio[{i - 1}]={ratios[i - 1]:.4f} > ratio[{i}]={ratios[i]:.4f}"
|
||||
|
||||
def test_ratio_index_matches_numpy_array(self):
|
||||
rg = _build_plain_group()
|
||||
for i, r in enumerate(rg.data):
|
||||
assert math.isclose(rg.ratio[i], r.ratio, rel_tol=1e-9), (
|
||||
f"rg.ratio[{i}]={rg.ratio[i]:.6f} != data[{i}].ratio={r.ratio:.6f}"
|
||||
)
|
||||
|
||||
# -- key index positions -------------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"idx,expected_hw",
|
||||
[
|
||||
# First (most portrait) → 1:4 ratio bucket
|
||||
(0, (512, 2048)),
|
||||
# Centre square bucket
|
||||
(16, (1024, 1024)),
|
||||
# Most landscape bucket → 4:1 ratio
|
||||
(32, (2048, 512)),
|
||||
],
|
||||
)
|
||||
def test_landmark_index(self, idx, expected_hw):
|
||||
rg = _build_plain_group()
|
||||
r = rg.data[idx]
|
||||
assert (r.h, r.w) == expected_hw, f"idx {idx}: expected {expected_hw}, got ({r.h}, {r.w})"
|
||||
|
||||
def test_center_is_square(self):
|
||||
rg = _build_plain_group()
|
||||
r = rg.data[16]
|
||||
assert r.h == r.w == _BASE
|
||||
|
||||
def test_first_and_last_are_extremes(self):
|
||||
rg = _build_plain_group()
|
||||
# first → 512×2048 (ratio 0.25)
|
||||
assert rg.data[0].ratio < 0.3
|
||||
# last → 2048×512 (ratio 4.0)
|
||||
assert rg.data[-1].ratio > 3.9
|
||||
|
||||
# -- width × height matching ---------------------------------------------
|
||||
|
||||
def test_every_resolution_has_same_area_product(self):
|
||||
"""Symmetry: stepping loop produces the same product set in both
|
||||
directions, so (h × w) tuples should appear symmetrically."""
|
||||
rg = _build_plain_group()
|
||||
products = {(r.h, r.w) for r in rg.data}
|
||||
# For every (h, w) there should exist (w, h) or a nearby equivalent
|
||||
for h, w in products:
|
||||
if h == w:
|
||||
continue # square — trivially symmetric
|
||||
assert (w, h) in products or any(abs(rr.h - w) <= _STEP and abs(rr.w - h) <= _STEP for rr in rg.data), (
|
||||
f"no symmetric counterpart for {h}×{w}"
|
||||
)
|
||||
|
||||
def test_product_of_wh_divisible_by_256(self):
|
||||
"""Each bucket's area should be a multiple of 256 (16×16),
|
||||
which is the VAE down-sample grid."""
|
||||
rg = _build_plain_group()
|
||||
for r in rg.data:
|
||||
assert (r.h * r.w) % 256 == 0, f"{r.h}×{r.w} = {r.h * r.w} not divisible by 256"
|
||||
|
||||
# -- custom base_size / step / align -------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("base_size", [512, 768, 1280])
|
||||
def test_count_scales_with_base_size(self, base_size):
|
||||
step = base_size // 16
|
||||
rg = ResolutionGroup(base_size=base_size, step=step, align=16, extra_resolutions=None)
|
||||
expected = 1 + 2 * (base_size // step) # centre + 2×phases
|
||||
assert len(rg.data) == expected, f"base_size={base_size}: expected {expected}, got {len(rg.data)}"
|
||||
|
||||
def test_custom_step(self):
|
||||
"""A smaller step produces more buckets."""
|
||||
rg_small = ResolutionGroup(base_size=1024, step=32, align=16, extra_resolutions=None)
|
||||
rg_default = _build_plain_group()
|
||||
assert len(rg_small.data) > len(rg_default.data), "smaller step should produce more buckets"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS matching
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_EXTRA_SQUARES = {"512x512", "640x640", "768x768", "896x896"}
|
||||
_EXTRA_NON_SQUARES = {"1024x768", "1280x720", "768x1024", "720x1280"}
|
||||
|
||||
|
||||
def _build_extra_group() -> ResolutionGroup:
|
||||
return ResolutionGroup(
|
||||
base_size=_BASE,
|
||||
extra_resolutions=[_make_reso(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS],
|
||||
)
|
||||
|
||||
|
||||
class TestExtraResolutionMatching:
|
||||
"""Verify extra resolutions are correctly integrated into
|
||||
the ``ResolutionGroup``."""
|
||||
|
||||
# -- square extras (ratio = 1.0) -----------------------------------------
|
||||
|
||||
def test_square_extras_appended_to_base_square(self):
|
||||
"""512², 640², 768², 896² all have ratio 1.0 and should be
|
||||
appended to the (1024, 1024) Resolution via ``append()``."""
|
||||
rg = _build_extra_group()
|
||||
# Find the 1024×1024 entry
|
||||
square_reso = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
|
||||
expected_extra = {(512, 512), (640, 640), (768, 768), (896, 896)}
|
||||
assert square_reso.extra_res == expected_extra, f"square extra_res: {square_reso.extra_res}"
|
||||
|
||||
def test_base_1024_unchanged(self):
|
||||
"""The 1024×1024 entry itself still reports h=1024, w=1024."""
|
||||
rg = _build_extra_group()
|
||||
r = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
|
||||
assert r.h == 1024 and r.w == 1024
|
||||
|
||||
def test_squares_not_standalone_entries(self):
|
||||
"""Square extras must NOT appear as top-level ``Resolution`` entries;
|
||||
they live inside ``extra_res`` only."""
|
||||
rg = _build_extra_group()
|
||||
standalone = {(r.h, r.w) for r in rg.data}
|
||||
for square in [(512, 512), (640, 640), (768, 768), (896, 896)]:
|
||||
assert square not in standalone, f"{square} should be in extra_res, not a standalone entry"
|
||||
|
||||
# -- non-square extras (ratio ≠ 1.0) ------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("reso_str", sorted(_EXTRA_NON_SQUARES))
|
||||
def test_non_square_extra_is_standalone_entry(self, reso_str):
|
||||
"""1024×768, 1280×720, 768×1024, 720×1280 are appended as
|
||||
new ``Resolution`` entries because no training bucket shares
|
||||
their exact ratio."""
|
||||
r = _make_reso(reso_str)
|
||||
rg = _build_extra_group()
|
||||
matching = [d for d in rg.data if d.h == r.h and d.w == r.w]
|
||||
assert len(matching) == 1, f"{reso_str} should be in rg.data exactly once"
|
||||
assert matching[0].extra_res == set(), f"{reso_str} should have empty extra_res, got {matching[0].extra_res}"
|
||||
|
||||
def test_non_square_extras_empty_extra_res(self):
|
||||
"""Non-square extras have no own extra_res."""
|
||||
rg = _build_extra_group()
|
||||
for r in rg.data:
|
||||
if (r.h, r.w) in {(1024, 768), (1280, 720), (768, 1024), (720, 1280)}:
|
||||
assert r.extra_res == set()
|
||||
|
||||
# -- total count ---------------------------------------------------------
|
||||
|
||||
def test_total_count(self):
|
||||
"""33 training + 4 non-square extras = 37."""
|
||||
rg = _build_extra_group()
|
||||
assert len(rg.data) == 37, f"expected 37, got {len(rg.data)}"
|
||||
|
||||
def test_extra_count(self):
|
||||
"""Exactly 8 extra resolutions (4 square + 4 non-square)."""
|
||||
rg = _build_extra_group()
|
||||
reso_set = _reso_set(rg)
|
||||
extra_in_set = reso_set - {(r.h, r.w) for r in _build_plain_group().data}
|
||||
assert len(extra_in_set) == 8, f"expected 8 extra, got {len(extra_in_set)}: {extra_in_set}"
|
||||
|
||||
# -- index stability -----------------------------------------------------
|
||||
|
||||
def test_indices_0_32_unchanged(self):
|
||||
"""Indices 0–32 remain the training buckets (sorted by ratio)."""
|
||||
plain = _build_plain_group()
|
||||
rg = _build_extra_group()
|
||||
for i in range(33):
|
||||
assert (rg.data[i].h, rg.data[i].w) == (plain.data[i].h, plain.data[i].w), (
|
||||
f"idx {i}: extra group has ({rg.data[i].h},{rg.data[i].w}), "
|
||||
f"expected ({plain.data[i].h},{plain.data[i].w})"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"reso_str,expected_idx",
|
||||
[
|
||||
("1024x768", 33),
|
||||
("1280x720", 34),
|
||||
("768x1024", 35),
|
||||
("720x1280", 36),
|
||||
],
|
||||
)
|
||||
def test_non_square_extra_indices(self, reso_str, expected_idx):
|
||||
"""Non-square extras are appended in order at indices 33-36."""
|
||||
r = _make_reso(reso_str)
|
||||
rg = _build_extra_group()
|
||||
idx = next(i for i, d in enumerate(rg.data) if d.h == r.h and d.w == r.w)
|
||||
assert idx == expected_idx, f"{reso_str}: expected idx {expected_idx}, got {idx}"
|
||||
|
||||
# -- ratio lookup via get_base_size_and_ratio_index ----------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"h,w,expected_idx",
|
||||
[
|
||||
(1024, 768, 33),
|
||||
(1280, 720, 34),
|
||||
(768, 1024, 35),
|
||||
(720, 1280, 36),
|
||||
(1024, 1024, 16),
|
||||
(512, 512, 16),
|
||||
(640, 640, 16),
|
||||
(768, 768, 16),
|
||||
(896, 896, 16),
|
||||
],
|
||||
)
|
||||
def test_ratio_index_for_resolutions(self, h, w, expected_idx):
|
||||
"""``get_base_size_and_ratio_index`` must return the correct
|
||||
ratio index for both training-bucket and extra resolutions."""
|
||||
rg = _build_extra_group()
|
||||
_, idx = rg.get_base_size_and_ratio_index(w, h)
|
||||
assert idx == expected_idx, f"({w},{h}): expected idx {expected_idx}, got {idx}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Resolution.match — area-based matching within a ratio group
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResolutionMatch:
|
||||
"""Verify ``Resolution.match`` picks the best-fit alternative from
|
||||
``extra_res`` based on area proximity."""
|
||||
|
||||
def test_no_extra_returns_self(self):
|
||||
r = _make_reso("1024x768")
|
||||
assert r.match(800, 600) == (768, 1024)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"target_w,target_h,expected",
|
||||
[
|
||||
# Exact area match in extra_res
|
||||
(512, 512, (512, 512)),
|
||||
(640, 640, (640, 640)),
|
||||
(768, 768, (768, 768)),
|
||||
(896, 896, (896, 896)),
|
||||
# Closer to 1024² than to any extra square
|
||||
(1024, 1024, (1024, 1024)),
|
||||
(1000, 1000, (1024, 1024)),
|
||||
# Small square picks 512²
|
||||
(256, 256, (512, 512)),
|
||||
(400, 400, (512, 512)),
|
||||
# 700² → closest to 640² (diff 129600) vs 768² (diff 99856)
|
||||
# 768*768 - 700*700 = 589824 - 490000 = 99824
|
||||
# 640*640 - 700*700 = 409600 - 490000 = -80400 → abs=80400
|
||||
# So 640² is closer in area
|
||||
(700, 700, (640, 640)),
|
||||
# 800² → closest to 768² (diff 49984) vs 896² (diff 162944)
|
||||
(800, 800, (768, 768)),
|
||||
],
|
||||
)
|
||||
def test_square_extra_match(self, target_w, target_h, expected):
|
||||
"""Build a Resolution like the 1024² entry with square extras,
|
||||
then verify area-based matching."""
|
||||
r = Resolution(1024, 1024)
|
||||
for s in [(512, 512), (640, 640), (768, 768), (896, 896)]:
|
||||
r.extra_res.add(s)
|
||||
result = r.match(target_w, target_h)
|
||||
assert result == expected, f"match({target_w},{target_h}) = {result}, expected {expected}"
|
||||
|
||||
def test_match_returns_tuple(self):
|
||||
r = Resolution(1024, 1024)
|
||||
r.extra_res.add((512, 512))
|
||||
result = r.match(512, 512)
|
||||
assert isinstance(result, tuple)
|
||||
assert len(result) == 2
|
||||
assert all(isinstance(v, int) for v in result)
|
||||
|
||||
def test_match_does_not_mutate_resolution(self):
|
||||
r = Resolution(1024, 1024)
|
||||
r.extra_res.add((512, 512))
|
||||
extra_before = set(r.extra_res)
|
||||
r.match(300, 300)
|
||||
assert r.extra_res == extra_before
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ResolutionGroup.get_target_size — end-to-end resolution selection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetTargetSize:
|
||||
"""Verify ``get_target_size`` uses ratio-then-area logic correctly."""
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def rg(self) -> ResolutionGroup:
|
||||
return _build_extra_group()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"w,h,expected_min_area,expected_max_area",
|
||||
[
|
||||
# Square targets should land on square variants
|
||||
(1024, 1024, 1024 * 1024, 1024 * 1024),
|
||||
(512, 512, 512 * 512, 512 * 512),
|
||||
(640, 640, 640 * 640, 640 * 640),
|
||||
(768, 768, 768 * 768, 768 * 768),
|
||||
(896, 896, 896 * 896, 896 * 896),
|
||||
],
|
||||
)
|
||||
def test_square_target_returns_square(self, rg, w, h, expected_min_area, expected_max_area):
|
||||
rw, rh = rg.get_target_size(w, h)
|
||||
area = rw * rh
|
||||
assert expected_min_area <= area <= expected_max_area, f"({w},{h}) → ({rw},{rh}) area={area}"
|
||||
|
||||
def test_training_bucket_self_maps(self, rg):
|
||||
"""A training bucket should map to itself (or an extra_res variant
|
||||
with the same ratio)."""
|
||||
for r in rg.data:
|
||||
rw, rh = rg.get_target_size(r.w, r.h)
|
||||
# Either exact match or the ratio must be preserved
|
||||
assert abs(rh / rw - r.ratio) < 0.01, f"({r.w},{r.h}) → ({rw},{rh}) ratio drifted from {r.ratio:.4f}"
|
||||
|
||||
def test_wild_ratio_still_finds_bucket(self, rg):
|
||||
"""Even an unusual aspect ratio should map to some valid bucket."""
|
||||
for w, h in [(800, 600), (1200, 900), (1600, 400), (400, 1600)]:
|
||||
rw, rh = rg.get_target_size(w, h)
|
||||
assert rw > 0 and rh > 0
|
||||
assert rw % 16 == 0 and rh % 16 == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. get_cached_resolution_group & ar2diffusion integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetCachedResolutionGroup:
|
||||
"""Verify ``get_cached_resolution_group`` returns a correctly
|
||||
configured ``ResolutionGroup`` and respects the LRU cache."""
|
||||
|
||||
def test_returns_resolution_group(self):
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
assert isinstance(rg, ResolutionGroup)
|
||||
|
||||
def test_base_size_propagated(self):
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
assert rg.base_size == 1024
|
||||
|
||||
def test_has_extra_resolutions(self):
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
reso_set = _reso_set(rg)
|
||||
for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS:
|
||||
r = _make_reso(s)
|
||||
assert (r.h, r.w) in reso_set, f"{s} missing from cached group"
|
||||
|
||||
def test_cache_same_base_size_same_object(self):
|
||||
"""LRU cache: same base_size → same object identity."""
|
||||
# Clear any prior cache entries
|
||||
get_cached_resolution_group.cache_clear()
|
||||
rg1 = get_cached_resolution_group(base_size=1024)
|
||||
rg2 = get_cached_resolution_group(base_size=1024)
|
||||
assert rg1 is rg2, "same base_size must return cached object"
|
||||
|
||||
def test_cache_different_base_size_different_object(self):
|
||||
get_cached_resolution_group.cache_clear()
|
||||
rg1 = get_cached_resolution_group(base_size=1024)
|
||||
rg2 = get_cached_resolution_group(base_size=768)
|
||||
assert rg1 is not rg2, "different base_size must return different objects"
|
||||
|
||||
def test_cache_size_limit(self):
|
||||
"""Cache is bounded at 4 entries (``maxsize=4``)."""
|
||||
get_cached_resolution_group.cache_clear()
|
||||
for bs in [512, 768, 1024, 1280, 1536]:
|
||||
get_cached_resolution_group(base_size=bs)
|
||||
info = get_cached_resolution_group.cache_info()
|
||||
assert info.currsize <= 4, f"cache grew to {info.currsize}, max is 4"
|
||||
# All 5 unique base_size values caused cache misses
|
||||
assert info.misses >= 5, f"expected >=5 misses for 5 unique base_size values, got {info.misses}"
|
||||
|
||||
def test_resolution_indexing_works(self):
|
||||
"""Indices 0..len(rg)-1 must be accessible."""
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
for i in range(len(rg)):
|
||||
r = rg[i]
|
||||
assert isinstance(r, Resolution)
|
||||
assert r.h > 0 and r.w > 0
|
||||
|
||||
@pytest.mark.parametrize("base_size", [512, 1024])
|
||||
def test_ratio_index_roundtrip_via_group_length(self, base_size):
|
||||
"""Every ratio index in [0, len(rg)) must resolve to a valid
|
||||
Resolution with height/width within bounds."""
|
||||
rg = get_cached_resolution_group(base_size=base_size)
|
||||
for i in range(len(rg)):
|
||||
r = rg[i]
|
||||
assert isinstance(r.height, int) and r.height > 0
|
||||
assert isinstance(r.width, int) and r.width > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ar2diffusion — integration tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Token IDs matching the real prompt_utils values
|
||||
_RATIO_0 = 128044
|
||||
_RATIO_32 = 128076
|
||||
_RATIO_33 = 130103
|
||||
_RATIO_36 = 130106
|
||||
|
||||
|
||||
def _fake_ar_output(
|
||||
*,
|
||||
cumulative_token_ids: list[int] | None = None,
|
||||
text: str = "",
|
||||
) -> MagicMock:
|
||||
"""Minimal stub for a vLLM ``RequestOutput`` with one completion."""
|
||||
choice = MagicMock()
|
||||
choice.cumulative_token_ids = cumulative_token_ids or []
|
||||
choice.text = text
|
||||
output = MagicMock()
|
||||
output.outputs = [choice]
|
||||
return output
|
||||
|
||||
|
||||
class TestAr2DiffusionIntegration:
|
||||
"""Verify that ``ar2diffusion`` calls ``get_cached_resolution_group``
|
||||
and obtains a correctly configured ``ResolutionGroup`` whose indices map
|
||||
to the expected training-bucket or extra-resolution sizes."""
|
||||
|
||||
def _call_ar2diffusion(
|
||||
self,
|
||||
token_ids: list[int],
|
||||
prompt_kwargs: dict[str, Any] | None = None,
|
||||
) -> tuple[Any, Any]:
|
||||
"""Invoke ``ar2diffusion`` with a single fake AR output.
|
||||
|
||||
Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
|
||||
assert on the resolved height/width.
|
||||
"""
|
||||
from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
|
||||
ar2diffusion,
|
||||
)
|
||||
|
||||
pk = dict(prompt_kwargs or {})
|
||||
pk.setdefault("prompt", "a beautiful sunset")
|
||||
pk.setdefault("height", 1024)
|
||||
pk.setdefault("width", 1024)
|
||||
pk.setdefault("image_base_size", 1024)
|
||||
|
||||
# Build a prompt-like object that behaves like the real OmniTokensPrompt
|
||||
class _FakePrompt:
|
||||
def __init__(self, d):
|
||||
self._d = d
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._d[key]
|
||||
|
||||
def get(self, key, default=None):
|
||||
return self._d.get(key, default)
|
||||
|
||||
def _asdict(self):
|
||||
return dict(self._d)
|
||||
|
||||
prompt = _FakePrompt(pk)
|
||||
ar_output = _fake_ar_output(cumulative_token_ids=token_ids)
|
||||
result = ar2diffusion(
|
||||
source_outputs=[ar_output],
|
||||
prompt=[prompt],
|
||||
requires_multimodal_data=False,
|
||||
)
|
||||
return result
|
||||
|
||||
# -- ratio token ↔ resolution resolution ---------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ratio_idx,expected_hw",
|
||||
[
|
||||
# Training buckets
|
||||
(0, (512, 2048)),
|
||||
(16, (1024, 1024)),
|
||||
(32, (2048, 512)),
|
||||
# Non-square extras
|
||||
(33, (1024, 768)),
|
||||
(34, (1280, 720)),
|
||||
(35, (768, 1024)),
|
||||
(36, (720, 1280)),
|
||||
],
|
||||
)
|
||||
def test_ratio_index_resolves_to_correct_size(self, ratio_idx, expected_hw):
|
||||
"""When the AR emits ``<img_ratio_N>``, the DiT input must carry
|
||||
the training-bucket width/height for that index."""
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
reso = rg[ratio_idx]
|
||||
assert (reso.height, reso.width) == expected_hw, (
|
||||
f"ratio_idx {ratio_idx}: expected {expected_hw}, got ({reso.height}, {reso.width})"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("ratio_idx", [0, 16, 32, 33, 34, 35, 36])
|
||||
def test_ar2diffusion_extracts_height_width(self, ratio_idx):
|
||||
"""End-to-end: simulate AR output containing ``<img_ratio_{idx}>``
|
||||
and assert the DiT inputs carry the expected size."""
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
expected_h = rg[ratio_idx].height
|
||||
expected_w = rg[ratio_idx].width
|
||||
|
||||
token_id = _RATIO_0 + ratio_idx if ratio_idx <= 32 else _RATIO_33 + (ratio_idx - 33)
|
||||
result = self._call_ar2diffusion(
|
||||
token_ids=[150000, 150001, token_id],
|
||||
prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
diffusion_input = result[0]
|
||||
assert diffusion_input["height"] == expected_h, (
|
||||
f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
|
||||
)
|
||||
assert diffusion_input["width"] == expected_w, (
|
||||
f"ratio_idx={ratio_idx}: expected width {expected_w}, got {diffusion_input['width']}"
|
||||
)
|
||||
|
||||
def test_no_ratio_token_falls_back_to_prompt_size(self):
|
||||
"""If no ratio token is present, height/width from the prompt
|
||||
are kept unchanged."""
|
||||
result = self._call_ar2diffusion(
|
||||
token_ids=[150000, 150001], # no ratio tokens
|
||||
prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
diffusion_input = result[0]
|
||||
assert diffusion_input["height"] == 768
|
||||
assert diffusion_input["width"] == 1024
|
||||
|
||||
def test_out_of_range_ratio_index_falls_back(self):
|
||||
"""If ratio_idx ≥ len(reso_group), ar2diffusion must keep the
|
||||
prompt size (fallback path)."""
|
||||
with patch(
|
||||
"vllm_omni.model_executor.stage_input_processors.hunyuan_image3._extract_ratio_index",
|
||||
return_value=999, # out of range
|
||||
):
|
||||
result = self._call_ar2diffusion(
|
||||
token_ids=[150000],
|
||||
prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
assert result[0]["height"] == 512
|
||||
assert result[0]["width"] == 512
|
||||
|
||||
# -- cache behaviour -----------------------------------------------------
|
||||
|
||||
def test_ar2diffusion_uses_cached_group(self):
|
||||
"""Each invocation of ``ar2diffusion`` calls
|
||||
``get_cached_resolution_group``, which returns a shared cached object."""
|
||||
get_cached_resolution_group.cache_clear()
|
||||
rg_before = get_cached_resolution_group(base_size=1024)
|
||||
|
||||
token_id = _RATIO_0 + 16 # <img_ratio_16> = 1024×1024
|
||||
self._call_ar2diffusion(
|
||||
token_ids=[token_id],
|
||||
prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
|
||||
rg_after = get_cached_resolution_group(base_size=1024)
|
||||
assert rg_before is rg_after, "get_cached_resolution_group should return the same cached object"
|
||||
|
||||
def test_correct_reso_group_used_for_base_size(self):
|
||||
"""Verify ``ar2diffusion`` constructs or retrieves the
|
||||
ResolutionGroup for the correct ``image_base_size``."""
|
||||
get_cached_resolution_group.cache_clear()
|
||||
|
||||
# First request with a non-default base_size
|
||||
token_id = _RATIO_0 # <img_ratio_0>
|
||||
self._call_ar2diffusion(
|
||||
token_ids=[token_id],
|
||||
prompt_kwargs={"height": 512, "width": 512, "image_base_size": 512},
|
||||
)
|
||||
|
||||
# The cached group for base_size=512 should exist
|
||||
rg512 = get_cached_resolution_group(base_size=512)
|
||||
assert rg512.base_size == 512
|
||||
|
||||
# And base_size=1024 should also be separately accessible
|
||||
rg1024 = get_cached_resolution_group(base_size=1024)
|
||||
assert rg1024.base_size == 1024
|
||||
assert rg512 is not rg1024
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Edge cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
def test_resolution_string_parsing(self):
|
||||
"""Resolution accepts both "HxW" strings and (h, w) tuples."""
|
||||
r1 = Resolution("512x2048")
|
||||
assert r1.h == 512 and r1.w == 2048
|
||||
|
||||
r2 = Resolution(768, 1024)
|
||||
assert r2.h == 768 and r2.w == 1024
|
||||
|
||||
r3 = Resolution(1024) # single int → square
|
||||
assert r3.h == 1024 and r3.w == 1024
|
||||
|
||||
def test_resolution_extra_res_initially_empty(self):
|
||||
r = Resolution(1024, 1024)
|
||||
assert r.extra_res == set()
|
||||
assert r.match(512, 512) == (1024, 1024)
|
||||
|
||||
def test_resolution_append_deduplicates(self):
|
||||
"""Appending the same resolution twice is a no-op on the set."""
|
||||
r = Resolution(1024, 1024)
|
||||
r.append(Resolution(512, 512))
|
||||
r.append(Resolution(512, 512))
|
||||
assert len(r.extra_res) == 1
|
||||
|
||||
def test_all_final_resolutions_are_16_aligned(self):
|
||||
"""Every reachable resolution (including extras) must be
|
||||
divisible by the VAE down-sample factor of 16."""
|
||||
rg = _build_extra_group()
|
||||
for h, w in _reso_set(rg):
|
||||
assert h % 16 == 0, f"height {h} not 16-aligned in {h}×{w}"
|
||||
assert w % 16 == 0, f"width {w} not 16-aligned in {h}×{w}"
|
||||
|
||||
def test_ratio_array_length_matches_data(self):
|
||||
rg = _build_extra_group()
|
||||
assert len(rg.ratio) == len(rg.data)
|
||||
|
||||
def test_resolution_repr(self):
|
||||
r = Resolution(1024, 1024)
|
||||
rep = repr(r)
|
||||
assert "1024x1024" in rep
|
||||
|
||||
r2 = Resolution(1024, 1024)
|
||||
r2.extra_res.add((512, 512))
|
||||
rep2 = repr(r2)
|
||||
assert "512x512" in rep2
|
||||
@@ -97,6 +97,260 @@ class TestPipelineIndependence:
|
||||
assert modules.resident_names == ["vocoder"]
|
||||
assert len(modules.vaes) == 2
|
||||
|
||||
def test_ltx23_pipeline_has_diffusion_pipeline_profiler_mixin(self):
|
||||
"""LTX23Pipeline must support lightweight stage timing."""
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
||||
from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin
|
||||
|
||||
assert issubclass(LTX23Pipeline, DiffusionPipelineProfilerMixin)
|
||||
|
||||
|
||||
class TestLTX23VaeDecodeParallel:
|
||||
"""Test LTX-2.3 video VAE tiled parallel helpers without loading weights."""
|
||||
|
||||
def test_ltx23_video_vae_is_distributed_tile_only_class(self):
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
||||
DistributedAutoencoderKLLTX2Video,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import DistributedVaeMixin
|
||||
|
||||
assert issubclass(DistributedAutoencoderKLLTX2Video, DistributedVaeMixin)
|
||||
assert not hasattr(DistributedAutoencoderKLLTX2Video, "patch_split")
|
||||
|
||||
def test_ltx23_video_vae_tile_split_uses_native_ltx23_tile_geometry(self):
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
||||
DistributedAutoencoderKLLTX2Video,
|
||||
)
|
||||
|
||||
vae = SimpleNamespace(
|
||||
spatial_compression_ratio=32,
|
||||
tile_sample_min_height=512,
|
||||
tile_sample_min_width=512,
|
||||
tile_sample_stride_height=448,
|
||||
tile_sample_stride_width=448,
|
||||
temporal_compression_ratio=8,
|
||||
dtype=torch.float32,
|
||||
)
|
||||
|
||||
z = torch.zeros(1, 2, 5, 16, 24)
|
||||
tasks, grid_spec = DistributedAutoencoderKLLTX2Video.tile_split(vae, z)
|
||||
|
||||
assert grid_spec.grid_shape == (2, 2)
|
||||
assert grid_spec.split_dims == (3, 4)
|
||||
assert grid_spec.tile_spec["sample_height"] == 512
|
||||
assert grid_spec.tile_spec["sample_width"] == 768
|
||||
assert grid_spec.tile_spec["blend_height"] == 64
|
||||
assert grid_spec.tile_spec["blend_width"] == 64
|
||||
assert grid_spec.tile_spec["max_tile_output_shape"] == (1, 3, 33, 512, 512)
|
||||
assert grid_spec.tile_spec["tile_output_shapes"] == {
|
||||
0: (1, 3, 33, 512, 512),
|
||||
1: (1, 3, 33, 512, 320),
|
||||
2: (1, 3, 33, 64, 512),
|
||||
3: (1, 3, 33, 64, 320),
|
||||
}
|
||||
assert [task.grid_coord for task in tasks] == [(0, 0), (0, 1), (1, 0), (1, 1)]
|
||||
assert [tuple(task.tensor.shape) for task in tasks] == [
|
||||
(1, 2, 5, 16, 16),
|
||||
(1, 2, 5, 16, 10),
|
||||
(1, 2, 5, 2, 16),
|
||||
(1, 2, 5, 2, 10),
|
||||
]
|
||||
assert [task.workload for task in tasks] == [5 * 16 * 16, 5 * 16 * 10, 5 * 2 * 16, 5 * 2 * 10]
|
||||
|
||||
def test_ltx23_video_vae_tile_merge_blends_and_crops_like_tiled_decode(self):
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import (
|
||||
DistributedAutoencoderKLLTX2Video,
|
||||
)
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import GridSpec
|
||||
|
||||
class FakeVae:
|
||||
def __init__(self):
|
||||
self.blend_calls = []
|
||||
|
||||
def clear_cache(self):
|
||||
pass
|
||||
|
||||
def blend_v(self, _previous, current, blend_height):
|
||||
self.blend_calls.append(("v", blend_height))
|
||||
return current
|
||||
|
||||
def blend_h(self, _previous, current, blend_width):
|
||||
self.blend_calls.append(("h", blend_width))
|
||||
return current
|
||||
|
||||
fake_vae = FakeVae()
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(3, 4),
|
||||
grid_shape=(2, 2),
|
||||
tile_spec={
|
||||
"sample_height": 10,
|
||||
"sample_width": 10,
|
||||
"blend_height": 1,
|
||||
"blend_width": 2,
|
||||
"tile_sample_stride_height": 5,
|
||||
"tile_sample_stride_width": 5,
|
||||
},
|
||||
)
|
||||
tiles = {
|
||||
(0, 0): torch.full((1, 3, 2, 6, 6), 1.0),
|
||||
(0, 1): torch.full((1, 3, 2, 6, 6), 2.0),
|
||||
(1, 0): torch.full((1, 3, 2, 6, 6), 3.0),
|
||||
(1, 1): torch.full((1, 3, 2, 6, 6), 4.0),
|
||||
}
|
||||
|
||||
merged = DistributedAutoencoderKLLTX2Video.tile_merge(fake_vae, tiles, grid_spec)
|
||||
|
||||
assert merged.shape == (1, 3, 2, 10, 10)
|
||||
assert fake_vae.blend_calls == [("h", 2), ("v", 1), ("v", 1), ("h", 2)]
|
||||
torch.testing.assert_close(merged[:, :, :, :5, :5], torch.ones(1, 3, 2, 5, 5))
|
||||
torch.testing.assert_close(merged[:, :, :, :5, 5:], torch.full((1, 3, 2, 5, 5), 2.0))
|
||||
torch.testing.assert_close(merged[:, :, :, 5:, :5], torch.full((1, 3, 2, 5, 5), 3.0))
|
||||
torch.testing.assert_close(merged[:, :, :, 5:, 5:], torch.full((1, 3, 2, 5, 5), 4.0))
|
||||
|
||||
def test_ltx23_video_vae_tiled_decode_dispatches_to_tile_operator(self):
|
||||
from vllm_omni.diffusion.distributed.autoencoders import autoencoder_kl_ltx2
|
||||
|
||||
z = torch.zeros(1, 2, 1, 16, 24)
|
||||
expected = torch.ones(1, 3, 1, 512, 768)
|
||||
seen = {}
|
||||
|
||||
class FakeExecutor:
|
||||
def execute(self, tensor, operator, broadcast_result=True):
|
||||
seen["tensor"] = tensor
|
||||
seen["operator"] = operator
|
||||
seen["broadcast_result"] = broadcast_result
|
||||
return expected
|
||||
|
||||
vae = SimpleNamespace(distributed_executor=FakeExecutor(), is_distributed_enabled=lambda: True)
|
||||
vae.tile_split = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_split.__get__(vae)
|
||||
vae.tile_exec = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_exec.__get__(vae)
|
||||
vae.tile_merge = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tile_merge.__get__(vae)
|
||||
|
||||
output = autoencoder_kl_ltx2.DistributedAutoencoderKLLTX2Video.tiled_decode(
|
||||
vae,
|
||||
z,
|
||||
temb=torch.tensor(0.5),
|
||||
return_dict=False,
|
||||
)
|
||||
|
||||
assert len(output) == 1
|
||||
assert output[0] is expected
|
||||
assert seen["tensor"] is z
|
||||
assert seen["broadcast_result"] is False
|
||||
assert seen["operator"].split.__name__ == "tile_split"
|
||||
assert seen["operator"].merge.__name__ == "tile_merge"
|
||||
|
||||
def test_ltx23_vae_executor_gathers_known_tile_shapes_and_returns_empty_on_non_rank0(self):
|
||||
from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_ltx2 import LTX2VaeExecutor
|
||||
from vllm_omni.diffusion.distributed.autoencoders.distributed_vae_executor import (
|
||||
DistributedOperator,
|
||||
GridSpec,
|
||||
TileTask,
|
||||
)
|
||||
|
||||
z = torch.zeros(1, 1, 1, 1, 1)
|
||||
tile_output_shapes = {
|
||||
0: (1, 1, 1, 2, 2),
|
||||
1: (1, 1, 1, 2, 1),
|
||||
2: (1, 1, 1, 1, 2),
|
||||
3: (1, 1, 1, 1, 1),
|
||||
}
|
||||
tasks = [
|
||||
TileTask(0, (0, 0), z, workload=4),
|
||||
TileTask(1, (0, 1), z, workload=2),
|
||||
TileTask(2, (1, 0), z, workload=2),
|
||||
TileTask(3, (1, 1), z, workload=1),
|
||||
]
|
||||
grid_spec = GridSpec(
|
||||
split_dims=(3, 4),
|
||||
grid_shape=(2, 2),
|
||||
tile_spec={
|
||||
"max_tile_output_shape": (1, 1, 1, 2, 2),
|
||||
"tile_output_shapes": tile_output_shapes,
|
||||
},
|
||||
output_dtype=torch.float32,
|
||||
)
|
||||
seen = {}
|
||||
|
||||
def exec_tile(task):
|
||||
return torch.full(tile_output_shapes[task.tile_id], float(task.tile_id + 1))
|
||||
|
||||
def merge_tiles(coord_tensor_map, passed_grid_spec):
|
||||
seen["merged_shapes"] = {coord: tuple(tile.shape) for coord, tile in coord_tensor_map.items()}
|
||||
assert passed_grid_spec is grid_spec
|
||||
return torch.stack(
|
||||
[
|
||||
coord_tensor_map[(0, 0)].flatten()[0],
|
||||
coord_tensor_map[(0, 1)].flatten()[0],
|
||||
coord_tensor_map[(1, 0)].flatten()[0],
|
||||
coord_tensor_map[(1, 1)].flatten()[0],
|
||||
]
|
||||
)
|
||||
|
||||
operator = DistributedOperator(split=lambda _z: (tasks, grid_spec), exec=exec_tile, merge=merge_tiles)
|
||||
|
||||
rank0_executor = object.__new__(LTX2VaeExecutor)
|
||||
rank0_executor.parallel_size = 2
|
||||
rank0_executor.world_size = 2
|
||||
rank0_executor.rank = 0
|
||||
|
||||
def gather_rank0(local_tile_tensor):
|
||||
assigned = rank0_executor._balance_tasks(tasks, 2)
|
||||
rank1_results = [(task.tile_id, exec_tile(task)) for task in assigned[1]]
|
||||
rank1_tile_tensor = rank0_executor._pack_local_tiles_without_meta(
|
||||
rank1_results,
|
||||
list(local_tile_tensor.shape),
|
||||
z.device,
|
||||
torch.float32,
|
||||
)
|
||||
seen["rank0_gather_shape"] = tuple(local_tile_tensor.shape)
|
||||
return [local_tile_tensor, rank1_tile_tensor]
|
||||
|
||||
def fail_final_sync(*_args, **_kwargs):
|
||||
raise AssertionError("broadcast_result=False should not sync the final result")
|
||||
|
||||
rank0_executor.gather_tensors = gather_rank0
|
||||
rank0_executor._sync_final_result = fail_final_sync
|
||||
|
||||
rank0_result = rank0_executor.execute(z, operator, broadcast_result=False)
|
||||
|
||||
torch.testing.assert_close(rank0_result, torch.tensor([1.0, 2.0, 3.0, 4.0]))
|
||||
assert seen["rank0_gather_shape"] == (2, 1, 1, 1, 2, 2)
|
||||
assert seen["merged_shapes"] == {
|
||||
(0, 0): (1, 1, 1, 2, 2),
|
||||
(0, 1): (1, 1, 1, 2, 1),
|
||||
(1, 0): (1, 1, 1, 1, 2),
|
||||
(1, 1): (1, 1, 1, 1, 1),
|
||||
}
|
||||
|
||||
non_rank0_executor = object.__new__(LTX2VaeExecutor)
|
||||
non_rank0_executor.parallel_size = 2
|
||||
non_rank0_executor.world_size = 2
|
||||
non_rank0_executor.rank = 1
|
||||
|
||||
def gather_rank1(local_tile_tensor):
|
||||
seen["rank1_gather_shape"] = tuple(local_tile_tensor.shape)
|
||||
return None
|
||||
|
||||
def fail_non_rank0_merge(*_args, **_kwargs):
|
||||
raise AssertionError("non-rank0 should not merge gathered tiles")
|
||||
|
||||
non_rank0_executor.gather_tensors = gather_rank1
|
||||
non_rank0_executor._sync_final_result = fail_final_sync
|
||||
|
||||
empty_result = non_rank0_executor.execute(
|
||||
z,
|
||||
DistributedOperator(
|
||||
split=lambda _z: (tasks, grid_spec),
|
||||
exec=exec_tile,
|
||||
merge=fail_non_rank0_merge,
|
||||
),
|
||||
broadcast_result=False,
|
||||
)
|
||||
|
||||
assert tuple(empty_result.shape) == (0,)
|
||||
assert seen["rank1_gather_shape"] == (2, 1, 1, 1, 2, 2)
|
||||
|
||||
|
||||
class TestCFGParallelHelpers:
|
||||
"""Test LTX-2.3 CFG helper math without loading model weights."""
|
||||
|
||||
@@ -9,6 +9,7 @@ import torch
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.ipc import (
|
||||
_SHM_TENSOR_THRESHOLD,
|
||||
DIFFUSION_RPC_RESULT_ENVELOPE,
|
||||
_pack_value_if_large,
|
||||
_unpack_if_shm_handle,
|
||||
pack_diffusion_output_shm,
|
||||
@@ -81,6 +82,31 @@ def test_diffusion_output_list_tensors_round_trip_through_shm() -> None:
|
||||
torch.testing.assert_close(output.output[1], frames[1])
|
||||
|
||||
|
||||
def test_rpc_result_envelope_diffusion_output_round_trips_through_shm() -> None:
|
||||
tensor = torch.arange(300_000, dtype=torch.float32)
|
||||
envelope = {
|
||||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||||
"result": DiffusionOutput(output=tensor),
|
||||
"rank_statuses": [{"rank": 0, "ok": True}],
|
||||
}
|
||||
|
||||
packed = pack_diffusion_output_shm(envelope)
|
||||
|
||||
assert packed is envelope
|
||||
result = packed["result"]
|
||||
assert isinstance(result, DiffusionOutput)
|
||||
assert result.output["__tensor_shm__"] is True
|
||||
assert packed["rank_statuses"] == [{"rank": 0, "ok": True}]
|
||||
|
||||
unpacked = unpack_diffusion_output_shm(packed)
|
||||
|
||||
assert unpacked is envelope
|
||||
result = unpacked["result"]
|
||||
assert isinstance(result, DiffusionOutput)
|
||||
torch.testing.assert_close(result.output, tensor)
|
||||
assert unpacked["rank_statuses"] == [{"rank": 0, "ok": True}]
|
||||
|
||||
|
||||
def test_pack_value_keeps_tensor_at_threshold_inline() -> None:
|
||||
tensor = torch.arange(
|
||||
_SHM_TENSOR_THRESHOLD // torch.empty((), dtype=torch.float32).element_size(),
|
||||
|
||||
@@ -17,8 +17,10 @@ from vllm.v1.engine.exceptions import EngineDeadError
|
||||
from vllm_omni.diffusion.data import DiffusionOutput
|
||||
from vllm_omni.diffusion.diffusion_engine import DiffusionEngine
|
||||
from vllm_omni.diffusion.executor.multiproc_executor import MultiprocDiffusionExecutor
|
||||
from vllm_omni.diffusion.ipc import DIFFUSION_RPC_RESULT_ENVELOPE
|
||||
from vllm_omni.diffusion.sched import RequestScheduler
|
||||
from vllm_omni.diffusion.stage_diffusion_proc import StageDiffusionProc
|
||||
from vllm_omni.diffusion.worker.diffusion_worker import WorkerProc
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
@@ -305,6 +307,80 @@ class TestSerialEngineOperations:
|
||||
assert len(results) == 1
|
||||
assert results[0].error == "rank0"
|
||||
|
||||
def test_collective_rpc_all_rank_status_error_propagation(self):
|
||||
engine, _, _, res_q = _make_engine(num_gpus=2)
|
||||
|
||||
res_q.put(
|
||||
{
|
||||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||||
"method": "add_lora",
|
||||
"result": True,
|
||||
"rank_statuses": [
|
||||
{"rank": 0, "ok": True, "bool_result": True},
|
||||
{
|
||||
"rank": 1,
|
||||
"ok": False,
|
||||
"error": "rank1 boom",
|
||||
"error_type": "RuntimeError",
|
||||
"traceback": "rank1 traceback",
|
||||
},
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError) as excinfo:
|
||||
engine.collective_rpc("add_lora")
|
||||
error = str(excinfo.value)
|
||||
assert "rank 1" in error
|
||||
assert "rank1 boom" in error
|
||||
assert "rank1 traceback" in error
|
||||
|
||||
def test_collective_rpc_all_rank_bool_false_is_aggregated(self):
|
||||
engine, _, _, res_q = _make_engine(num_gpus=2)
|
||||
|
||||
res_q.put(
|
||||
{
|
||||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||||
"method": "remove_lora",
|
||||
"result": True,
|
||||
"rank_statuses": [
|
||||
{"rank": 0, "ok": True, "bool_result": True},
|
||||
{"rank": 1, "ok": True, "bool_result": False},
|
||||
],
|
||||
}
|
||||
)
|
||||
|
||||
assert engine.collective_rpc("remove_lora") == [False]
|
||||
|
||||
def test_collective_rpc_collects_rank_status_only_for_control_plane_all_rank_rpc(self):
|
||||
executor, req_q, res_q = _make_executor(num_gpus=2)
|
||||
|
||||
res_q.put(_tagged_output("forward"))
|
||||
result = executor.collective_rpc(
|
||||
"execute_stepwise",
|
||||
unique_reply_rank=0,
|
||||
exec_all_ranks=True,
|
||||
)
|
||||
forward_rpc = req_q.get_nowait()
|
||||
|
||||
assert result.error == "forward"
|
||||
assert forward_rpc["exec_all_ranks"] is True
|
||||
assert forward_rpc["collect_rank_status"] is False
|
||||
|
||||
res_q.put(
|
||||
{
|
||||
"type": DIFFUSION_RPC_RESULT_ENVELOPE,
|
||||
"method": "remove_lora",
|
||||
"result": True,
|
||||
"rank_statuses": [{"rank": 0, "ok": True, "bool_result": True}],
|
||||
}
|
||||
)
|
||||
assert executor.collective_rpc("remove_lora") == [True]
|
||||
control_rpc = req_q.get_nowait()
|
||||
|
||||
assert control_rpc["exec_all_ranks"] is True
|
||||
assert control_rpc["collect_rank_status"] is True
|
||||
|
||||
def test_serial_add_req_then_collective_rpc(self):
|
||||
engine, _, req_q, res_q = _make_engine()
|
||||
wt = _start_worker(req_q, res_q, count=2)
|
||||
@@ -348,6 +424,115 @@ class TestSerialEngineOperations:
|
||||
engine.collective_rpc("anything")
|
||||
|
||||
|
||||
class TestWorkerProcRpcRankStatus:
|
||||
def _make_worker_proc(self, has_result_mq: bool = True):
|
||||
proc = object.__new__(WorkerProc)
|
||||
proc.gpu_id = 0
|
||||
proc.result_mq = object() if has_result_mq else None
|
||||
proc.worker = SimpleNamespace(execute_method=Mock(return_value=True))
|
||||
return proc
|
||||
|
||||
def test_execute_rpc_returns_rank_status_envelope(self, monkeypatch):
|
||||
proc = self._make_worker_proc()
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: True)
|
||||
monkeypatch.setattr(torch.distributed, "get_world_size", lambda: 2)
|
||||
|
||||
def _all_gather_object(out, local):
|
||||
out[0] = local
|
||||
out[1] = {
|
||||
"rank": 1,
|
||||
"ok": False,
|
||||
"error": "rank1 boom",
|
||||
"error_type": "RuntimeError",
|
||||
"traceback": "trace",
|
||||
"bool_result": None,
|
||||
}
|
||||
|
||||
monkeypatch.setattr(torch.distributed, "all_gather_object", _all_gather_object)
|
||||
|
||||
result, should_reply = proc.execute_rpc(
|
||||
{
|
||||
"method": "remove_lora",
|
||||
"args": (),
|
||||
"kwargs": {},
|
||||
"output_rank": 0,
|
||||
"exec_all_ranks": True,
|
||||
"collect_rank_status": True,
|
||||
}
|
||||
)
|
||||
|
||||
assert should_reply is True
|
||||
assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
|
||||
assert result["result"] is True
|
||||
assert result["rank_statuses"][0]["rank"] == 0
|
||||
assert result["rank_statuses"][1]["rank"] == 1
|
||||
assert result["rank_statuses"][1]["ok"] is False
|
||||
|
||||
def test_execute_rpc_local_exception_is_reported_in_envelope(self, monkeypatch):
|
||||
proc = self._make_worker_proc()
|
||||
proc.worker.execute_method = Mock(side_effect=RuntimeError("local boom"))
|
||||
monkeypatch.setattr(torch.distributed, "is_initialized", lambda: False)
|
||||
|
||||
result, should_reply = proc.execute_rpc(
|
||||
{
|
||||
"method": "add_lora",
|
||||
"args": (),
|
||||
"kwargs": {},
|
||||
"output_rank": 0,
|
||||
"exec_all_ranks": True,
|
||||
"collect_rank_status": True,
|
||||
}
|
||||
)
|
||||
|
||||
assert should_reply is True
|
||||
assert result["type"] == DIFFUSION_RPC_RESULT_ENVELOPE
|
||||
assert len(result["rank_statuses"]) == 1
|
||||
status = result["rank_statuses"][0]
|
||||
assert status["rank"] == 0
|
||||
assert status["ok"] is False
|
||||
assert status["error"] == "local boom"
|
||||
assert status["error_type"] == "RuntimeError"
|
||||
assert status["bool_result"] is None
|
||||
assert "local boom" in status["traceback"]
|
||||
|
||||
def test_execute_rpc_rejects_collect_rank_status_without_all_ranks(self):
|
||||
proc = self._make_worker_proc()
|
||||
|
||||
with pytest.raises(ValueError, match="collect_rank_status requires exec_all_ranks=True"):
|
||||
proc.execute_rpc(
|
||||
{
|
||||
"method": "ping",
|
||||
"args": (),
|
||||
"kwargs": {},
|
||||
"output_rank": 0,
|
||||
"exec_all_ranks": False,
|
||||
"collect_rank_status": True,
|
||||
}
|
||||
)
|
||||
|
||||
proc.worker.execute_method.assert_not_called()
|
||||
|
||||
def test_execute_rpc_non_collect_exception_preserves_original_type(self):
|
||||
proc = self._make_worker_proc()
|
||||
original = ValueError("local boom")
|
||||
proc.worker.execute_method = Mock(side_effect=original)
|
||||
|
||||
with pytest.raises(ValueError) as excinfo:
|
||||
proc.execute_rpc(
|
||||
{
|
||||
"method": "bad",
|
||||
"args": (),
|
||||
"kwargs": {},
|
||||
"output_rank": 0,
|
||||
"exec_all_ranks": False,
|
||||
"collect_rank_status": False,
|
||||
}
|
||||
)
|
||||
|
||||
assert excinfo.value is original
|
||||
|
||||
|
||||
# ───────── error handling: EngineDeadError propagation through layers ─────
|
||||
|
||||
|
||||
|
||||
@@ -145,7 +145,7 @@ def test_shm_connector_flow(mocker: MockerFixture):
|
||||
Using real SharedMemoryConnector (inline mode for simplicity).
|
||||
"""
|
||||
# 1. Setup Connector
|
||||
config = {"shm_threshold_bytes": 1024} # Large threshold to use inline
|
||||
config = {"shm_threshold_bytes": 1024, "inline_small_payloads": True}
|
||||
connector = SharedMemoryConnector(config)
|
||||
connectors_map = {("0", "1"): connector}
|
||||
|
||||
|
||||
@@ -75,18 +75,21 @@ def test_create_unknown_connector():
|
||||
|
||||
@pytest.fixture
|
||||
def shm_connector():
|
||||
config = {"shm_threshold_bytes": 100} # Small threshold for testing
|
||||
config = {"shm_threshold_bytes": 100, "inline_small_payloads": True}
|
||||
return SharedMemoryConnector(config)
|
||||
|
||||
|
||||
def test_put_get_inline(shm_connector):
|
||||
"""Test transfer for small data (currently always uses SHM)."""
|
||||
"""Test inline transfer for small data."""
|
||||
data = {"small": "data"}
|
||||
|
||||
success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
|
||||
assert success is True
|
||||
assert "shm" in metadata
|
||||
assert "inline_bytes" in metadata
|
||||
assert "shm" not in metadata
|
||||
assert "size" in metadata
|
||||
assert shm_connector._metrics["inline_writes"] == 1
|
||||
assert shm_connector._metrics["shm_writes"] == 0
|
||||
|
||||
# Retrieve
|
||||
retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
|
||||
|
||||
@@ -200,9 +200,9 @@ def test_save_async(build_adapter):
|
||||
request = _req("req-1", RequestStatus.WAITING, external_req_id="external-1")
|
||||
|
||||
adapter.custom_process_next_stage_input_func = lambda **kwargs: {"x": [1], "finished": False}
|
||||
adapter.save_async(pooling_output=None, request=request)
|
||||
adapter.save_async(multimodal_output=None, request=request)
|
||||
adapter.custom_process_next_stage_input_func = lambda **kwargs: {}
|
||||
adapter.save_async(pooling_output=None, request=request)
|
||||
adapter.save_async(multimodal_output=None, request=request)
|
||||
|
||||
task = adapter._pending_save_reqs.popleft()
|
||||
assert task["is_finished"] is False
|
||||
@@ -214,7 +214,7 @@ def test_save_async_uses_confirmed_tokens_for_async_scheduler_watermark(build_ad
|
||||
request.num_computed_tokens = 10
|
||||
request.num_output_placeholders = 2
|
||||
|
||||
adapter.save_async(pooling_output=None, request=request)
|
||||
adapter.save_async(multimodal_output=None, request=request)
|
||||
|
||||
assert adapter.requests_num_chunks_sent["external-async"] == 8
|
||||
assert len(adapter._pending_save_reqs) == 1
|
||||
@@ -236,7 +236,7 @@ def test_send_single_request_struct_without_meta_does_not_crash(build_adapter, m
|
||||
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
|
||||
|
||||
adapter._send_single_request(
|
||||
{"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
|
||||
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
|
||||
)
|
||||
|
||||
assert cleanup_calls == [] # no terminal cleanup; meta.finished is false
|
||||
@@ -256,7 +256,7 @@ def test_send_single_request_empty_struct_goes_on_wire(build_adapter, monkeypatc
|
||||
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
|
||||
|
||||
adapter._send_single_request(
|
||||
{"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
|
||||
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": False}
|
||||
)
|
||||
|
||||
assert connector.put.called
|
||||
@@ -274,7 +274,7 @@ def test_send_single_request_struct_preserves_segment_finished(build_adapter, mo
|
||||
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: None)
|
||||
|
||||
adapter._send_single_request(
|
||||
{"pooling_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
|
||||
{"multimodal_output": None, "request": request, "is_finished": False, "is_segment_finished": True}
|
||||
)
|
||||
|
||||
sent_payload = connector.put.call_args.kwargs["data"]
|
||||
@@ -289,13 +289,13 @@ def test_save_async_skips_stale_resumable_chunk_until_dedup_is_reset(build_adapt
|
||||
request.num_computed_tokens = 0
|
||||
adapter.requests_num_chunks_sent["ext-stream"] = 111
|
||||
|
||||
adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
|
||||
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
|
||||
|
||||
assert len(adapter._pending_save_reqs) == 0
|
||||
assert adapter.requests_num_chunks_sent["ext-stream"] == 111
|
||||
|
||||
adapter.requests_num_chunks_sent.pop("ext-stream")
|
||||
adapter.save_async(pooling_output=None, request=request, is_segment_finished=False)
|
||||
adapter.save_async(multimodal_output=None, request=request, is_segment_finished=False)
|
||||
|
||||
assert len(adapter._pending_save_reqs) == 1
|
||||
assert adapter.requests_num_chunks_sent["ext-stream"] == 0
|
||||
@@ -312,7 +312,7 @@ def test_send_single_request_cleans_up_after_finished_payload(build_adapter, mon
|
||||
monkeypatch.setattr(adapter, "cleanup", lambda *a, **kw: cleanup_calls.append((a, kw)))
|
||||
|
||||
adapter._send_single_request(
|
||||
{"pooling_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
|
||||
{"multimodal_output": None, "request": request, "is_finished": True, "is_segment_finished": True}
|
||||
)
|
||||
|
||||
assert len(cleanup_calls) == 1
|
||||
@@ -387,14 +387,15 @@ def test_process_and_restore_queues(build_adapter):
|
||||
running_req = _req("r1", RequestStatus.RUNNING)
|
||||
waiting_queue = DummyWaitingQueue([waiting_req])
|
||||
running_queue = [running_req]
|
||||
scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert running_req.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert waiting_queue == []
|
||||
assert running_queue == []
|
||||
|
||||
adapter.restore_queues(waiting_queue, running_queue)
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert waiting_queue == [waiting_req]
|
||||
assert running_queue == [running_req]
|
||||
assert adapter.waiting_for_chunk_waiting_requests == deque()
|
||||
@@ -897,6 +898,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
sched.requests = {}
|
||||
sched.running = []
|
||||
sched.waiting = []
|
||||
|
||||
with patch.object(VLLMScheduler, "finish_requests", _super_finish):
|
||||
OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
|
||||
@@ -1170,3 +1173,164 @@ def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
|
||||
assert eco.outputs[0].finish_reason == "stop"
|
||||
assert eco.outputs[0].kv_transfer_params is None
|
||||
assert scheduler._pending_finish_reqs == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# Zombie purge tests (regression for vllm-project/vllm-omni#3736)
|
||||
# ---------------------------------------------------------------
|
||||
|
||||
|
||||
def test_process_pending_chunks_purges_zombies_in_running_deque(
|
||||
build_adapter,
|
||||
):
|
||||
"""A request aborted while parked in ``waiting_for_chunk_running_requests``
|
||||
is no longer tracked by ``scheduler.requests`` once
|
||||
``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop
|
||||
that zombie *before* ``restore_queues`` would re-inject it onto the
|
||||
running queue.
|
||||
|
||||
Regression for https://github.com/vllm-project/vllm-omni/issues/3736
|
||||
(engine-core ``KeyError`` on aborts under chunk-transfer pressure).
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
adapter.waiting_for_chunk_running_requests.append(live_req)
|
||||
adapter.waiting_for_chunk_running_requests.append(zombie_req)
|
||||
# Mirror state the adapter would carry for an in-flight request so we can
|
||||
# later assert ``cleanup_receiver`` actually fired against the zombie.
|
||||
adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING
|
||||
adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
# Simulate the post-abort scheduler: only ``live-1`` is still tracked.
|
||||
scheduler_requests = {live_req.request_id: live_req}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
|
||||
# 1. Zombie was popped from the deque.
|
||||
assert zombie_req not in adapter.waiting_for_chunk_running_requests
|
||||
# 2. Live request is still in the deque.
|
||||
assert live_req in adapter.waiting_for_chunk_running_requests
|
||||
# 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping
|
||||
# and registers the id as cancelled so a late load/poll is dropped too).
|
||||
assert zombie_req.request_id not in adapter.requests_origin_status
|
||||
assert zombie_req.request_id in adapter._cancelled_load_reqs
|
||||
# 4. Live request's bookkeeping is untouched.
|
||||
assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING
|
||||
assert live_req.request_id not in adapter._cancelled_load_reqs
|
||||
|
||||
# 5. ``restore_queues`` (which the scheduler runs in its ``finally``
|
||||
# clause) now only re-injects the live request -- the zombie is
|
||||
# gone, so the worker's ``_update_states`` cannot crash on it.
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert running_queue == [live_req]
|
||||
|
||||
|
||||
def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter):
|
||||
"""Zombie purge applies symmetrically to the waiting-side deque.
|
||||
|
||||
Regression for https://github.com/vllm-project/vllm-omni/issues/3736 --
|
||||
aborted requests that landed in ``waiting_for_chunk_waiting_requests``
|
||||
must also be removed before ``restore_queues`` re-injects them, since
|
||||
``restore_queues`` uses ``waiting_queue.add_request`` rather than
|
||||
``running_queue.extend`` for that path.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK)
|
||||
adapter.waiting_for_chunk_waiting_requests.append(live_req)
|
||||
adapter.waiting_for_chunk_waiting_requests.append(zombie_req)
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
scheduler_requests = {live_req.request_id: live_req}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
|
||||
assert live_req in adapter.waiting_for_chunk_waiting_requests
|
||||
assert zombie_req not in adapter.waiting_for_chunk_waiting_requests
|
||||
assert zombie_req.request_id in adapter._cancelled_load_reqs
|
||||
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert waiting_queue == [live_req]
|
||||
|
||||
|
||||
def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter):
|
||||
"""Interleaved live/zombie ordering -- pin the ``popleft`` + append-live
|
||||
in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests:
|
||||
survivor order must match insertion order, and ``cleanup_receiver`` must
|
||||
run exactly once per zombie.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK)
|
||||
for req in (live1, zombie1, live2, zombie2):
|
||||
adapter.waiting_for_chunk_running_requests.append(req)
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
scheduler_requests = {live1.request_id: live1, live2.request_id: live2}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
|
||||
assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2]
|
||||
assert zombie1.request_id in adapter._cancelled_load_reqs
|
||||
assert zombie2.request_id in adapter._cancelled_load_reqs
|
||||
|
||||
|
||||
def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
|
||||
build_adapter,
|
||||
):
|
||||
"""Race window guard: an abort can fire between ``process_pending_chunks``
|
||||
and the scheduler's ``finally``-clause ``restore_queues`` call. The
|
||||
second purge inside ``restore_queues`` must drop the now-untracked
|
||||
request so it does not get re-injected onto ``running_queue``.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK)
|
||||
adapter.waiting_for_chunk_running_requests.append(req)
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
|
||||
# Tick N: process_pending_chunks runs while req is still tracked.
|
||||
scheduler_requests: dict = {req.request_id: req}
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert req in adapter.waiting_for_chunk_running_requests
|
||||
|
||||
# Mid-tick: abort fires and the scheduler's free path deletes the entry.
|
||||
del scheduler_requests[req.request_id]
|
||||
|
||||
# finally: restore_queues sees the now-untracked req and must drop it
|
||||
# instead of blindly extending it onto running_queue.
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert running_queue == []
|
||||
assert req.request_id in adapter._cancelled_load_reqs
|
||||
|
||||
|
||||
def test_purge_is_noop_on_empty_deques(build_adapter):
|
||||
"""Empty deques short-circuit -- guards against any accidental
|
||||
``IndexError`` from ``popleft`` on an empty deque if a future caller
|
||||
shadows the empty check.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
|
||||
assert len(adapter.waiting_for_chunk_running_requests) == 0
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={})
|
||||
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
|
||||
assert len(adapter.waiting_for_chunk_running_requests) == 0
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
|
||||
assert running_queue == []
|
||||
assert waiting_queue == []
|
||||
@@ -13,7 +13,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@pytest.fixture()
|
||||
def connector():
|
||||
c = SharedMemoryConnector({"shm_threshold_bytes": 64})
|
||||
c = SharedMemoryConnector({"shm_threshold_bytes": 0})
|
||||
yield c
|
||||
c.close()
|
||||
|
||||
|
||||
@@ -29,7 +29,46 @@ from tests.e2e.accuracy.helpers import (
|
||||
)
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner, OmniServer
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import (
|
||||
build_prompt_tokens,
|
||||
resolve_stop_token_ids,
|
||||
)
|
||||
|
||||
# CoT structural markers that MUST appear in AR output text.
|
||||
# The AR trajectory under think_recaption mode is:
|
||||
# <think> ... </think> <recaption> ... <relation_1> ... </relation_1>
|
||||
# <relation_2> ... </relation_2> </recaption>
|
||||
# Test deploy config sets include_stop_str_in_output=True and
|
||||
# skip_special_tokens=False, so ALL opening and closing markers
|
||||
# should be present in the detokenized text.
|
||||
_THINK_CLOSE = "</" + "think>" # avoid Python slash-t-as-TAB pitfall
|
||||
_COT_REQUIRED_MARKERS = (
|
||||
"<recaption>",
|
||||
"</recaption>",
|
||||
"<relation_1>",
|
||||
"</relation_1>",
|
||||
"<relation_2>",
|
||||
"</relation_2>",
|
||||
_THINK_CLOSE,
|
||||
)
|
||||
|
||||
|
||||
def _assert_cot_structural_markers(cot_text: str, label: str) -> None:
|
||||
"""Assert that the AR CoT output contains all required structural tags.
|
||||
|
||||
Guards against regressions where detokenizer stop-token stripping
|
||||
or truncation logic loses CoT content. Under the test deploy config
|
||||
(include_stop_str_in_output=True, skip_special_tokens=False),
|
||||
ALL markers should be present in the detokenized text.
|
||||
"""
|
||||
missing = [m for m in _COT_REQUIRED_MARKERS if m not in cot_text]
|
||||
assert not missing, (
|
||||
f"[{label}] CoT text missing required structural markers: {missing}. "
|
||||
f"This typically indicates that stop-token stripping or "
|
||||
f"_truncate_at_cot_end truncation removed CoT content. "
|
||||
f"CoT text length={len(cot_text)}, first 100 chars: {repr(cot_text[:100])}"
|
||||
)
|
||||
|
||||
|
||||
os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
|
||||
|
||||
@@ -129,6 +168,7 @@ _DEPLOY_CONFIG = {
|
||||
"stop_token_ids": [128025],
|
||||
"detokenize": True,
|
||||
"skip_special_tokens": False,
|
||||
"include_stop_str_in_output": True,
|
||||
},
|
||||
},
|
||||
{
|
||||
@@ -400,6 +440,8 @@ def test_image_to_image_alignment_online(accuracy_artifact_root: Path, accuracy_
|
||||
]
|
||||
print("[ONLINE] " + tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
|
||||
|
||||
_assert_cot_structural_markers(online_cot, "ONLINE")
|
||||
|
||||
assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
|
||||
f"[ONLINE] COT semantic similarity {cot_results['cot_semantic_sim']:.4f} below threshold {THRESHOLDS['cot_semantic_sim']}"
|
||||
)
|
||||
@@ -508,6 +550,8 @@ def test_image_to_image_alignment(accuracy_artifact_root: Path, accuracy_assets_
|
||||
|
||||
print(tabulate(table, headers=["Metric", "Value", "L20x Reference"], tablefmt="grid"))
|
||||
|
||||
_assert_cot_structural_markers(omni_cot, "OFFLINE")
|
||||
|
||||
assert cot_results["cot_semantic_sim"] >= THRESHOLDS["cot_semantic_sim"], (
|
||||
f"COT semantic similarity {cot_results['cot_semantic_sim']:.4f} is below threshold {THRESHOLDS['cot_semantic_sim']}"
|
||||
)
|
||||
|
||||
@@ -0,0 +1,37 @@
|
||||
# Stage config for Step-Audio2 CI tests
|
||||
# Single Thinker stage (audio/text understanding → text). ASR/text-output
|
||||
# tests don't need Token2Wav; a two-stage config would deadlock on stage-1 skip.
|
||||
|
||||
# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: StepAudio2ThinkerForConditionalGeneration
|
||||
tensor_parallel_size: 1
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
max_model_len: 4096
|
||||
max_num_batched_tokens: 4096
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.8
|
||||
skip_mm_profiling: true
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: text # Output text + audio tokens
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.7
|
||||
top_p: 0.9
|
||||
top_k: -1
|
||||
max_tokens: 256
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.05
|
||||
@@ -51,7 +51,7 @@ test_params = [(m, c) for m in models for c in [get_eager_config()]]
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
|
||||
def test_text_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
@@ -67,7 +67,7 @@ def test_text_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
|
||||
def test_image_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
@@ -84,7 +84,7 @@ def test_image_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
|
||||
def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
@@ -103,7 +103,7 @@ def test_audio_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
|
||||
def test_video_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
@@ -120,7 +120,7 @@ def test_video_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params_thinker, indirect=True)
|
||||
def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
@@ -140,7 +140,7 @@ def test_mixed_to_text(omni_runner, omni_runner_handler) -> None:
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
|
||||
def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
@@ -156,7 +156,7 @@ def test_text_to_audio(omni_runner, omni_runner_handler) -> None:
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
|
||||
def test_image_to_audio(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
|
||||
@@ -0,0 +1,297 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""End-to-end offline inference tests for Ming-omni-tts."""
|
||||
|
||||
import asyncio
|
||||
import uuid
|
||||
from collections.abc import Mapping
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from transformers import AutoTokenizer
|
||||
from vllm import SamplingParams
|
||||
from vllm.inputs import tokens_input
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
from vllm_omni import AsyncOmni
|
||||
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
KEY_MAX_DECODE_STEPS,
|
||||
SAMPLE_RATE,
|
||||
TEXT_EOS_TOKEN_ID,
|
||||
)
|
||||
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import DEFAULT_PROMPT, build_ming_dense_prompt
|
||||
|
||||
MODEL = "inclusionAI/Ming-omni-tts-0.5B"
|
||||
DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
|
||||
TEST_TEXT = "我会一直在这里陪着你,直到你慢慢地沉入那个最温柔的梦里。"
|
||||
TEST_INSTRUCTION = "轻柔的ASMR耳语,慢速,贴近麦克风"
|
||||
MIN_AUDIO_SAMPLES = 1000
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def ming_tokenizer():
|
||||
return AutoTokenizer.from_pretrained(MODEL, trust_remote_code=False)
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def ming_engine():
|
||||
with OmniRunner(
|
||||
MODEL,
|
||||
deploy_config=DEPLOY_CONFIG,
|
||||
stage_init_timeout=300,
|
||||
enforce_eager=True,
|
||||
) as runner:
|
||||
yield runner.omni
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def async_omni_engine():
|
||||
engine = AsyncOmni(
|
||||
model=MODEL,
|
||||
deploy_config=DEPLOY_CONFIG,
|
||||
stage_init_timeout=300,
|
||||
enforce_eager=True,
|
||||
)
|
||||
yield engine
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
def _build_prompt(
|
||||
tokenizer,
|
||||
*,
|
||||
text: str = TEST_TEXT,
|
||||
instruction=TEST_INSTRUCTION,
|
||||
use_zero_spk_emb: bool = True,
|
||||
) -> dict:
|
||||
prompt_dict = build_ming_dense_prompt(
|
||||
tokenizer,
|
||||
prompt=DEFAULT_PROMPT,
|
||||
text=text,
|
||||
instruction=instruction,
|
||||
runtime_controls={KEY_MAX_DECODE_STEPS: 200},
|
||||
use_zero_spk_emb=use_zero_spk_emb,
|
||||
)
|
||||
prompt = tokens_input(prompt_token_ids=prompt_dict["prompt_token_ids"])
|
||||
prompt["prompt"] = prompt_dict["prompt"]
|
||||
prompt["text"] = prompt_dict["text"]
|
||||
prompt["additional_information"] = prompt_dict["additional_information"]
|
||||
return prompt
|
||||
|
||||
|
||||
def _sampling_params_list() -> list[SamplingParams]:
|
||||
return [
|
||||
SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=201,
|
||||
stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
|
||||
),
|
||||
SamplingParams(temperature=0.0, max_tokens=1),
|
||||
]
|
||||
|
||||
|
||||
def _flatten_audio(audio) -> torch.Tensor:
|
||||
if isinstance(audio, list):
|
||||
parts = [torch.as_tensor(item, dtype=torch.float32).reshape(-1).cpu() for item in audio]
|
||||
parts = [item for item in parts if item.numel() > 0]
|
||||
if not parts:
|
||||
return torch.zeros((0,), dtype=torch.float32)
|
||||
return torch.cat(parts, dim=0)
|
||||
return torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
|
||||
|
||||
|
||||
def _extract_audio(multimodal_output: Mapping[str, Any]) -> torch.Tensor:
|
||||
assert isinstance(multimodal_output, Mapping), f"Expected Mapping, got {type(multimodal_output)}"
|
||||
audio = multimodal_output.get("audio", multimodal_output.get("model_outputs"))
|
||||
assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}"
|
||||
waveform = _flatten_audio(audio)
|
||||
if waveform.numel() == 0:
|
||||
raise RuntimeError("Generated audio waveform is empty")
|
||||
return waveform
|
||||
|
||||
|
||||
def _extract_sample_rate(multimodal_output: Mapping[str, Any]) -> int:
|
||||
sample_rate = multimodal_output.get("sr")
|
||||
if sample_rate is None:
|
||||
raise RuntimeError("Expected multimodal_output['sr']")
|
||||
if isinstance(sample_rate, list):
|
||||
sample_rate = sample_rate[-1]
|
||||
if hasattr(sample_rate, "item"):
|
||||
sample_rate = sample_rate.item()
|
||||
return int(sample_rate)
|
||||
|
||||
|
||||
def _extract_final_audio_outputs(outputs):
|
||||
final_outputs = []
|
||||
for item in outputs:
|
||||
if getattr(item, "final_output_type", None) == "audio":
|
||||
final_outputs.append(item)
|
||||
continue
|
||||
request_output = getattr(item, "request_output", None)
|
||||
if request_output is None:
|
||||
continue
|
||||
multimodal_output = getattr(request_output, "multimodal_output", None)
|
||||
if isinstance(multimodal_output, Mapping):
|
||||
final_outputs.append(item)
|
||||
continue
|
||||
completions = getattr(request_output, "outputs", None) or []
|
||||
if any(isinstance(getattr(completion, "multimodal_output", None), Mapping) for completion in completions):
|
||||
final_outputs.append(item)
|
||||
return final_outputs
|
||||
|
||||
|
||||
def _extract_multimodal_output(output) -> Mapping[str, Any]:
|
||||
multimodal_output = getattr(output, "multimodal_output", None)
|
||||
if isinstance(multimodal_output, Mapping):
|
||||
return multimodal_output
|
||||
|
||||
request_output = getattr(output, "request_output", None)
|
||||
if request_output is not None:
|
||||
multimodal_output = getattr(request_output, "multimodal_output", None)
|
||||
if isinstance(multimodal_output, Mapping):
|
||||
return multimodal_output
|
||||
completions = getattr(request_output, "outputs", None) or []
|
||||
for completion in completions:
|
||||
multimodal_output = getattr(completion, "multimodal_output", None)
|
||||
if isinstance(multimodal_output, Mapping):
|
||||
return multimodal_output
|
||||
|
||||
raise AssertionError("No multimodal audio output found in Ming generate results")
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_ming_tts_offline_basic(ming_engine, ming_tokenizer) -> None:
|
||||
"""Test blocking Ming generation through Omni."""
|
||||
outputs = ming_engine.generate(
|
||||
prompts=[_build_prompt(ming_tokenizer)],
|
||||
sampling_params_list=_sampling_params_list(),
|
||||
py_generator=False,
|
||||
)
|
||||
final_outputs = _extract_final_audio_outputs(outputs)
|
||||
assert len(final_outputs) == 1, f"Expected one final audio output, got {len(final_outputs)}"
|
||||
multimodal_output = _extract_multimodal_output(final_outputs[0])
|
||||
waveform = _extract_audio(multimodal_output)
|
||||
sample_rate = _extract_sample_rate(multimodal_output)
|
||||
assert waveform.ndim == 1
|
||||
assert waveform.shape[0] == waveform.numel()
|
||||
assert waveform.numel() > MIN_AUDIO_SAMPLES
|
||||
assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
|
||||
assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_ming_tts_speaker_conditioning_differs(ming_engine, ming_tokenizer) -> None:
|
||||
"""Test that different Ming speaker controls produce different waveform outputs."""
|
||||
style_outputs = ming_engine.generate(
|
||||
prompts=[_build_prompt(ming_tokenizer)],
|
||||
sampling_params_list=_sampling_params_list(),
|
||||
py_generator=False,
|
||||
)
|
||||
ip_outputs = ming_engine.generate(
|
||||
prompts=[_build_prompt(ming_tokenizer, text=TEST_TEXT, instruction={"IP": "灵小甄"}, use_zero_spk_emb=True)],
|
||||
sampling_params_list=_sampling_params_list(),
|
||||
py_generator=False,
|
||||
)
|
||||
|
||||
style_final_outputs = _extract_final_audio_outputs(style_outputs)
|
||||
ip_final_outputs = _extract_final_audio_outputs(ip_outputs)
|
||||
assert len(style_final_outputs) == 1, "No style audio output produced"
|
||||
assert len(ip_final_outputs) == 1, "No IP audio output produced"
|
||||
|
||||
style_waveform = _extract_audio(_extract_multimodal_output(style_final_outputs[0]))
|
||||
ip_waveform = _extract_audio(_extract_multimodal_output(ip_final_outputs[0]))
|
||||
assert style_waveform.numel() > MIN_AUDIO_SAMPLES
|
||||
assert ip_waveform.numel() > MIN_AUDIO_SAMPLES
|
||||
assert np.max(np.abs(style_waveform.numpy())) > 0.01, "Style audio appears silent"
|
||||
assert np.max(np.abs(ip_waveform.numpy())) > 0.01, "IP audio appears silent"
|
||||
|
||||
overlap = min(int(style_waveform.numel()), int(ip_waveform.numel()))
|
||||
mean_abs_diff = torch.mean(torch.abs(style_waveform[:overlap] - ip_waveform[:overlap])).item()
|
||||
assert style_waveform.shape != ip_waveform.shape or mean_abs_diff > 1e-4, (
|
||||
"Speaker-conditioned outputs should differ, but style and IP waveforms were effectively identical"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_ming_tts_multiple_prompts_queued(ming_engine, ming_tokenizer) -> None:
|
||||
"""Regression: supported max_num_seqs=1 config must still drain queued prompts."""
|
||||
prompts = [
|
||||
_build_prompt(
|
||||
ming_tokenizer,
|
||||
text="第一条语音用于验证 Ming dense 队列中的长请求可以完成。",
|
||||
instruction="平静自然的旁白,语速中等",
|
||||
),
|
||||
_build_prompt(
|
||||
ming_tokenizer,
|
||||
text="第二条短语音也必须生成完成。",
|
||||
instruction="清晰明亮的女声",
|
||||
),
|
||||
]
|
||||
outputs = ming_engine.generate(
|
||||
prompts=prompts,
|
||||
sampling_params_list=_sampling_params_list(),
|
||||
py_generator=False,
|
||||
)
|
||||
final_outputs = _extract_final_audio_outputs(outputs)
|
||||
assert len(final_outputs) == len(prompts), f"Expected {len(prompts)} audio outputs, got {len(final_outputs)}"
|
||||
for i, output in enumerate(final_outputs):
|
||||
waveform = _extract_audio(_extract_multimodal_output(output))
|
||||
duration_s = waveform.numel() / SAMPLE_RATE
|
||||
assert 0.1 < duration_s < 30.0, f"Request {i} audio duration out of range: {duration_s:.2f}s"
|
||||
assert np.max(np.abs(waveform.numpy())) > 0.01, f"Request {i} audio appears silent"
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
def test_ming_tts_offline_streaming(async_omni_engine, ming_tokenizer) -> None:
|
||||
"""Test async_chunk streaming Ming generation through AsyncOmni."""
|
||||
|
||||
async def _run() -> None:
|
||||
all_audio_chunks = []
|
||||
accumulated_samples = 0
|
||||
chunk_idx = 0
|
||||
sample_rate = None
|
||||
async for stage_output in async_omni_engine.generate(
|
||||
prompt=_build_prompt(ming_tokenizer),
|
||||
request_id=str(uuid.uuid4()),
|
||||
sampling_params_list=_sampling_params_list(),
|
||||
):
|
||||
multimodal_output = stage_output.multimodal_output or {}
|
||||
audio = multimodal_output.get("audio")
|
||||
if "sr" in multimodal_output:
|
||||
sample_rate = _extract_sample_rate(multimodal_output)
|
||||
if audio is None:
|
||||
continue
|
||||
finished = stage_output.finished
|
||||
if isinstance(audio, torch.Tensor):
|
||||
if finished:
|
||||
audio_chunk = audio[accumulated_samples:].float().detach().cpu()
|
||||
else:
|
||||
audio_chunk = audio.float().detach().cpu()
|
||||
elif isinstance(audio, list):
|
||||
audio_chunk = _flatten_audio(audio)
|
||||
else:
|
||||
audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
|
||||
accumulated_samples += int(audio_chunk.numel())
|
||||
chunk_idx += 1
|
||||
if audio_chunk.numel() > 0:
|
||||
all_audio_chunks.append(audio_chunk)
|
||||
assert all_audio_chunks, "No streaming audio chunks received"
|
||||
waveform = torch.cat(all_audio_chunks, dim=0)
|
||||
assert waveform.numel() > MIN_AUDIO_SAMPLES
|
||||
assert np.max(np.abs(waveform.numpy())) > 0.01, "Audio appears silent"
|
||||
assert sample_rate is not None, "Streaming path did not return a sample rate"
|
||||
assert sample_rate == SAMPLE_RATE, f"Expected Ming output sample rate {SAMPLE_RATE}, got {sample_rate}"
|
||||
|
||||
asyncio.run(_run())
|
||||
@@ -32,6 +32,7 @@ _OMNI_RUNNER_PARAM = (
|
||||
)
|
||||
|
||||
pytestmark = [
|
||||
pytest.mark.skip(reason="issue#4361"),
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.tts,
|
||||
pytest.mark.parametrize("omni_runner", [_OMNI_RUNNER_PARAM], indirect=True),
|
||||
|
||||
@@ -0,0 +1,160 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Expanded end-to-end test for SenseNova-U1 in offline mode.
|
||||
|
||||
This test file primarily covers end-to-end tests for LoRA support.
|
||||
|
||||
Validates that LoRA adapters are correctly loaded, applied with controllable
|
||||
scale, and cleanly deactivated. Uses a synthetic rank-4 adapter targeting the
|
||||
first decoder layer's gen-path (``_mot_gen``) output projection, so the
|
||||
perturbation is exercised during the denoising loop.
|
||||
|
||||
Assertions:
|
||||
(a) LoRA at scale=1.0 visibly changes the output (diff > 0.5)
|
||||
(b) scale=2.0 produces a different delta than scale=1.0 (scale sensitivity)
|
||||
(c) The delta is bounded (not corrupted)
|
||||
(d) Deactivating LoRA exactly restores the baseline (diff == 0)
|
||||
|
||||
Observed on H100 (seed=42, 8 steps, 512x512): diff_1x~62.2, diff_2x~76.9,
|
||||
diff_restored=0.0. Thresholds match test_bagel_expansion.py and hold with margin.
|
||||
"""
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
from PIL import Image
|
||||
from safetensors.torch import save_file
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.lora.utils import stable_lora_int_id
|
||||
|
||||
MODEL = "SenseNova/SenseNova-U1-8B-MoT"
|
||||
DEFAULT_PROMPT = "a photo of a cat sitting on a laptop keyboard"
|
||||
SIZE = (512, 512) # (width, height); small for fast E2E
|
||||
|
||||
# SenseNova-U1-8B-MoT hidden_size (Qwen3-based). o_proj is square (hidden x hidden)
|
||||
# regardless of GQA, so this is the robust LoRA target. Verify against config.json.
|
||||
_LORA_DIM = 4096
|
||||
_LORA_MODULE = "language_model.model.layers.0.self_attn.o_proj_mot_gen"
|
||||
_LORA_RANK = 4
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers (reused from test_sensenova_u1_text2img.py patterns)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _build_sampling_params(
|
||||
lora_request: LoRARequest | None = None,
|
||||
lora_scale: float = 1.0,
|
||||
) -> OmniDiffusionSamplingParams:
|
||||
return OmniDiffusionSamplingParams(
|
||||
height=SIZE[1],
|
||||
width=SIZE[0],
|
||||
seed=42,
|
||||
num_inference_steps=8,
|
||||
extra_args={
|
||||
"cfg_scale": 4.0,
|
||||
"cfg_norm": "none",
|
||||
"timestep_shift": 3.0,
|
||||
"cfg_interval": (0.0, 1.0),
|
||||
"batch_size": 1,
|
||||
"think": False,
|
||||
"t_eps": 0.02,
|
||||
},
|
||||
lora_request=lora_request,
|
||||
lora_scale=lora_scale,
|
||||
)
|
||||
|
||||
|
||||
def _extract_generated_image(omni_outputs: list) -> Image.Image | None:
|
||||
for req_output in omni_outputs:
|
||||
if images := getattr(req_output, "images", None):
|
||||
return images[0]
|
||||
return None
|
||||
|
||||
|
||||
def _generate(omni: Omni, lora_request: LoRARequest | None = None, lora_scale: float = 1.0) -> Image.Image:
|
||||
outputs = list(
|
||||
omni.generate(
|
||||
prompts={"prompt": DEFAULT_PROMPT, "modalities": ["image"]},
|
||||
sampling_params_list=_build_sampling_params(lora_request, lora_scale),
|
||||
)
|
||||
)
|
||||
img = _extract_generated_image(outputs)
|
||||
assert img is not None, "No image generated"
|
||||
return img
|
||||
|
||||
|
||||
def _make_file_lora_request(adapter_dir: Path) -> LoRARequest:
|
||||
"""Write synthetic adapter to disk and return a file-backed LoRARequest."""
|
||||
adapter_dir.mkdir(parents=True, exist_ok=True)
|
||||
gen = torch.Generator().manual_seed(42)
|
||||
lora_a = torch.randn((_LORA_RANK, _LORA_DIM), dtype=torch.float32, generator=gen) * 0.1
|
||||
lora_b = torch.randn((_LORA_DIM, _LORA_RANK), dtype=torch.float32, generator=gen) * 0.5
|
||||
save_file(
|
||||
{
|
||||
f"base_model.model.{_LORA_MODULE}.lora_A.weight": lora_a,
|
||||
f"base_model.model.{_LORA_MODULE}.lora_B.weight": lora_b,
|
||||
},
|
||||
str(adapter_dir / "adapter_model.safetensors"),
|
||||
)
|
||||
(adapter_dir / "adapter_config.json").write_text(
|
||||
json.dumps({"r": _LORA_RANK, "lora_alpha": _LORA_RANK, "target_modules": [_LORA_MODULE]}),
|
||||
encoding="utf-8",
|
||||
)
|
||||
lora_dir = str(adapter_dir)
|
||||
return LoRARequest(lora_name="test_file", lora_int_id=stable_lora_int_id(lora_dir), lora_path=lora_dir)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Test
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.full_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_sensenova_u1_lora_scale_and_deactivation(tmp_path) -> None:
|
||||
"""Validate LoRA effect, scale sensitivity, bounded perturbation, and clean deactivation."""
|
||||
lora_request = _make_file_lora_request(tmp_path / "sensenova_lora")
|
||||
|
||||
with OmniRunner(MODEL, stage_configs_path=None) as runner:
|
||||
omni = runner.omni
|
||||
baseline = _generate(omni)
|
||||
img_1x = _generate(omni, lora_request, lora_scale=1.0)
|
||||
img_2x = _generate(omni, lora_request, lora_scale=2.0)
|
||||
restored = _generate(omni)
|
||||
|
||||
baseline_arr = np.array(baseline, dtype=np.int16)
|
||||
diff_1x = np.abs(baseline_arr - np.array(img_1x, dtype=np.int16)).mean()
|
||||
diff_2x = np.abs(baseline_arr - np.array(img_2x, dtype=np.int16)).mean()
|
||||
diff_restored = np.abs(baseline_arr - np.array(restored, dtype=np.int16)).mean()
|
||||
|
||||
# (a) Adapter has visible effect at both scales
|
||||
assert diff_1x > 0.5, f"LoRA scale=1.0 had no visible effect: diff={diff_1x}"
|
||||
assert diff_2x > 0.5, f"LoRA scale=2.0 had no visible effect: diff={diff_2x}"
|
||||
|
||||
# (b) Different scales produce different outputs
|
||||
assert not np.isclose(diff_1x, diff_2x, atol=1.0), (
|
||||
f"LoRA scale has no effect: diff_1x={diff_1x:.2f}, diff_2x={diff_2x:.2f}"
|
||||
)
|
||||
|
||||
# (c) Output is not corrupted
|
||||
assert diff_1x < 80, f"LoRA output looks corrupted: diff_1x={diff_1x}"
|
||||
assert diff_2x < 120, f"LoRA output looks corrupted: diff_2x={diff_2x}"
|
||||
|
||||
# (d) Deactivation fully restores base model
|
||||
assert diff_restored == 0.0, f"Base model not restored after LoRA deactivation: diff={diff_restored}"
|
||||
@@ -0,0 +1,303 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E tests for Step-Audio2 model with audio input and audio output.
|
||||
|
||||
These tests verify:
|
||||
1. Processor loads correctly and handles audio input
|
||||
2. Field configuration for vLLM batching (flat_from_sizes)
|
||||
3. Tensor shapes are correct (3D not 4D for profiling)
|
||||
4. Basic E2E inference works with audio input
|
||||
"""
|
||||
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.helpers.runtime import OmniRunner
|
||||
|
||||
models = ["stepfun-ai/Step-Audio-2-mini"]
|
||||
|
||||
stage_configs = [str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")]
|
||||
|
||||
test_params = [(model, stage_config) for model in models for stage_config in stage_configs]
|
||||
|
||||
|
||||
def create_dummy_audio(sample_rate: int = 16000, duration_sec: float = 1.0) -> tuple[np.ndarray, int]:
|
||||
"""Create a dummy audio signal (sine wave) for testing."""
|
||||
t = np.linspace(0, duration_sec, int(sample_rate * duration_sec), dtype=np.float32)
|
||||
audio = 0.5 * np.sin(2 * np.pi * 440 * t)
|
||||
return audio, sample_rate
|
||||
|
||||
|
||||
class StepAudio2OmniRunner(OmniRunner):
|
||||
"""Test runner specialized for Step-Audio2 model."""
|
||||
|
||||
def get_omni_inputs(
|
||||
self,
|
||||
prompts: list[str] | str,
|
||||
system_prompt: str | None = None,
|
||||
audios=None,
|
||||
**kwargs,
|
||||
) -> list[dict]:
|
||||
"""Construct Omni input format for Step-Audio2."""
|
||||
if system_prompt is None:
|
||||
system_prompt = "你是一个语音对话助手,能够理解音频输入并生成语音回复。"
|
||||
|
||||
audio_placeholder = "<audio_patch>"
|
||||
|
||||
if isinstance(prompts, str):
|
||||
prompts = [prompts]
|
||||
|
||||
num_prompts = len(prompts)
|
||||
|
||||
if audios is None:
|
||||
audios_list = [None] * num_prompts
|
||||
elif isinstance(audios, list):
|
||||
audios_list = audios if len(audios) == num_prompts else [audios[0]] * num_prompts
|
||||
else:
|
||||
audios_list = [audios] * num_prompts
|
||||
|
||||
omni_inputs = []
|
||||
for i, prompt_text in enumerate(prompts):
|
||||
user_content = ""
|
||||
multi_modal_data = {}
|
||||
|
||||
audio = audios_list[i]
|
||||
if audio is not None:
|
||||
user_content += audio_placeholder
|
||||
multi_modal_data["audio"] = audio
|
||||
|
||||
user_content += prompt_text
|
||||
|
||||
full_prompt = (
|
||||
f"<|im_start|>system\n{system_prompt}<|im_end|>\n"
|
||||
f"<|im_start|>user\n{user_content}<|im_end|>\n"
|
||||
f"<|im_start|>assistant\n"
|
||||
)
|
||||
|
||||
input_dict = {"prompt": full_prompt}
|
||||
if multi_modal_data:
|
||||
input_dict["multi_modal_data"] = multi_modal_data
|
||||
|
||||
omni_inputs.append(input_dict)
|
||||
|
||||
return omni_inputs
|
||||
|
||||
|
||||
@pytest.fixture(scope="session")
|
||||
def step_audio2_runner():
|
||||
return StepAudio2OmniRunner
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_processor_loads_correctly():
|
||||
"""Test that Step-Audio2 processor can be imported and initialized."""
|
||||
from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
|
||||
StepAudio2MultiModalProcessor,
|
||||
StepAudio2ProcessingInfo,
|
||||
StepAudio2Processor,
|
||||
)
|
||||
|
||||
assert StepAudio2MultiModalProcessor is not None
|
||||
assert StepAudio2ProcessingInfo is not None
|
||||
assert StepAudio2Processor is not None
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_audio_preprocessing():
|
||||
"""Test that audio preprocessing produces correct tensor shapes."""
|
||||
from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
|
||||
log_mel_spectrogram,
|
||||
padding_mels,
|
||||
)
|
||||
|
||||
audio, sample_rate = create_dummy_audio(sample_rate=16000, duration_sec=1.0)
|
||||
|
||||
mel = log_mel_spectrogram(audio)
|
||||
|
||||
assert mel.ndim == 2, f"Expected 2D tensor, got {mel.ndim}D"
|
||||
assert mel.shape[0] == 128, f"Expected 128 mel bins, got {mel.shape[0]}"
|
||||
|
||||
mels = [mel, mel]
|
||||
padded_mels, mel_lens = padding_mels(mels)
|
||||
|
||||
assert padded_mels.ndim == 3, f"Expected 3D tensor, got {padded_mels.ndim}D"
|
||||
assert padded_mels.shape[0] == 2, f"Expected batch size 2, got {padded_mels.shape[0]}"
|
||||
assert padded_mels.shape[1] == 128, f"Expected 128 mel bins, got {padded_mels.shape[1]}"
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_feature_length_calculation():
|
||||
"""Test that audio feature length calculation is correct."""
|
||||
from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
|
||||
calculate_audio_feature_length,
|
||||
)
|
||||
|
||||
assert calculate_audio_feature_length(1000) == 125
|
||||
|
||||
assert calculate_audio_feature_length(100) > 0
|
||||
|
||||
assert calculate_audio_feature_length(1) >= 1
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_token_config_constants():
|
||||
"""Test that token configuration constants are set correctly."""
|
||||
from vllm_omni.model_executor.models.step_audio2.step_audio2_constants import (
|
||||
DEFAULT_TOKEN_CONFIG,
|
||||
STEP_AUDIO2_AUDIO_END,
|
||||
STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID,
|
||||
STEP_AUDIO2_AUDIO_START,
|
||||
STEP_AUDIO2_AUDIO_VOCAB_SIZE,
|
||||
STEP_AUDIO2_TEXT_MAX,
|
||||
)
|
||||
|
||||
assert STEP_AUDIO2_TEXT_MAX < STEP_AUDIO2_AUDIO_START, "Text tokens should come before audio tokens"
|
||||
assert STEP_AUDIO2_AUDIO_END == STEP_AUDIO2_AUDIO_START + STEP_AUDIO2_AUDIO_VOCAB_SIZE - 1
|
||||
assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID > STEP_AUDIO2_TEXT_MAX
|
||||
assert STEP_AUDIO2_AUDIO_PATCH_TOKEN_ID < STEP_AUDIO2_AUDIO_START
|
||||
|
||||
assert DEFAULT_TOKEN_CONFIG.text_max == STEP_AUDIO2_TEXT_MAX
|
||||
assert DEFAULT_TOKEN_CONFIG.audio_start == STEP_AUDIO2_AUDIO_START
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_mm_field_config_structure():
|
||||
"""Test that multimodal field config is properly structured for vLLM batching."""
|
||||
from vllm.multimodal.inputs import MultiModalFieldConfig
|
||||
|
||||
audio_lens = torch.tensor([100, 200, 150], dtype=torch.int32)
|
||||
|
||||
field_config = MultiModalFieldConfig.flat_from_sizes("audio", audio_lens)
|
||||
|
||||
assert field_config is not None
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_audio_encoder_output_shape():
|
||||
"""Test that audio encoder produces correct output shapes."""
|
||||
from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
|
||||
Adaptor,
|
||||
AudioEncoder,
|
||||
)
|
||||
|
||||
encoder = AudioEncoder(n_mels=128, n_ctx=1500, n_state=512, n_head=8, n_layer=6)
|
||||
|
||||
batch_size = 2
|
||||
n_mels = 128
|
||||
time_steps = 400
|
||||
x = torch.randn(batch_size, n_mels, time_steps)
|
||||
x_len = torch.tensor([time_steps, time_steps // 2], dtype=torch.int32)
|
||||
|
||||
encoded, encoded_lens = encoder(x, x_len)
|
||||
|
||||
assert encoded.ndim == 3, f"Expected 3D tensor, got {encoded.ndim}D"
|
||||
assert encoded.shape[0] == batch_size
|
||||
assert encoded.shape[2] == 512
|
||||
|
||||
adapter = Adaptor(n_state=512, n_hidden=4096, kernel_size=3, stride=2)
|
||||
adapted = adapter(encoded)
|
||||
|
||||
assert adapted.ndim == 3
|
||||
assert adapted.shape[0] == batch_size
|
||||
assert adapted.shape[2] == 4096
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_token_separation():
|
||||
"""Test that token separation works correctly."""
|
||||
from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
|
||||
StepAudio2ThinkerForConditionalGeneration,
|
||||
)
|
||||
|
||||
token_ids = [100, 200, 151700, 300, 151800, 400]
|
||||
|
||||
text_tokens, audio_tokens = StepAudio2ThinkerForConditionalGeneration.separate_tokens(token_ids)
|
||||
|
||||
assert text_tokens == [100, 200, 300, 400]
|
||||
assert audio_tokens == [4, 104]
|
||||
|
||||
|
||||
@pytest.mark.core_model
|
||||
def test_has_audio_output():
|
||||
"""Test detection of audio tokens in output."""
|
||||
from vllm_omni.model_executor.models.step_audio2.step_audio2_thinker import (
|
||||
StepAudio2ThinkerForConditionalGeneration,
|
||||
)
|
||||
|
||||
text_only = [100, 200, 300]
|
||||
assert not StepAudio2ThinkerForConditionalGeneration.has_audio_output(text_only)
|
||||
|
||||
with_audio = [100, 200, 151700, 300]
|
||||
assert StepAudio2ThinkerForConditionalGeneration.has_audio_output(with_audio)
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_audio_to_text_and_audio(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None:
|
||||
"""Test processing audio input and generating text + audio output."""
|
||||
model, stage_config_path = test_config
|
||||
|
||||
with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
|
||||
audio = create_dummy_audio(sample_rate=16000, duration_sec=2.0)
|
||||
|
||||
outputs = runner.generate_multimodal(
|
||||
prompts="请复述这段音频的内容。",
|
||||
audios=audio,
|
||||
)
|
||||
|
||||
assert len(outputs) > 0
|
||||
|
||||
text_output = None
|
||||
for stage_output in outputs:
|
||||
if stage_output.final_output_type == "text":
|
||||
text_output = stage_output
|
||||
break
|
||||
|
||||
assert text_output is not None
|
||||
assert text_output.request_output is not None
|
||||
text_content = text_output.request_output.outputs[0].text
|
||||
assert text_content is not None
|
||||
assert len(text_content.strip()) > 0
|
||||
|
||||
audio_output = None
|
||||
for stage_output in outputs:
|
||||
if stage_output.final_output_type == "audio":
|
||||
audio_output = stage_output
|
||||
break
|
||||
|
||||
if audio_output is not None:
|
||||
assert audio_output.request_output is not None
|
||||
audio_tensor = audio_output.request_output.outputs[0].multimodal_output.get("audio")
|
||||
if audio_tensor is not None:
|
||||
assert audio_tensor.numel() > 0
|
||||
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.parametrize("test_config", test_params)
|
||||
def test_text_only_input(step_audio2_runner: type[StepAudio2OmniRunner], test_config: tuple[str, str]) -> None:
|
||||
"""Test processing text-only input (no audio)."""
|
||||
model, stage_config_path = test_config
|
||||
|
||||
with step_audio2_runner(model, seed=42, stage_configs_path=stage_config_path) as runner:
|
||||
outputs = runner.generate_multimodal(
|
||||
prompts="你好,请用中文回答。",
|
||||
audios=None,
|
||||
)
|
||||
|
||||
assert len(outputs) > 0
|
||||
|
||||
text_output = None
|
||||
for stage_output in outputs:
|
||||
if stage_output.final_output_type == "text":
|
||||
text_output = stage_output
|
||||
break
|
||||
|
||||
assert text_output is not None
|
||||
assert text_output.request_output is not None
|
||||
text_content = text_output.request_output.outputs[0].text
|
||||
assert text_content is not None
|
||||
assert len(text_content.strip()) > 0
|
||||
+2
-1
@@ -1,6 +1,7 @@
|
||||
"""E2E test for VoxCPM2 native AR offline inference."""
|
||||
|
||||
import os
|
||||
from collections.abc import Mapping
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
@@ -23,7 +24,7 @@ def voxcpm2_engine():
|
||||
|
||||
def _extract_audio(multimodal_output: dict) -> torch.Tensor:
|
||||
"""Extract the final complete audio tensor from multimodal output."""
|
||||
assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}"
|
||||
assert isinstance(multimodal_output, (dict, Mapping)), f"Expected dict/Mapping, got {type(multimodal_output)}"
|
||||
|
||||
# Output processor accumulates per-step audio chunks under "audio".
|
||||
audio = multimodal_output.get("audio")
|
||||
@@ -0,0 +1,81 @@
|
||||
# Stage config for Step-Audio2 CI tests
|
||||
# Stage 0: Thinker (audio understanding → text + audio tokens)
|
||||
# Stage 1: Token2Wav (audio tokens → waveform)
|
||||
|
||||
# The following config has been verified on 2x 24GB GPU (L4/RTX3090/RTX4090).
|
||||
stage_args:
|
||||
- stage_id: 0
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
model_stage: thinker
|
||||
model_arch: StepAudio2ThinkerForConditionalGeneration
|
||||
tensor_parallel_size: 1
|
||||
hf_overrides:
|
||||
architectures: ["StepAudio2ThinkerForConditionalGeneration"]
|
||||
|
||||
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler
|
||||
max_model_len: 1024
|
||||
max_num_batched_tokens: 1024
|
||||
max_num_seqs: 1
|
||||
gpu_memory_utilization: 0.7
|
||||
skip_mm_profiling: true
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
engine_output_type: text # Output text + audio tokens
|
||||
is_comprehension: true
|
||||
final_output: true
|
||||
final_output_type: text
|
||||
default_sampling_params:
|
||||
temperature: 0.7
|
||||
top_p: 0.9
|
||||
top_k: -1
|
||||
max_tokens: 256
|
||||
seed: 42
|
||||
detokenize: true
|
||||
repetition_penalty: 1.05
|
||||
|
||||
- stage_id: 1
|
||||
runtime:
|
||||
process: true
|
||||
devices: "0"
|
||||
max_batch_size: 1
|
||||
engine_args:
|
||||
|
||||
model_stage: token2wav
|
||||
model_arch: StepAudio2Token2WavModel
|
||||
worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker
|
||||
scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler
|
||||
hf_overrides:
|
||||
architectures: ["StepAudio2Token2WavModel"]
|
||||
gpu_memory_utilization: 0.2
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
max_num_batched_tokens: 1024
|
||||
engine_output_type: audio # Output waveform
|
||||
engine_input_source: [0]
|
||||
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.step_audio2.thinker2token2wav
|
||||
final_output: true
|
||||
final_output_type: audio
|
||||
default_sampling_params:
|
||||
temperature: 0.0
|
||||
top_p: 1.0
|
||||
top_k: -1
|
||||
max_tokens: 1
|
||||
seed: 42
|
||||
detokenize: false
|
||||
|
||||
runtime:
|
||||
enabled: true
|
||||
defaults:
|
||||
window_size: -1
|
||||
max_inflight: 1
|
||||
edges:
|
||||
- from: 0
|
||||
to: 1
|
||||
window_size: -1
|
||||
@@ -0,0 +1,92 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""
|
||||
E2E online serving tests for Ming-omni-tts (dense) model.
|
||||
Tests text-to-audio via /v1/audio/speech endpoint.
|
||||
"""
|
||||
|
||||
import os
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
import pytest
|
||||
|
||||
from tests.helpers.mark import hardware_test
|
||||
from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
pytestmark = [pytest.mark.advanced_model, pytest.mark.omni]
|
||||
|
||||
MODEL = "inclusionAI/Ming-omni-tts-0.5B"
|
||||
DEPLOY_CONFIG = get_deploy_config_path("ming_tts.yaml")
|
||||
|
||||
no_async_chunk_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=DEPLOY_CONFIG,
|
||||
server_args=["--enforce-eager", "--no-async-chunk"],
|
||||
),
|
||||
id="no_async_chunk",
|
||||
)
|
||||
]
|
||||
|
||||
async_chunk_params = [
|
||||
pytest.param(
|
||||
OmniServerParams(
|
||||
model=MODEL,
|
||||
stage_config_path=DEPLOY_CONFIG,
|
||||
server_args=["--enforce-eager"],
|
||||
),
|
||||
id="async_chunk",
|
||||
)
|
||||
]
|
||||
|
||||
|
||||
def get_prompt(prompt_type="zh"):
|
||||
prompts = {
|
||||
"zh": "今天天气真不错,适合出去散散步。",
|
||||
"zh_short": "这款产品的名字,叫变态坑爹牛肉丸。",
|
||||
}
|
||||
return prompts.get(prompt_type, prompts["zh"])
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", no_async_chunk_params, indirect=True)
|
||||
def test_text_to_audio_non_streaming_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Deploy Setting: ming_tts.yaml with --no-async-chunk
|
||||
Input Modal: text
|
||||
Output Modal: audio
|
||||
Input Setting: stream=False
|
||||
Datasets: two concurrent requests
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": get_prompt("zh"),
|
||||
"stream": False,
|
||||
"response_format": "wav",
|
||||
"timeout": 300.0,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config, request_num=2)
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", async_chunk_params, indirect=True)
|
||||
def test_text_to_audio_streaming_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Deploy Setting: ming_tts.yaml (async_chunk=true)
|
||||
Input Modal: text + voice
|
||||
Output Modal: audio (streamed)
|
||||
Input Setting: stream=True
|
||||
Datasets: single request
|
||||
"""
|
||||
request_config = {
|
||||
"model": omni_server.model,
|
||||
"input": get_prompt("zh_short"),
|
||||
"voice": "灵小甄",
|
||||
"stream": True,
|
||||
"response_format": "wav",
|
||||
"timeout": 300.0,
|
||||
}
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
@@ -25,11 +25,11 @@ from tests.helpers.runtime import OmniServerParams
|
||||
from tests.helpers.stage_config import get_deploy_config_path
|
||||
|
||||
# ``omni`` for all tests; 001/002 also get ``full_model`` via ``TestMossTtsNanoFull`` (003 is core+advanced only, no full_model).
|
||||
pytestmark = [pytest.mark.full_model, pytest.mark.tts]
|
||||
|
||||
_SKIP_ISSUE_3168 = pytest.mark.skip(
|
||||
reason="https://github.com/vllm-project/vllm-omni/issues/3168",
|
||||
)
|
||||
pytestmark = [
|
||||
pytest.mark.skip(reason="issue#4361"),
|
||||
pytest.mark.full_model,
|
||||
pytest.mark.tts,
|
||||
]
|
||||
|
||||
MODEL = "OpenMOSS-Team/MOSS-TTS-Nano"
|
||||
REF_AUDIO_URL = "https://raw.githubusercontent.com/OpenMOSS/MOSS-TTS-Nano/main/assets/audio/zh_1.wav"
|
||||
|
||||
@@ -0,0 +1,368 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""
|
||||
E2E tests for Step-Audio2 online serving with concurrent requests.
|
||||
|
||||
Tests:
|
||||
1. Single request via API
|
||||
2. Concurrent throughput (1/2/4/8 concurrent requests)
|
||||
3. Latency percentiles (p50/p95/p99)
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import os
|
||||
import socket
|
||||
import subprocess
|
||||
import sys
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import pytest
|
||||
from vllm.utils.network_utils import get_open_port
|
||||
|
||||
MODEL = "stepfun-ai/Step-Audio-2-mini"
|
||||
STAGE_CONFIG = str(Path(__file__).parent / "stage_configs" / "step_audio2_ci.yaml")
|
||||
|
||||
os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
|
||||
|
||||
class OmniServer:
|
||||
"""Omniserver for vLLM-Omni tests."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: str,
|
||||
serve_args: list[str],
|
||||
*,
|
||||
env_dict: dict[str, str] | None = None,
|
||||
) -> None:
|
||||
self.model = model
|
||||
self.serve_args = serve_args
|
||||
self.env_dict = env_dict
|
||||
self.proc: subprocess.Popen | None = None
|
||||
self.host = "127.0.0.1"
|
||||
self.port = get_open_port()
|
||||
|
||||
def _start_server(self) -> None:
|
||||
"""Start the vLLM-Omni server subprocess."""
|
||||
env = os.environ.copy()
|
||||
env["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn"
|
||||
if self.env_dict is not None:
|
||||
env.update(self.env_dict)
|
||||
|
||||
cmd = [
|
||||
sys.executable,
|
||||
"-m",
|
||||
"vllm_omni.entrypoints.cli.main",
|
||||
"serve",
|
||||
self.model,
|
||||
"--omni",
|
||||
"--host",
|
||||
self.host,
|
||||
"--port",
|
||||
str(self.port),
|
||||
] + self.serve_args
|
||||
|
||||
print(f"Launching OmniServer with: {' '.join(cmd)}")
|
||||
self.proc = subprocess.Popen(
|
||||
cmd,
|
||||
env=env,
|
||||
cwd=os.path.dirname(os.path.dirname(os.path.abspath(__file__))),
|
||||
)
|
||||
|
||||
max_wait = 600 # 10 minutes
|
||||
start_time = time.time()
|
||||
while time.time() - start_time < max_wait:
|
||||
try:
|
||||
with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
|
||||
sock.settimeout(1)
|
||||
result = sock.connect_ex((self.host, self.port))
|
||||
if result == 0:
|
||||
print(f"Server ready on {self.host}:{self.port}")
|
||||
return
|
||||
except Exception:
|
||||
pass
|
||||
time.sleep(2)
|
||||
|
||||
raise RuntimeError(f"Server failed to start within {max_wait} seconds")
|
||||
|
||||
def __enter__(self):
|
||||
self._start_server()
|
||||
return self
|
||||
|
||||
def __exit__(self, exc_type, exc_val, exc_tb):
|
||||
if self.proc:
|
||||
self.proc.terminate()
|
||||
try:
|
||||
self.proc.wait(timeout=30)
|
||||
except subprocess.TimeoutExpired:
|
||||
self.proc.kill()
|
||||
self.proc.wait()
|
||||
|
||||
|
||||
def create_dummy_audio_base64(duration_sec: float = 5.0, sample_rate: int = 16000) -> str:
|
||||
"""Create dummy audio as base64."""
|
||||
import base64
|
||||
import io
|
||||
import wave
|
||||
|
||||
num_samples = int(duration_sec * sample_rate)
|
||||
t = np.linspace(0, duration_sec, num_samples, dtype=np.float32)
|
||||
audio = (0.5 * np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
|
||||
|
||||
buffer = io.BytesIO()
|
||||
with wave.open(buffer, "wb") as wav_file:
|
||||
wav_file.setnchannels(1)
|
||||
wav_file.setsampwidth(2)
|
||||
wav_file.setframerate(sample_rate)
|
||||
wav_file.writeframes(audio.tobytes())
|
||||
|
||||
return base64.b64encode(buffer.getvalue()).decode("utf-8")
|
||||
|
||||
|
||||
def create_s2st_request(audio_base64: str) -> dict:
|
||||
"""Create S2ST request payload."""
|
||||
# Use audio_url format with base64 data URL
|
||||
audio_url = f"data:audio/wav;base64,{audio_base64}"
|
||||
return {
|
||||
"model": MODEL,
|
||||
"messages": [
|
||||
{"role": "system", "content": "请仔细聆听这段语音,然后复述其内容。"},
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{
|
||||
"type": "audio_url",
|
||||
"audio_url": {"url": audio_url},
|
||||
},
|
||||
{
|
||||
"type": "text",
|
||||
"text": "<tts_start>",
|
||||
},
|
||||
],
|
||||
},
|
||||
],
|
||||
"max_tokens": 512,
|
||||
"temperature": 0.7,
|
||||
"top_p": 0.9,
|
||||
}
|
||||
|
||||
|
||||
async def send_request(session: aiohttp.ClientSession, url: str, payload: dict) -> tuple[float, int, str]:
|
||||
"""Send single request and return (latency, num_tokens, response_text)."""
|
||||
start = time.perf_counter()
|
||||
async with session.post(url, json=payload) as response:
|
||||
result = await response.json()
|
||||
latency = time.perf_counter() - start
|
||||
|
||||
if "choices" in result and len(result["choices"]) > 0:
|
||||
text = result["choices"][0].get("message", {}).get("content", "")
|
||||
tokens = result.get("usage", {}).get("completion_tokens", len(text.split()))
|
||||
else:
|
||||
text = str(result)
|
||||
tokens = 0
|
||||
|
||||
return latency, tokens, text
|
||||
|
||||
|
||||
async def benchmark_concurrent(
|
||||
url: str,
|
||||
payload: dict,
|
||||
num_concurrent: int,
|
||||
num_requests: int,
|
||||
warmup: int = 2,
|
||||
) -> dict:
|
||||
"""Run concurrent benchmark with true sustained concurrency.
|
||||
|
||||
Uses semaphore to maintain exactly num_concurrent requests in flight at all times,
|
||||
rather than batch-and-wait approach. Measures real wall clock time.
|
||||
"""
|
||||
semaphore = asyncio.Semaphore(num_concurrent)
|
||||
|
||||
async def bounded_request(session: aiohttp.ClientSession):
|
||||
async with semaphore:
|
||||
return await send_request(session, url, payload)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
if warmup > 0:
|
||||
print(f" Warming up ({warmup} requests)...", end=" ", flush=True)
|
||||
for _ in range(warmup):
|
||||
try:
|
||||
await send_request(session, url, payload)
|
||||
except Exception as e:
|
||||
print(f"Warmup failed: {e}")
|
||||
print("done")
|
||||
|
||||
print(f" Running {num_requests} requests with concurrency={num_concurrent}...", end=" ", flush=True)
|
||||
wall_start = time.perf_counter()
|
||||
|
||||
tasks = [bounded_request(session) for _ in range(num_requests)]
|
||||
results = await asyncio.gather(*tasks, return_exceptions=True)
|
||||
|
||||
wall_time = time.perf_counter() - wall_start
|
||||
print(f"done ({wall_time:.2f}s)")
|
||||
|
||||
all_latencies = []
|
||||
all_tokens = []
|
||||
errors = 0
|
||||
|
||||
for result in results:
|
||||
if isinstance(result, Exception):
|
||||
print(f" Request failed: {result}")
|
||||
errors += 1
|
||||
continue
|
||||
latency, tokens, _ = result
|
||||
all_latencies.append(latency)
|
||||
all_tokens.append(tokens)
|
||||
|
||||
if not all_latencies:
|
||||
return {"error": "All requests failed"}
|
||||
|
||||
latencies = np.array(all_latencies)
|
||||
total_tokens = sum(all_tokens)
|
||||
successful_requests = len(all_latencies)
|
||||
|
||||
return {
|
||||
"num_concurrent": num_concurrent,
|
||||
"num_requests": num_requests,
|
||||
"successful_requests": successful_requests,
|
||||
"errors": errors,
|
||||
"wall_time": wall_time,
|
||||
"p50_latency": float(np.percentile(latencies, 50)),
|
||||
"p95_latency": float(np.percentile(latencies, 95)),
|
||||
"p99_latency": float(np.percentile(latencies, 99)),
|
||||
"mean_latency": float(np.mean(latencies)),
|
||||
"min_latency": float(np.min(latencies)),
|
||||
"max_latency": float(np.max(latencies)),
|
||||
"throughput_req_per_sec": successful_requests / wall_time,
|
||||
"throughput_tokens_per_sec": total_tokens / wall_time,
|
||||
}
|
||||
|
||||
|
||||
class TestStepAudio2OnlineServing:
|
||||
"""Online serving tests for Step-Audio2."""
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def server(self):
|
||||
"""Start vLLM-Omni server."""
|
||||
with OmniServer(MODEL, ["--stage-configs-path", STAGE_CONFIG]) as omni_server:
|
||||
yield f"http://{omni_server.host}:{omni_server.port}/v1/chat/completions"
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_single_request(self, server):
|
||||
"""Test single request works."""
|
||||
audio_base64 = create_dummy_audio_base64(duration_sec=2.0)
|
||||
payload = create_s2st_request(audio_base64)
|
||||
|
||||
async with aiohttp.ClientSession() as session:
|
||||
latency, tokens, text = await send_request(session, server, payload)
|
||||
|
||||
assert latency > 0
|
||||
assert tokens >= 0
|
||||
print(f"Single request: latency={latency:.2f}s, tokens={tokens}")
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize("num_concurrent", [1, 2, 4])
|
||||
async def test_concurrent_throughput(self, server, num_concurrent):
|
||||
"""Test concurrent throughput at different concurrency levels."""
|
||||
audio_base64 = create_dummy_audio_base64(duration_sec=5.0)
|
||||
payload = create_s2st_request(audio_base64)
|
||||
|
||||
num_requests = num_concurrent * 3
|
||||
|
||||
stats = await benchmark_concurrent(server, payload, num_concurrent, num_requests)
|
||||
|
||||
assert "error" not in stats
|
||||
print(f"\nConcurrency={num_concurrent}:")
|
||||
print(f" p50={stats['p50_latency']:.2f}s, p95={stats['p95_latency']:.2f}s")
|
||||
print(f" throughput={stats['throughput_req_per_sec']:.2f} req/s")
|
||||
|
||||
|
||||
async def run_throughput_benchmark(
|
||||
server_url: str,
|
||||
audio_duration: float = 5.0,
|
||||
concurrency_levels: list[int] = [1, 2, 4, 8],
|
||||
requests_per_level: int = 10,
|
||||
warmup: int = 2,
|
||||
):
|
||||
"""Run full throughput benchmark."""
|
||||
audio_base64 = create_dummy_audio_base64(duration_sec=audio_duration)
|
||||
payload = create_s2st_request(audio_base64)
|
||||
|
||||
print(f"{'=' * 70}")
|
||||
print("Step-Audio2 Online Serving Throughput Benchmark")
|
||||
print(f"{'=' * 70}")
|
||||
print(f"Server: {server_url}")
|
||||
print(f"Audio duration: {audio_duration}s")
|
||||
print(f"Requests per level: {requests_per_level}")
|
||||
print(f"Warmup requests: {warmup}")
|
||||
print(f"Concurrency levels: {concurrency_levels}")
|
||||
print(f"{'=' * 70}")
|
||||
|
||||
results = []
|
||||
for num_concurrent in concurrency_levels:
|
||||
print(f"\n[Concurrency={num_concurrent}]")
|
||||
stats = await benchmark_concurrent(server_url, payload, num_concurrent, requests_per_level, warmup)
|
||||
results.append(stats)
|
||||
|
||||
if "error" not in stats:
|
||||
print(
|
||||
f" Result: p50={stats['p50_latency']:.2f}s, "
|
||||
f"p95={stats['p95_latency']:.2f}s, "
|
||||
f"throughput={stats['throughput_req_per_sec']:.2f} req/s, "
|
||||
f"{stats['throughput_tokens_per_sec']:.1f} tok/s"
|
||||
)
|
||||
|
||||
print(f"\n{'=' * 80}")
|
||||
print("SUMMARY")
|
||||
print(f"{'=' * 80}")
|
||||
print(
|
||||
f"{'Conc':<6} {'Reqs':<6} {'Err':<5} {'Wall(s)':<9} {'p50':<8} {'p95':<8} {'p99':<8} {'Req/s':<8} {'Tok/s':<8}"
|
||||
)
|
||||
print("-" * 80)
|
||||
|
||||
for stats in results:
|
||||
if "error" in stats:
|
||||
print(f"{stats.get('num_concurrent', '?'):<6} ERROR: {stats['error']}")
|
||||
continue
|
||||
print(
|
||||
f"{stats['num_concurrent']:<6} "
|
||||
f"{stats['successful_requests']:<6} "
|
||||
f"{stats['errors']:<5} "
|
||||
f"{stats['wall_time']:<9.2f} "
|
||||
f"{stats['p50_latency']:<8.2f} "
|
||||
f"{stats['p95_latency']:<8.2f} "
|
||||
f"{stats['p99_latency']:<8.2f} "
|
||||
f"{stats['throughput_req_per_sec']:<8.2f} "
|
||||
f"{stats['throughput_tokens_per_sec']:<8.1f}"
|
||||
)
|
||||
|
||||
print("=" * 80)
|
||||
return results
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Step-Audio2 Online Serving Benchmark")
|
||||
parser.add_argument(
|
||||
"--server-url", type=str, default="http://localhost:8000/v1/chat/completions", help="vLLM server URL"
|
||||
)
|
||||
parser.add_argument("--audio-duration", type=float, default=5.0, help="Input audio duration in seconds")
|
||||
parser.add_argument("--concurrency", type=str, default="1,2,4,8", help="Comma-separated concurrency levels to test")
|
||||
parser.add_argument("--requests-per-level", type=int, default=10, help="Number of requests per concurrency level")
|
||||
parser.add_argument("--warmup", type=int, default=2, help="Number of warmup requests before each benchmark")
|
||||
|
||||
args = parser.parse_args()
|
||||
concurrency_levels = [int(x) for x in args.concurrency.split(",")]
|
||||
|
||||
asyncio.run(
|
||||
run_throughput_benchmark(
|
||||
args.server_url,
|
||||
args.audio_duration,
|
||||
concurrency_levels,
|
||||
args.requests_per_level,
|
||||
args.warmup,
|
||||
)
|
||||
)
|
||||
@@ -73,6 +73,7 @@ def test_voice_clone_streaming_001(omni_server, openai_client) -> None:
|
||||
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
@pytest.mark.skip(reason="#4335")
|
||||
def test_response_format_001(omni_server, openai_client) -> None:
|
||||
"""
|
||||
Test voice-clone non-stream PCM output via OpenAI API.
|
||||
|
||||
@@ -91,3 +91,90 @@ class TestVoxtralTTSFixedVoice:
|
||||
"timeout": 120.0,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_speed(self, omni_server, openai_client) -> None:
|
||||
"""Request with speed parameters"""
|
||||
speeds = [0.5, 1, 1.5, 2, 2.5]
|
||||
for speed in speeds:
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "The boy was there when the sun rose.",
|
||||
"voice": "casual_female",
|
||||
"language": "English",
|
||||
"response_format": "wav",
|
||||
"timeout": 120.0,
|
||||
"speed": speed,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_instructions(self, omni_server, openai_client) -> None:
|
||||
"""Request with instructions parameters"""
|
||||
instructions = [
|
||||
"Speak formally",
|
||||
"Speak angrily",
|
||||
"Deliver with a sad voice",
|
||||
"Speak with a chirpy happy voice",
|
||||
]
|
||||
for instruction in instructions:
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "The boy was there when the sun rose.",
|
||||
"voice": "casual_female",
|
||||
"language": "English",
|
||||
"response_format": "wav",
|
||||
"timeout": 120.0,
|
||||
"instructions": instruction,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_response_formats(self, omni_server, openai_client) -> None:
|
||||
"""Test TTS with different response formats"""
|
||||
response_formats = ["wav", "mp3"]
|
||||
for response_format in response_formats:
|
||||
openai_client.send_audio_speech_request(
|
||||
{
|
||||
"model": omni_server.model,
|
||||
"input": "Testing various response formats.",
|
||||
"voice": "casual_male",
|
||||
"language": "English",
|
||||
"response_format": response_format,
|
||||
"timeout": 120.0,
|
||||
}
|
||||
)
|
||||
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=1)
|
||||
def test_speech_batches(self, omni_server, openai_client) -> None:
|
||||
"""Test TTS batches"""
|
||||
items = [
|
||||
{"input": "The birch canoe slid on the smooth planks."},
|
||||
{"input": "Glue the sheet to the dark blue background."},
|
||||
{"input": "It's easy to tell the depth of a well."},
|
||||
{"input": "These days a chicken leg is a rare dish."},
|
||||
{"input": "Rice is often served in round bowls."},
|
||||
]
|
||||
|
||||
openai_client.send_audio_speech_batch_http_request(
|
||||
{
|
||||
"json": {
|
||||
"model": omni_server.model,
|
||||
"items": items,
|
||||
"voice": "casual_male",
|
||||
"language": "English",
|
||||
"response_format": "wav",
|
||||
},
|
||||
"timeout": 120.0,
|
||||
}
|
||||
)
|
||||
@@ -0,0 +1,59 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for _ensure_tensor_values boundary enforcement."""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.worker.gpu_ar_model_runner import _ensure_tensor_values
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_tensors_pass_through():
|
||||
"""Tensor values pass through unchanged."""
|
||||
t = torch.randn(10)
|
||||
result = _ensure_tensor_values({"audio": t})
|
||||
assert torch.equal(result["audio"], t)
|
||||
|
||||
|
||||
def test_scalar_int_wrapped():
|
||||
"""Integer scalars are wrapped in torch.tensor."""
|
||||
result = _ensure_tensor_values({"sr": 24000})
|
||||
assert isinstance(result["sr"], torch.Tensor)
|
||||
assert result["sr"].item() == 24000
|
||||
|
||||
|
||||
def test_scalar_float_wrapped():
|
||||
"""Float scalars are wrapped in torch.tensor."""
|
||||
result = _ensure_tensor_values({"gain": 0.5})
|
||||
assert isinstance(result["gain"], torch.Tensor)
|
||||
assert result["gain"].item() == pytest.approx(0.5)
|
||||
|
||||
|
||||
def test_list_converted():
|
||||
"""Lists of numbers are converted to tensors."""
|
||||
result = _ensure_tensor_values({"ids": [1, 2, 3]})
|
||||
assert isinstance(result["ids"], torch.Tensor)
|
||||
assert list(result["ids"]) == [1, 2, 3]
|
||||
|
||||
|
||||
def test_non_tensorizable_dropped():
|
||||
"""Values that cannot be tensorized are dropped with warning."""
|
||||
result = _ensure_tensor_values({"ok": torch.randn(5), "bad": {"nested": "dict"}})
|
||||
assert "ok" in result
|
||||
assert "bad" not in result
|
||||
|
||||
|
||||
def test_mixed_payload():
|
||||
"""Mixed tensor + scalar + list payload is properly sanitized."""
|
||||
payload = {
|
||||
"audio": torch.randn(100),
|
||||
"sr": 24000,
|
||||
"codes": [1, 2, 3],
|
||||
"hidden": torch.randn(10, 64),
|
||||
}
|
||||
result = _ensure_tensor_values(payload)
|
||||
assert len(result) == 4
|
||||
for v in result.values():
|
||||
assert isinstance(v, torch.Tensor)
|
||||
@@ -161,3 +161,113 @@ async def test_diffusion_error_output_routed_as_finished(orchestrator_factory) -
|
||||
finally:
|
||||
orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
|
||||
orchestrator_fixture.thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_diffusion_client_error_output_propagates_status_code(orchestrator_factory) -> None:
|
||||
"""A client-error OmniRequestOutput (e.g. a guardrail 4xx) from a diffusion
|
||||
stage must surface as an ErrorMessage that carries status_code/error_type,
|
||||
so the frontend can map it to the correct HTTP response instead of a 500.
|
||||
"""
|
||||
stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
|
||||
orchestrator_fixture = orchestrator_factory([stage0])
|
||||
params = OmniDiffusionSamplingParams()
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-blocked",
|
||||
prompt={"prompt": "blocked"},
|
||||
original_prompt={"prompt": "blocked"},
|
||||
sampling_params_list=[params],
|
||||
final_stage_id=0,
|
||||
)
|
||||
|
||||
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
|
||||
|
||||
# Push a client-error output from the diffusion stage.
|
||||
stage0.push_diffusion_output(
|
||||
OmniRequestOutput.from_error(
|
||||
"req-blocked",
|
||||
"Input was blocked by Cosmos3 guardrails.",
|
||||
status_code=400,
|
||||
error_type="BadRequestError",
|
||||
)
|
||||
)
|
||||
|
||||
msg = await _get_any_output_message(orchestrator_fixture)
|
||||
|
||||
assert isinstance(msg, ErrorMessage)
|
||||
assert msg.request_id == "req-blocked"
|
||||
assert msg.stage_id == 0
|
||||
assert msg.fatal is False
|
||||
assert msg.status_code == 400
|
||||
assert msg.error_type == "BadRequestError"
|
||||
assert msg.error == "Input was blocked by Cosmos3 guardrails."
|
||||
|
||||
# Request state should be cleaned up.
|
||||
await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states)
|
||||
finally:
|
||||
orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
|
||||
orchestrator_fixture.thread.join(timeout=5)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "error_type"),
|
||||
[
|
||||
pytest.param(429, "RateLimitError", id="429-too-many-requests"),
|
||||
pytest.param(413, "PayloadTooLargeError", id="413-payload-too-large"),
|
||||
pytest.param(422, "UnprocessableEntityError", id="422-unprocessable-entity"),
|
||||
pytest.param(403, "PermissionDeniedError", id="403-forbidden"),
|
||||
],
|
||||
)
|
||||
async def test_diffusion_client_error_output_propagates_non_400_status(
|
||||
orchestrator_factory,
|
||||
status_code: int,
|
||||
error_type: str,
|
||||
) -> None:
|
||||
"""A diffusion-stage client error with a non-400 4xx status must be serialized
|
||||
into an ErrorMessage that preserves that exact status_code/error_type.
|
||||
"""
|
||||
stage0 = FakeStageClient(stage_type="diffusion", final_output=True, final_output_type="image")
|
||||
orchestrator_fixture = orchestrator_factory([stage0])
|
||||
params = OmniDiffusionSamplingParams()
|
||||
|
||||
try:
|
||||
await _enqueue_add_request(
|
||||
orchestrator_fixture,
|
||||
request_id="req-blocked",
|
||||
prompt={"prompt": "blocked"},
|
||||
original_prompt={"prompt": "blocked"},
|
||||
sampling_params_list=[params],
|
||||
final_stage_id=0,
|
||||
)
|
||||
|
||||
await _wait_for(lambda: len(stage0.add_request_calls) == 1)
|
||||
|
||||
# Push a non-400 client-error output from the diffusion stage.
|
||||
stage0.push_diffusion_output(
|
||||
OmniRequestOutput.from_error(
|
||||
"req-blocked",
|
||||
"client error from diffusion stage",
|
||||
status_code=status_code,
|
||||
error_type=error_type,
|
||||
)
|
||||
)
|
||||
|
||||
msg = await _get_any_output_message(orchestrator_fixture)
|
||||
|
||||
assert isinstance(msg, ErrorMessage)
|
||||
assert msg.request_id == "req-blocked"
|
||||
assert msg.stage_id == 0
|
||||
assert msg.fatal is False
|
||||
assert msg.status_code == status_code
|
||||
assert msg.error_type == error_type
|
||||
assert msg.error == "client error from diffusion stage"
|
||||
|
||||
# Request state should be cleaned up.
|
||||
await _wait_for(lambda: "req-blocked" not in orchestrator_fixture.orchestrator.request_states)
|
||||
finally:
|
||||
orchestrator_fixture.request_sync_q.put_nowait(ShutdownRequestMessage())
|
||||
orchestrator_fixture.thread.join(timeout=5)
|
||||
@@ -81,7 +81,7 @@ def test_multimodal_payload_and_completion_output():
|
||||
assert p is not None
|
||||
assert "waveform" in p.tensors and torch.equal(p.primary_tensor, data["waveform"])
|
||||
assert p.metadata["sample_rate"] == 16000
|
||||
assert not p.is_empty and len(p) == 1
|
||||
assert not p.is_empty and len(p) == 2 # 1 tensor + 1 metadata
|
||||
|
||||
# None/empty returns None
|
||||
assert MultimodalPayload.from_dict(None) is None
|
||||
@@ -125,7 +125,7 @@ def test_output_modality_printed_examples(capsys):
|
||||
print(f" tensors keys : {list(p.tensors.keys())}")
|
||||
print(f" primary_tensor: shape={p.primary_tensor.shape}, dtype={p.primary_tensor.dtype}")
|
||||
print(f" metadata : {p.metadata}")
|
||||
print(f" is_empty={p.is_empty}, len={len(p)}")
|
||||
print(f" is_empty={p.is_empty}, len={len(p)}") # len = tensors + metadata
|
||||
print(f" from_dict(None) -> {MultimodalPayload.from_dict(None)}")
|
||||
print(f" from_dict({{}}) -> {MultimodalPayload.from_dict({})}")
|
||||
|
||||
|
||||
@@ -209,3 +209,64 @@ def test_cumulative_token_ids_is_a_copy():
|
||||
s = _make_state(RequestOutputKind.DELTA)
|
||||
out = s._new_completion_output([42], None, None)
|
||||
assert out.cumulative_token_ids is not _DETOK.output_token_ids
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# No-detokenizer regression tests (Concern 2)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_NO_DETOK_STATE_KWARGS = {
|
||||
**_DEFAULT_STATE_KWARGS,
|
||||
"detokenizer": None,
|
||||
"logprobs_processor": None,
|
||||
}
|
||||
|
||||
|
||||
def _make_no_detok_state(output_kind: RequestOutputKind):
|
||||
return OmniRequestState(**_NO_DETOK_STATE_KWARGS, output_kind=output_kind)
|
||||
|
||||
|
||||
def test_no_detokenizer_completion_output():
|
||||
"""_new_completion_output works when detokenizer is None (generation stages)."""
|
||||
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
|
||||
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
||||
out = s._new_completion_output([1, 2], None, None)
|
||||
# Should still produce output with empty text and the token_ids as-is
|
||||
assert out.text == ""
|
||||
assert list(out.token_ids) == [1, 2]
|
||||
assert hasattr(out, "cumulative_token_ids")
|
||||
assert list(out.cumulative_token_ids) == [1, 2]
|
||||
# Multimodal data should be attached
|
||||
assert AUDIO in out.multimodal_output
|
||||
|
||||
|
||||
def test_no_detokenizer_make_request_output():
|
||||
"""make_request_output works without detokenizer when multimodal data is present."""
|
||||
s = _make_no_detok_state(RequestOutputKind.CUMULATIVE)
|
||||
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
||||
result = s.make_request_output([], None, FinishReason.STOP, None)
|
||||
assert result is not None
|
||||
assert not isinstance(result, PoolingRequestOutput)
|
||||
assert AUDIO in result.outputs[0].multimodal_output
|
||||
|
||||
|
||||
def test_no_detokenizer_stream_interval_skipped():
|
||||
"""stream_interval > 1 does not crash when detokenizer is None."""
|
||||
kwargs = {**_NO_DETOK_STATE_KWARGS, "stream_interval": 5}
|
||||
s = OmniRequestState(**kwargs, output_kind=RequestOutputKind.DELTA)
|
||||
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
||||
# Should not assert; stream_interval logic is simply skipped
|
||||
result = s.make_request_output([], None, FinishReason.STOP, None)
|
||||
assert result is not None
|
||||
|
||||
|
||||
def test_no_detokenizer_final_only():
|
||||
"""FINAL_ONLY with no detokenizer returns None until finish."""
|
||||
s = _make_no_detok_state(RequestOutputKind.FINAL_ONLY)
|
||||
s.add_multimodal_tensor(torch.randn(10), mm_type=AUDIO)
|
||||
# Non-finish step should return None
|
||||
assert s.make_request_output([], None, None, None) is None
|
||||
# Finish step should return output
|
||||
result = s.make_request_output([], None, FinishReason.STOP, None)
|
||||
assert result is not None
|
||||
assert AUDIO in result.outputs[0].multimodal_output
|
||||
@@ -0,0 +1,88 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Msgpack roundtrip tests for OmniEngineCoreOutputs with multimodal_output.
|
||||
|
||||
Validates that tensor-only payloads survive msgspec encode/decode and that
|
||||
non-tensor values fail decoding (enforcing the wire invariant).
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.v1.serial_utils import MsgpackDecoder, MsgpackEncoder
|
||||
|
||||
from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _roundtrip(outputs: OmniEngineCoreOutputs) -> OmniEngineCoreOutputs:
|
||||
encoder = MsgpackEncoder()
|
||||
encoded = encoder.encode(outputs)
|
||||
decoder = MsgpackDecoder(OmniEngineCoreOutputs)
|
||||
return decoder.decode(encoded)
|
||||
|
||||
|
||||
def test_tensor_only_roundtrip():
|
||||
"""Tensor-only multimodal_output survives msgpack roundtrip."""
|
||||
audio = torch.randn(16000)
|
||||
sr = torch.tensor(24000)
|
||||
eco = OmniEngineCoreOutput(
|
||||
request_id="req-1",
|
||||
new_token_ids=[1, 2, 3],
|
||||
finish_reason=None,
|
||||
multimodal_output={"audio": audio, "sr": sr},
|
||||
)
|
||||
decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
|
||||
assert len(decoded.outputs) == 1
|
||||
out = decoded.outputs[0]
|
||||
assert out.request_id == "req-1"
|
||||
assert isinstance(out.multimodal_output, dict)
|
||||
assert torch.allclose(out.multimodal_output["audio"], audio)
|
||||
assert out.multimodal_output["sr"].item() == 24000
|
||||
|
||||
|
||||
def test_empty_multimodal_roundtrip():
|
||||
"""None multimodal_output roundtrips correctly."""
|
||||
eco = OmniEngineCoreOutput(
|
||||
request_id="req-2",
|
||||
new_token_ids=[4],
|
||||
finish_reason=None,
|
||||
multimodal_output=None,
|
||||
)
|
||||
decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
|
||||
assert decoded.outputs[0].multimodal_output is None
|
||||
|
||||
|
||||
def test_multiple_tensor_keys_roundtrip():
|
||||
"""Multiple tensor keys survive roundtrip."""
|
||||
hidden = torch.randn(10, 4096)
|
||||
codes = torch.randint(0, 1024, (1, 100))
|
||||
eco = OmniEngineCoreOutput(
|
||||
request_id="req-3",
|
||||
new_token_ids=[],
|
||||
finish_reason=None,
|
||||
multimodal_output={"hidden": hidden, "codes": codes},
|
||||
)
|
||||
decoded = _roundtrip(OmniEngineCoreOutputs(outputs=[eco]))
|
||||
out = decoded.outputs[0]
|
||||
assert torch.allclose(out.multimodal_output["hidden"], hidden)
|
||||
assert torch.equal(out.multimodal_output["codes"], codes)
|
||||
|
||||
|
||||
def test_non_tensor_value_fails_decode():
|
||||
"""Non-tensor values in multimodal_output cause decode failure.
|
||||
|
||||
This validates that the wire field must be tensor-only and that
|
||||
the _ensure_tensor_values boundary in model runners is necessary.
|
||||
"""
|
||||
eco = OmniEngineCoreOutput(
|
||||
request_id="req-bad",
|
||||
new_token_ids=[],
|
||||
finish_reason=None,
|
||||
multimodal_output={"audio": torch.randn(10), "sr": 24000}, # scalar!
|
||||
)
|
||||
encoder = MsgpackEncoder()
|
||||
encoded = encoder.encode(OmniEngineCoreOutputs(outputs=[eco]))
|
||||
decoder = MsgpackDecoder(OmniEngineCoreOutputs)
|
||||
with pytest.raises(Exception):
|
||||
decoder.decode(encoded)
|
||||
@@ -0,0 +1,187 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for decode_audio_url helper."""
|
||||
|
||||
import asyncio
|
||||
import base64
|
||||
import os
|
||||
|
||||
import httpx
|
||||
import pytest
|
||||
|
||||
from vllm_omni.entrypoints.openai.errors import InvalidInputReferenceError
|
||||
from vllm_omni.entrypoints.openai.video_api_utils import decode_audio_url
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def _run(coro):
|
||||
return asyncio.run(coro)
|
||||
|
||||
|
||||
def _cleanup(path):
|
||||
if path and os.path.exists(path):
|
||||
os.unlink(path)
|
||||
|
||||
|
||||
class TestDecodeAudioUrlDataUrl:
|
||||
"""Tests for base64 data URL decoding."""
|
||||
|
||||
def test_valid_mp3_data_url(self):
|
||||
audio_bytes = b"\xff\xfb\x90\x00" * 100
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/mp3;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert os.path.isfile(path)
|
||||
assert path.endswith(".mp3")
|
||||
with open(path, "rb") as f:
|
||||
assert f.read() == audio_bytes
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_valid_wav_data_url(self):
|
||||
audio_bytes = b"RIFF" + b"\x00" * 100
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/wav;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert path.endswith(".wav")
|
||||
with open(path, "rb") as f:
|
||||
assert f.read() == audio_bytes
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_valid_mpeg_data_url_uses_mp3_suffix(self):
|
||||
audio_bytes = b"\xff\xfb\x90\x00" * 50
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/mpeg;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert path.endswith(".mp3")
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_valid_ogg_data_url(self):
|
||||
audio_bytes = b"OggS" + b"\x00" * 100
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/ogg;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert path.endswith(".ogg")
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_invalid_base64_raises(self):
|
||||
url = "data:audio/mp3;base64,!!!not-valid-base64!!!"
|
||||
|
||||
with pytest.raises(InvalidInputReferenceError, match="not valid base64"):
|
||||
_run(decode_audio_url(url))
|
||||
|
||||
def test_empty_audio_data_raises(self):
|
||||
url = "data:audio/mp3;base64,"
|
||||
|
||||
with pytest.raises(InvalidInputReferenceError, match="audio data is empty"):
|
||||
_run(decode_audio_url(url))
|
||||
|
||||
|
||||
class TestDecodeAudioUrlSuffixSanitization:
|
||||
"""Tests for MIME extension sanitization (path traversal prevention)."""
|
||||
|
||||
def test_path_traversal_in_mime_is_neutralized(self):
|
||||
audio_bytes = b"\x00" * 100
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/../..;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
# ".." contains non-alnum dots, so falls back to .wav
|
||||
assert path.endswith(".wav")
|
||||
assert os.path.dirname(path) == "/tmp"
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_non_alnum_extension_falls_back_to_wav(self):
|
||||
audio_bytes = b"\x00" * 100
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/foo-bar;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert path.endswith(".wav")
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_too_long_extension_falls_back_to_wav(self):
|
||||
audio_bytes = b"\x00" * 100
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/verylongextension;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert path.endswith(".wav")
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_valid_short_alnum_extension_is_used(self):
|
||||
audio_bytes = b"\x00" * 100
|
||||
b64 = base64.b64encode(audio_bytes).decode()
|
||||
url = f"data:audio/flac;base64,{b64}"
|
||||
|
||||
path = _run(decode_audio_url(url))
|
||||
try:
|
||||
assert path.endswith(".flac")
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
|
||||
class TestDecodeAudioUrlHttp:
|
||||
"""Tests for HTTP URL decoding."""
|
||||
|
||||
def test_valid_http_url(self, monkeypatch):
|
||||
audio_bytes = b"\xff\xfb\x90\x00" * 50
|
||||
fake_resp = httpx.Response(200, content=audio_bytes)
|
||||
|
||||
async def _mock_get(self, url):
|
||||
return fake_resp
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
|
||||
monkeypatch.setattr(httpx.Response, "raise_for_status", lambda self: None)
|
||||
|
||||
path = _run(decode_audio_url("https://example.com/audio.mp3"))
|
||||
try:
|
||||
assert os.path.isfile(path)
|
||||
assert path.endswith(".wav")
|
||||
with open(path, "rb") as f:
|
||||
assert f.read() == audio_bytes
|
||||
finally:
|
||||
_cleanup(path)
|
||||
|
||||
def test_http_error_raises(self, monkeypatch):
|
||||
async def _mock_get(self, url):
|
||||
raise httpx.HTTPStatusError("Not Found", request=httpx.Request("GET", url), response=httpx.Response(404))
|
||||
|
||||
monkeypatch.setattr(httpx.AsyncClient, "get", _mock_get)
|
||||
|
||||
with pytest.raises(InvalidInputReferenceError, match="failed to download audio"):
|
||||
_run(decode_audio_url("https://example.com/missing.mp3"))
|
||||
|
||||
|
||||
class TestDecodeAudioUrlInvalid:
|
||||
"""Tests for invalid URL schemes."""
|
||||
|
||||
def test_local_file_path_raises(self):
|
||||
with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
|
||||
_run(decode_audio_url("/path/to/local/file.mp3"))
|
||||
|
||||
def test_ftp_url_raises(self):
|
||||
with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
|
||||
_run(decode_audio_url("ftp://example.com/file.mp3"))
|
||||
|
||||
def test_empty_string_raises(self):
|
||||
with pytest.raises(InvalidInputReferenceError, match="must be an http.*URL or data URL"):
|
||||
_run(decode_audio_url(""))
|
||||
Loaded 100 of 303 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user