Compare commits
1
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
9fb3ee21a3 |
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.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# exception in upload_pipeline_with_skip_ci.sh (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.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.
|
||||
# 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.
|
||||
steps:
|
||||
- label: ":github: Resolve skip-ci & upload pipeline"
|
||||
key: upload-ci-pipeline
|
||||
commands:
|
||||
- "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
|
||||
- "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
key: upload-ready-pipeline
|
||||
if: __UPLOAD_READY_IF__
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -41,7 +41,7 @@ steps:
|
||||
key: upload-merge-pipeline
|
||||
if: __UPLOAD_MERGE_IF__
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -23,7 +23,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -1,308 +0,0 @@
|
||||
#!/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())
|
||||
@@ -0,0 +1,146 @@
|
||||
#!/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
|
||||
+148
-159
@@ -1,111 +1,54 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
|
||||
- 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 · 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"
|
||||
|
||||
- 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 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
|
||||
|
||||
- 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'
|
||||
"
|
||||
- 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))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
@@ -119,6 +62,42 @@ 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
|
||||
@@ -132,65 +111,50 @@ 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.
|
||||
- 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 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: "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"
|
||||
'
|
||||
- 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"
|
||||
'
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
@@ -231,3 +195,28 @@ 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"
|
||||
'
|
||||
+142
-154
@@ -1,42 +1,26 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
|
||||
- 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: "Custom Pipeline Test"
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
|
||||
- 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
|
||||
@@ -46,83 +30,39 @@ steps:
|
||||
commands:
|
||||
- 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 · 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: "Custom Pipeline Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
|
||||
|
||||
- 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'
|
||||
"
|
||||
- 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))
|
||||
'
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
@@ -136,6 +76,56 @@ 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
|
||||
@@ -159,56 +149,43 @@ 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"
|
||||
|
||||
- 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: "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: "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"
|
||||
'
|
||||
- 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"
|
||||
'
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
@@ -248,3 +225,14 @@ 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"
|
||||
'
|
||||
+25
-78
@@ -206,34 +206,7 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Ming-Omni-TTS Test"
|
||||
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
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/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:
|
||||
- |
|
||||
@@ -276,12 +249,6 @@ 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:
|
||||
- |
|
||||
@@ -324,12 +291,6 @@ 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'
|
||||
@@ -368,11 +329,6 @@ 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'
|
||||
@@ -410,14 +366,37 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
- 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"
|
||||
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/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
@@ -459,9 +438,6 @@ 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'
|
||||
@@ -500,9 +476,6 @@ 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'
|
||||
@@ -541,9 +514,6 @@ 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'
|
||||
@@ -582,13 +552,6 @@ 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'
|
||||
@@ -628,9 +591,6 @@ 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 '
|
||||
@@ -672,9 +632,6 @@ 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 "
|
||||
@@ -716,10 +673,6 @@ 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 "
|
||||
@@ -761,9 +714,6 @@ 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:
|
||||
- |
|
||||
@@ -839,9 +789,6 @@ 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,55 +568,6 @@ 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
|
||||
@@ -640,7 +591,6 @@ 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:
|
||||
@@ -1512,7 +1462,6 @@ 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
|
||||
@@ -1526,7 +1475,6 @@ 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
|
||||
|
||||
+61
-62
@@ -215,43 +215,6 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Entrypoints Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -sv tests/entrypoints/ -m 'core_model and cuda' --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
|
||||
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
depends_on: upload-ready-pipeline
|
||||
steps:
|
||||
@@ -276,11 +239,6 @@ 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 "
|
||||
@@ -315,14 +273,73 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
- 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"
|
||||
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/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
@@ -406,9 +423,6 @@ 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 "
|
||||
@@ -444,12 +458,6 @@ 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 "
|
||||
@@ -491,9 +499,6 @@ 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 "
|
||||
@@ -536,9 +541,6 @@ 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 "
|
||||
@@ -580,9 +582,6 @@ 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,13 +6,6 @@
|
||||
{% 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:
|
||||
@@ -38,15 +31,8 @@ steps:
|
||||
agents:
|
||||
queue: amd-cpu
|
||||
|
||||
{% 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 %}
|
||||
{% for step in steps %}
|
||||
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
|
||||
- label: "{{ step.agent_pool }}: {{ step.label }}"
|
||||
depends_on: amd-build
|
||||
agents:
|
||||
@@ -60,23 +46,8 @@ steps:
|
||||
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 %}
|
||||
@@ -90,52 +61,5 @@ steps:
|
||||
{% 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 %}
|
||||
{% if step.grade and step.grade == "Blocking" %}
|
||||
soft_fail: false
|
||||
{% else %}
|
||||
soft_fail: true
|
||||
{% endif%}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
{% endfor %}
|
||||
@@ -1,103 +0,0 @@
|
||||
---
|
||||
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)
|
||||
@@ -1,70 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,113 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,82 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,47 +0,0 @@
|
||||
# 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
|
||||
@@ -1,122 +0,0 @@
|
||||
# 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,8 +21,6 @@ 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,15 +1,24 @@
|
||||
<!-- markdownlint-disable -->
|
||||
PLEASE FILL IN THE PR DESCRIPTION HERE.
|
||||
**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.
|
||||
|
||||
## Purpose
|
||||
|
||||
## Test Plan
|
||||
|
||||
**vLLM Version:**
|
||||
|
||||
**vLLM-Omni Commit:**
|
||||
|
||||
## Test Result
|
||||
|
||||
**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.
|
||||
---
|
||||
<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>
|
||||
(anything written below this line will be removed by GitHub Actions)
|
||||
@@ -1,5 +1,3 @@
|
||||
# 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,9 +57,8 @@ 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. 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)
|
||||
- Omni-modality models (e.g. Qwen-Omni)
|
||||
- Multi-modality generation models (e.g. Qwen-Image)
|
||||
|
||||
## Getting Started
|
||||
|
||||
|
||||
@@ -33,15 +33,6 @@ 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
|
||||
|
||||
@@ -87,8 +87,6 @@ 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,13 +118,17 @@ 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](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:
|
||||
Before submitting a PR, run the precheck-pr skill with code agent for a self-review against project conventions:
|
||||
|
||||
```
|
||||
/precheck-pr
|
||||
```
|
||||
|
||||
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 [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) 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 skill in `.claude/skills/precheck-pr/` 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,101 +268,6 @@ 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.
|
||||
|
||||
<<<<<<< ready
|
||||
## 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 such as Simple Test, Diffusion unit tests, Engine/Model Executor, Distributed, Custom Pipeline, Entrypoints (L2), LoRA / Entrypoints (L3), and **all E2E jobs on L4 docker queues** (`gpu_1_queue`, `gpu_4_queue`).
|
||||
- **Diff-gated** (`source_file_dependencies` set): **H100 E2E jobs only** (`queue: mithril-h100-pool` under the E2E Test group)—typically the model’s pytest file(s) plus related `vllm_omni/` model, deploy, and stage paths.
|
||||
|
||||
Adding a new expensive H100 E2E step: list the test file(s) and the smallest set of source prefixes that should trigger it. Prefer **per-step** deps on leaf jobs rather than a broad group-level list unless every child shares the same paths.
|
||||
|
||||
#### YAML example
|
||||
|
||||
```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"
|
||||
```
|
||||
|
||||
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
|
||||
```
|
||||
>>>>>>> main
|
||||
|
||||
## Chapter 1: L1 & L2 Level Testing - Unit Testing and Basic End-to-End Verification
|
||||
|
||||
### 1.1 Testing Purpose
|
||||
@@ -397,7 +302,6 @@ 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, expensive **H100 E2E** 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) paths; L4 docker jobs and non-E2E steps still run. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
|
||||
- ***Execution Environment***: L1 uses ***CPU*** environment; L2 requires ***GPU*** environment.
|
||||
- ***Script Example***:
|
||||
|
||||
@@ -500,8 +404,7 @@ 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 (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—only **H100 E2E** steps are diff-gated; L4 docker E2E and other L3 groups always upload. See [Diff-aware Buildkite uploads](#diff-aware-buildkite-uploads-source_file_dependencies).
|
||||
- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch.
|
||||
- ***Execution Environment***: ***GPU*** servers.
|
||||
- ***Script Example***:
|
||||
|
||||
@@ -617,8 +520,6 @@ 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,8 +69,6 @@ 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,7 +8,6 @@ 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)
|
||||
|
||||
@@ -24,7 +23,6 @@ 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"`. |
|
||||
|
||||
|
||||
@@ -1,76 +0,0 @@
|
||||
# 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,7 +63,6 @@ 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,7 +114,6 @@ Content-Type: application/json
|
||||
| `instructions` | string | "" | Voice style/emotion instructions |
|
||||
| `max_new_tokens` | integer | 2048 | Maximum tokens to generate |
|
||||
| `initial_codec_chunk_frames` | integer | null | Per-request initial chunk size override for TTFA tuning. When null, IC is computed dynamically based on server load. |
|
||||
| `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.
|
||||
@@ -333,8 +332,7 @@ 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",
|
||||
"non_streaming_mode": true
|
||||
"ref_text": "Original transcript of the reference audio"
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
@@ -456,7 +454,6 @@ 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:
|
||||
|
||||
@@ -472,7 +469,6 @@ 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,7 +70,6 @@ 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 |
|
||||
@@ -162,25 +161,6 @@ 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
|
||||
@@ -209,7 +189,5 @@ 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** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
@@ -1,128 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,117 +0,0 @@
|
||||
# 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,20 +143,6 @@ 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,
|
||||
@@ -250,8 +236,8 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--resolution",
|
||||
type=int,
|
||||
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.",
|
||||
default=640,
|
||||
help="Bucket in (640, 1024) to determine the condition and output resolution",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
@@ -398,15 +384,6 @@ 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:
|
||||
@@ -487,7 +464,6 @@ 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()
|
||||
@@ -497,26 +473,22 @@ 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},
|
||||
},
|
||||
sampling_params,
|
||||
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,
|
||||
),
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -17,7 +17,6 @@ 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 |
|
||||
@@ -160,46 +159,6 @@ 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.
|
||||
|
||||
@@ -1,248 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,110 +0,0 @@
|
||||
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
|
||||
@@ -1,200 +0,0 @@
|
||||
# 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()
|
||||
@@ -1,277 +0,0 @@
|
||||
# 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:
|
||||
|
||||
@@ -1,115 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,44 +0,0 @@
|
||||
#!/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
|
||||
@@ -1,29 +0,0 @@
|
||||
#!/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
|
||||
@@ -16,7 +16,6 @@ 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` | ✓ | — | — | — |
|
||||
@@ -203,51 +202,6 @@ 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`.
|
||||
|
||||
@@ -1,113 +0,0 @@
|
||||
# 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.
|
||||
@@ -1,225 +0,0 @@
|
||||
# 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()
|
||||
@@ -1,76 +0,0 @@
|
||||
#!/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
|
||||
@@ -1,25 +0,0 @@
|
||||
#!/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,10 +408,7 @@ def create_app(api_base: str):
|
||||
await resp.aclose()
|
||||
await client.aclose()
|
||||
|
||||
# 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")
|
||||
return StreamingResponse(relay(), media_type="application/octet-stream")
|
||||
|
||||
# ── Gradio UI ────────────────────────────────────────────────
|
||||
voices = fetch_voices(api_base)
|
||||
|
||||
+1
-6
@@ -66,8 +66,7 @@ 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)
|
||||
"s3tokenizer>=0.3.0" # for tts test
|
||||
"kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
||||
]
|
||||
|
||||
demo = [
|
||||
@@ -293,9 +292,5 @@ 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,7 +52,6 @@ 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:**
|
||||
|
||||
@@ -104,28 +103,8 @@ 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.
|
||||
- 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.
|
||||
- `enforce_eager=true` for both stages (CUDA graph available for stage 0 but no throughput gain at batch>1).
|
||||
- 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).
|
||||
- Stage 0 full-decode CUDA graph is experimental; sampler, delay-state updates, staging, and request postprocess remain outside the graph.
|
||||
- Async streaming (chunk-based) not yet implemented; sync-only pipeline.
|
||||
@@ -33,7 +33,6 @@ 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,34 +171,6 @@ 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
|
||||
|
||||
@@ -1,82 +0,0 @@
|
||||
# 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,3 +22,4 @@ prettytable>=3.8.0
|
||||
aenum==3.1.16
|
||||
pyzmq>=25.0.0
|
||||
janus>=1.0.0
|
||||
pydub
|
||||
@@ -1,3 +1,2 @@
|
||||
-r common.txt
|
||||
onnxruntime>=1.23.2
|
||||
auto-round-lib
|
||||
@@ -20,7 +20,6 @@ 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
|
||||
|
||||
@@ -143,300 +142,3 @@ 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]
|
||||
@@ -1,158 +0,0 @@
|
||||
"""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,9 +46,6 @@ 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"]
|
||||
|
||||
@@ -1,74 +0,0 @@
|
||||
[
|
||||
{
|
||||
"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,66 +339,5 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"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,78 +198,19 @@ 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_invalid_requests(
|
||||
def test_realtime_rejected_when_async_chunk_enabled(
|
||||
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": ws_error_code,
|
||||
"err_message": err_message,
|
||||
"ws_error_code": "unsupported",
|
||||
"err_message": "async_chunk",
|
||||
}
|
||||
)
|
||||
@@ -319,7 +319,6 @@ 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(
|
||||
@@ -388,7 +387,6 @@ 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:
|
||||
@@ -431,7 +429,6 @@ 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:
|
||||
@@ -774,7 +771,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@pytest.mark.skip(reason="OOM injection causes the serving instance to exit")
|
||||
@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,38 +131,6 @@ 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
|
||||
|
||||
|
||||
@@ -1,728 +0,0 @@
|
||||
# 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,260 +97,6 @@ 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."""
|
||||
|
||||
@@ -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, "inline_small_payloads": True}
|
||||
config = {"shm_threshold_bytes": 1024} # Large threshold to use inline
|
||||
connector = SharedMemoryConnector(config)
|
||||
connectors_map = {("0", "1"): connector}
|
||||
|
||||
|
||||
@@ -75,21 +75,18 @@ def test_create_unknown_connector():
|
||||
|
||||
@pytest.fixture
|
||||
def shm_connector():
|
||||
config = {"shm_threshold_bytes": 100, "inline_small_payloads": True}
|
||||
config = {"shm_threshold_bytes": 100} # Small threshold for testing
|
||||
return SharedMemoryConnector(config)
|
||||
|
||||
|
||||
def test_put_get_inline(shm_connector):
|
||||
"""Test inline transfer for small data."""
|
||||
"""Test transfer for small data (currently always uses SHM)."""
|
||||
data = {"small": "data"}
|
||||
|
||||
success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
|
||||
assert success is True
|
||||
assert "inline_bytes" in metadata
|
||||
assert "shm" not in metadata
|
||||
assert "shm" 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)
|
||||
|
||||
@@ -387,15 +387,14 @@ 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, scheduler_requests=scheduler_requests)
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
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, scheduler_requests=scheduler_requests)
|
||||
adapter.restore_queues(waiting_queue, running_queue)
|
||||
assert waiting_queue == [waiting_req]
|
||||
assert running_queue == [running_req]
|
||||
assert adapter.waiting_for_chunk_waiting_requests == deque()
|
||||
@@ -898,8 +897,6 @@ 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)
|
||||
@@ -1173,164 +1170,3 @@ 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": 0})
|
||||
c = SharedMemoryConnector({"shm_threshold_bytes": 64})
|
||||
yield c
|
||||
c.close()
|
||||
|
||||
|
||||
@@ -29,46 +29,7 @@ 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,
|
||||
)
|
||||
|
||||
# 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])}"
|
||||
)
|
||||
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.prompt_utils import build_prompt_tokens, resolve_stop_token_ids
|
||||
|
||||
os.environ["DIFFUSION_ATTENTION_BACKEND"] = "TORCH_SDPA"
|
||||
|
||||
@@ -168,7 +129,6 @@ _DEPLOY_CONFIG = {
|
||||
"stop_token_ids": [128025],
|
||||
"detokenize": True,
|
||||
"skip_special_tokens": False,
|
||||
"include_stop_str_in_output": True,
|
||||
},
|
||||
},
|
||||
{
|
||||
@@ -440,8 +400,6 @@ 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']}"
|
||||
)
|
||||
@@ -550,8 +508,6 @@ 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']}"
|
||||
)
|
||||
|
||||
@@ -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", "rocm": "MI325"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100"}, 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", "rocm": "MI325"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100"}, 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", "rocm": "MI325"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100"}, 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", "rocm": "MI325"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100"}, 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", "rocm": "MI325"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100"}, 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", "rocm": "MI325"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100"}, 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", "rocm": "MI325"}, num_cards=4)
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=4)
|
||||
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
|
||||
def test_image_to_audio(omni_runner, omni_runner_handler) -> None:
|
||||
"""
|
||||
|
||||
@@ -1,295 +0,0 @@
|
||||
# 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
|
||||
|
||||
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: dict) -> torch.Tensor:
|
||||
assert isinstance(multimodal_output, dict), f"Expected dict, 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: dict) -> 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, dict):
|
||||
final_outputs.append(item)
|
||||
continue
|
||||
completions = getattr(request_output, "outputs", None) or []
|
||||
if any(isinstance(getattr(completion, "multimodal_output", None), dict) for completion in completions):
|
||||
final_outputs.append(item)
|
||||
return final_outputs
|
||||
|
||||
|
||||
def _extract_multimodal_output(output) -> dict:
|
||||
multimodal_output = getattr(output, "multimodal_output", None)
|
||||
if isinstance(multimodal_output, dict):
|
||||
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, dict):
|
||||
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, dict):
|
||||
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())
|
||||
@@ -1,160 +0,0 @@
|
||||
# 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}"
|
||||
File renamed without changes.
@@ -1,92 +0,0 @@
|
||||
# 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)
|
||||
@@ -91,90 +91,3 @@ 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,
|
||||
}
|
||||
)
|
||||
@@ -161,113 +161,3 @@ 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)
|
||||
@@ -1,187 +0,0 @@
|
||||
# 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(""))
|
||||
@@ -29,7 +29,6 @@ from vllm_omni.entrypoints.openai.image_api_utils import (
|
||||
parse_size,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
|
||||
from vllm_omni.errors import GuardrailViolationError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -588,24 +587,6 @@ def test_generate_single_image(test_client):
|
||||
assert test_client.app.state.engine_client.captured_prompt["modalities"] == ["image"]
|
||||
|
||||
|
||||
def test_generate_images_guardrail_error_returns_400(test_client, mock_async_diffusion):
|
||||
async def blocked_generate(**kwargs):
|
||||
raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.")
|
||||
yield MockGenerationResult([]) # pragma: no cover
|
||||
|
||||
mock_async_diffusion.generate = blocked_generate
|
||||
response = test_client.post(
|
||||
"/v1/images/generations",
|
||||
json={
|
||||
"prompt": "blocked prompt",
|
||||
"n": 1,
|
||||
"size": "1024x1024",
|
||||
},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails."
|
||||
|
||||
|
||||
def test_generate_images_async_omni_sampling_params(async_omni_test_client):
|
||||
"""Test AsyncOmni path uses per-stage sampling params."""
|
||||
response = async_omni_test_client.post(
|
||||
@@ -910,31 +891,6 @@ def test_image_edits_streaming_errors_on_empty_final_image(streaming_image_edit_
|
||||
assert payloads[2] == "[DONE]"
|
||||
|
||||
|
||||
def test_image_edits_streaming_guardrail_error_uses_400(streaming_image_edit_client):
|
||||
async def generate_guardrail_error(prompt, request_id, sampling_params=None, sampling_params_list=None, **kwargs):
|
||||
raise GuardrailViolationError("Input was blocked by Cosmos3 guardrails.")
|
||||
yield MockStageResult(stage_id=1, final_output_type="image", images=[]) # pragma: no cover
|
||||
|
||||
streaming_image_edit_client.app.state.engine_client.generate = generate_guardrail_error
|
||||
img_bytes = make_test_image_bytes((16, 16))
|
||||
response = streaming_image_edit_client.post(
|
||||
"/v1/images/edits",
|
||||
files=[("image", img_bytes)],
|
||||
data={
|
||||
"prompt": "blocked prompt",
|
||||
"stream": "true",
|
||||
},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
payloads = _parse_sse_payloads(response.text)
|
||||
assert payloads[0]["object"] == "error"
|
||||
assert payloads[0]["error"]["message"] == "Input was blocked by Cosmos3 guardrails."
|
||||
assert payloads[0]["error"]["type"] == "BadRequestError"
|
||||
assert payloads[0]["error"]["code"] == 400
|
||||
assert payloads[1] == "[DONE]"
|
||||
|
||||
|
||||
def test_image_edits_streaming_rejects_single_stage(test_client):
|
||||
img_bytes = make_test_image_bytes((16, 16))
|
||||
response = test_client.post(
|
||||
|
||||
@@ -31,7 +31,6 @@ from vllm_omni.entrypoints.openai.protocol.audio import (
|
||||
CreateAudio,
|
||||
OpenAICreateSpeechRequest,
|
||||
SpeechBatchItem,
|
||||
StreamingSpeechSessionConfig,
|
||||
)
|
||||
from vllm_omni.entrypoints.openai.serving_speech import (
|
||||
_TTS_LANGUAGES,
|
||||
@@ -1236,49 +1235,6 @@ class TestTTSMethods:
|
||||
assert params["language"] == ["English"]
|
||||
assert params["task_type"] == ["CustomVoice"]
|
||||
|
||||
def test_build_tts_params_base_non_streaming_mode_true(self, speech_server):
|
||||
"""Base task should pass through an explicit non_streaming_mode override."""
|
||||
req = OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
task_type="Base",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="reference",
|
||||
non_streaming_mode=True,
|
||||
)
|
||||
|
||||
params = speech_server._build_tts_params(req)
|
||||
|
||||
assert params["task_type"] == ["Base"]
|
||||
assert params["non_streaming_mode"] == [True]
|
||||
|
||||
def test_build_tts_params_base_omits_non_streaming_mode_by_default(self, speech_server):
|
||||
"""Base task should keep using the model default when no override is sent."""
|
||||
req = OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
task_type="Base",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="reference",
|
||||
)
|
||||
|
||||
params = speech_server._build_tts_params(req)
|
||||
|
||||
assert params["task_type"] == ["Base"]
|
||||
assert "non_streaming_mode" not in params
|
||||
|
||||
def test_build_tts_params_explicit_non_streaming_mode_overrides_voicedesign_default(self, speech_server):
|
||||
"""Explicit false should not be replaced by the VoiceDesign fallback."""
|
||||
req = OpenAICreateSpeechRequest(
|
||||
input="Hello",
|
||||
task_type="VoiceDesign",
|
||||
instructions="warm and calm",
|
||||
non_streaming_mode=False,
|
||||
)
|
||||
|
||||
params = speech_server._build_tts_params(req)
|
||||
|
||||
assert params["task_type"] == ["VoiceDesign"]
|
||||
assert params["non_streaming_mode"] == [False]
|
||||
|
||||
def test_load_supported_speakers(self, mocker: MockerFixture):
|
||||
"""Test _load_supported_speakers."""
|
||||
mock_engine_client = mocker.MagicMock()
|
||||
@@ -2067,34 +2023,6 @@ class TestMergeBatchItem:
|
||||
assert merged.task_type == "CustomVoice"
|
||||
assert merged.max_new_tokens == 512
|
||||
|
||||
def test_non_streaming_mode_batch_default_used(self):
|
||||
"""Batch-level non_streaming_mode should be used when item doesn't specify one."""
|
||||
batch = BatchSpeechRequest(
|
||||
items=[SpeechBatchItem(input="hi")],
|
||||
non_streaming_mode=True,
|
||||
)
|
||||
|
||||
merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
|
||||
|
||||
assert merged.non_streaming_mode is True
|
||||
|
||||
def test_non_streaming_mode_item_override_wins(self):
|
||||
"""Per-item false should override a true batch-level default."""
|
||||
batch = BatchSpeechRequest(
|
||||
items=[SpeechBatchItem(input="hi", non_streaming_mode=False)],
|
||||
non_streaming_mode=True,
|
||||
)
|
||||
|
||||
merged = OmniOpenAIServingSpeech._merge_batch_item(batch, batch.items[0])
|
||||
|
||||
assert merged.non_streaming_mode is False
|
||||
|
||||
|
||||
def test_streaming_speech_session_config_accepts_non_streaming_mode():
|
||||
config = StreamingSpeechSessionConfig(non_streaming_mode=True)
|
||||
|
||||
assert config.non_streaming_mode is True
|
||||
|
||||
|
||||
class TestAsyncOmniSupportedTasks:
|
||||
"""Test that AsyncOmni reports correct supported tasks based on output modalities."""
|
||||
@@ -2883,52 +2811,6 @@ class TestTTSAsyncOffloading:
|
||||
qwen3_tts_server._build_tts_params.assert_called_once()
|
||||
qwen3_tts_server._estimate_prompt_len_async.assert_awaited_once()
|
||||
|
||||
def test_prepare_speech_generation_qwen3_voicedesign_non_streaming_mode_false(
|
||||
self, qwen3_tts_server, mocker: MockerFixture
|
||||
):
|
||||
"""VoiceDesign explicit false should reach the model prompt additional_information."""
|
||||
qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
|
||||
qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
|
||||
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input="hello",
|
||||
task_type="VoiceDesign",
|
||||
instructions="warm and calm",
|
||||
non_streaming_mode=False,
|
||||
)
|
||||
_request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
|
||||
|
||||
assert tts_params["task_type"] == ["VoiceDesign"]
|
||||
assert tts_params["non_streaming_mode"] == [False]
|
||||
prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
|
||||
assert prompt["additional_information"] is tts_params
|
||||
assert prompt["additional_information"]["non_streaming_mode"] == [False]
|
||||
|
||||
def test_prepare_speech_generation_qwen3_base_non_streaming_mode_true(
|
||||
self, qwen3_tts_server, mocker: MockerFixture
|
||||
):
|
||||
"""Base explicit true should reach the model prompt additional_information."""
|
||||
qwen3_tts_server._validate_tts_request = mocker.MagicMock(return_value=None)
|
||||
qwen3_tts_server._resolve_ref_audio = mocker.AsyncMock(return_value=([0.0] * 48000, 24000))
|
||||
qwen3_tts_server._get_resolved_ref_audio_artifact_key = mocker.MagicMock(return_value=None)
|
||||
qwen3_tts_server._estimate_prompt_len_async = mocker.AsyncMock(return_value=512)
|
||||
|
||||
request = OpenAICreateSpeechRequest(
|
||||
input="hello",
|
||||
task_type="Base",
|
||||
ref_audio="data:audio/wav;base64,abc",
|
||||
ref_text="reference transcript",
|
||||
non_streaming_mode=True,
|
||||
)
|
||||
_request_id, _generator, tts_params = asyncio.run(qwen3_tts_server._prepare_speech_generation(request))
|
||||
|
||||
assert tts_params["task_type"] == ["Base"]
|
||||
assert tts_params["ref_text"] == ["reference transcript"]
|
||||
assert tts_params["non_streaming_mode"] == [True]
|
||||
prompt = qwen3_tts_server.engine_client.generate.call_args.kwargs["prompt"]
|
||||
assert prompt["additional_information"] is tts_params
|
||||
assert prompt["additional_information"]["non_streaming_mode"] == [True]
|
||||
|
||||
def test_qwen3_repeated_ref_audio_hot_path_sends_cache_key_without_waveform(self, qwen3_tts_server):
|
||||
"""After a ref artifact is marked ready, repeated requests avoid ref_audio payload IPC."""
|
||||
wav_list = [0.0] * 48000
|
||||
|
||||
@@ -31,7 +31,6 @@ from vllm_omni.entrypoints.openai.protocol.videos import (
|
||||
from vllm_omni.entrypoints.openai.serving_video import OmniOpenAIServingVideo
|
||||
from vllm_omni.entrypoints.openai.storage import LocalStorageManager
|
||||
from vllm_omni.entrypoints.openai.stores import AsyncDictStore, TaskRegistry
|
||||
from vllm_omni.errors import GuardrailViolationError
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -87,10 +86,8 @@ class BlockingVideoHandler:
|
||||
if self.stage_configs is None:
|
||||
self.stage_configs = stage_configs
|
||||
|
||||
async def generate_video_bytes(
|
||||
self, request, reference_id, *, reference_image=None, reference_video=None, reference_audio=None
|
||||
):
|
||||
del request, reference_id, reference_image, reference_video, reference_audio
|
||||
async def generate_video_bytes(self, request, reference_id, *, reference_image=None, reference_video=None):
|
||||
del request, reference_id, reference_image, reference_video
|
||||
self.started.set()
|
||||
try:
|
||||
await asyncio.Future()
|
||||
@@ -460,41 +457,6 @@ def test_cosmos3_reference_video_limit_caps_condition_frames_to_output_frames():
|
||||
assert api_server._reference_video_decode_spec(request, _cosmos3_stage_configs()).max_frames == 5
|
||||
|
||||
|
||||
def test_s2v_video_generation_with_audio_reference_form(test_client, mocker: MockerFixture):
|
||||
"""Speech-to-video: image + audio_reference (base64 data URL) passes audio path to multi_modal_data."""
|
||||
audio_bytes = b"\xff\xfb\x90\x00" * 50
|
||||
audio_b64 = base64.b64encode(audio_bytes).decode()
|
||||
audio_ref = json.dumps({"audio_url": f"data:audio/mp3;base64,{audio_b64}"})
|
||||
|
||||
mocker.patch(
|
||||
"vllm_omni.entrypoints.openai.serving_video._encode_video_bytes",
|
||||
return_value=b"fake-video",
|
||||
)
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
data={
|
||||
"prompt": "A person singing",
|
||||
"audio_reference": audio_ref,
|
||||
"width": "832",
|
||||
"height": "480",
|
||||
},
|
||||
files={"input_reference": ("face.png", _make_test_image_bytes((64, 64)), "image/png")},
|
||||
)
|
||||
|
||||
assert response.status_code == 200
|
||||
video_id = response.json()["id"]
|
||||
_wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value)
|
||||
|
||||
engine = test_client.app.state.openai_serving_video._engine_client
|
||||
prompt = engine.captured_prompt
|
||||
assert "multi_modal_data" in prompt
|
||||
assert "image" in prompt["multi_modal_data"]
|
||||
assert "audio" in prompt["multi_modal_data"]
|
||||
audio_path = prompt["multi_modal_data"]["audio"]
|
||||
assert isinstance(audio_path, str)
|
||||
assert audio_path.endswith(".mp3")
|
||||
|
||||
|
||||
def test_seconds_defaults_fps_and_frames(test_client, mocker: MockerFixture):
|
||||
fps_values = []
|
||||
|
||||
@@ -1020,25 +982,6 @@ def test_invalid_lora_returns_400(test_client):
|
||||
assert "lora object" in failed["error"]["message"].lower()
|
||||
|
||||
|
||||
def test_async_guardrail_error_returns_400_on_retrieve(test_client, mocker: MockerFixture):
|
||||
mocker.patch.object(
|
||||
OmniOpenAIServingVideo,
|
||||
"generate_video_bytes",
|
||||
side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."),
|
||||
)
|
||||
response = test_client.post("/v1/videos", data={"prompt": "blocked prompt"})
|
||||
assert response.status_code == 200
|
||||
|
||||
video_id = response.json()["id"]
|
||||
failed = _wait_for_status(test_client, video_id, VideoGenerationStatus.FAILED.value)
|
||||
assert failed["error"]["code"] == 400
|
||||
assert failed["error"]["message"] == "Input was blocked by Cosmos3 guardrails."
|
||||
|
||||
retrieve = test_client.get(f"/v1/videos/{video_id}")
|
||||
assert retrieve.status_code == 400
|
||||
assert retrieve.json()["error"]["code"] == 400
|
||||
|
||||
|
||||
def test_unsupported_image_reference_file_id_returns_400(test_client):
|
||||
response = test_client.post(
|
||||
"/v1/videos",
|
||||
@@ -1541,20 +1484,6 @@ def test_sync_generation_error_returns_500(test_client, mocker: MockerFixture):
|
||||
assert "GPU exploded" in response.json()["detail"]
|
||||
|
||||
|
||||
def test_sync_guardrail_error_returns_400(test_client, mocker: MockerFixture):
|
||||
mocker.patch.object(
|
||||
OmniOpenAIServingVideo,
|
||||
"generate_video_bytes",
|
||||
side_effect=GuardrailViolationError("Input was blocked by Cosmos3 guardrails."),
|
||||
)
|
||||
response = test_client.post(
|
||||
"/v1/videos/sync",
|
||||
data={"prompt": "blocked prompt"},
|
||||
)
|
||||
assert response.status_code == 400
|
||||
assert response.json()["detail"] == "Input was blocked by Cosmos3 guardrails."
|
||||
|
||||
|
||||
def test_sync_does_not_create_store_entry(test_client, mocker: MockerFixture):
|
||||
"""The sync endpoint should NOT leave any record in VIDEO_STORE — it is
|
||||
stateless by design."""
|
||||
|
||||
@@ -18,12 +18,6 @@ from vllm_omni.entrypoints.async_omni import AsyncOmni
|
||||
from vllm_omni.entrypoints.client_request_state import ClientRequestState
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.entrypoints.omni_base import OmniEngineDeadError
|
||||
from vllm_omni.errors import (
|
||||
OmniClientError,
|
||||
client_error_from_metadata,
|
||||
client_error_metadata,
|
||||
is_client_error_status,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
@@ -610,71 +604,6 @@ async def test_async_omni_propagates_fatal_error_context(monkeypatch: pytest.Mon
|
||||
assert getattr(exc_info.value, "error_stage_id") == 2
|
||||
|
||||
|
||||
def _enqueue_client_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
||||
engine.output_q.put_nowait(
|
||||
ErrorMessage(
|
||||
request_id=msg["request_id"],
|
||||
stage_id=2,
|
||||
error="Input was blocked by Cosmos3 guardrails.",
|
||||
status_code=400,
|
||||
error_type="BadRequestError",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_propagates_client_error_status(monkeypatch: pytest.MonkeyPatch):
|
||||
"""A non-fatal client error from the orchestrator must be routed to the
|
||||
requesting generate() call (not raised in the shared dispatcher) and
|
||||
surface as an OmniClientError carrying status_code/error_type."""
|
||||
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_client_error_message)
|
||||
_patch_engine(monkeypatch, engine)
|
||||
|
||||
app = AsyncOmni("dummy-model")
|
||||
try:
|
||||
with pytest.raises(OmniClientError) as exc_info:
|
||||
async for _ in app.generate(prompt="blocked", request_id="req-1"):
|
||||
pass
|
||||
finally:
|
||||
app.shutdown()
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.error_type == "BadRequestError"
|
||||
assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails."
|
||||
|
||||
|
||||
def _enqueue_server_error_message(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
||||
engine.output_q.put_nowait(
|
||||
ErrorMessage(
|
||||
request_id=msg["request_id"],
|
||||
stage_id=2,
|
||||
error="GPU exploded",
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
async def test_async_omni_propagates_server_error_as_runtime(monkeypatch: pytest.MonkeyPatch):
|
||||
"""A non-fatal error WITHOUT a 4xx status_code is a server fault: it must
|
||||
surface as a RuntimeError (-> HTTP 500), not an OmniClientError (-> 400).
|
||||
Guards against the fix over-broadly mapping every error to a client error."""
|
||||
engine = FakeAsyncOmniEngine(stage_metadata=THREE_STAGE_META, on_add_request=_enqueue_server_error_message)
|
||||
_patch_engine(monkeypatch, engine)
|
||||
|
||||
app = AsyncOmni("dummy-model")
|
||||
try:
|
||||
# OmniClientError subclasses ValueError, not RuntimeError, so matching
|
||||
# RuntimeError here also proves it was NOT raised as a client error.
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
async for _ in app.generate(prompt="boom", request_id="req-1"):
|
||||
pass
|
||||
finally:
|
||||
app.shutdown()
|
||||
|
||||
assert not isinstance(exc_info.value, OmniClientError)
|
||||
assert str(exc_info.value) == "GPU exploded"
|
||||
|
||||
|
||||
def test_omni_generate_py_generator_yields_final_outputs_for_each_request(monkeypatch: pytest.MonkeyPatch):
|
||||
sampling_params = [SamplingParams(max_tokens=8) for _ in range(3)]
|
||||
engine = FakeAsyncOmniEngine(
|
||||
@@ -832,140 +761,6 @@ def test_non_fatal_error_raises_runtime():
|
||||
base._handle_output_message(msg)
|
||||
|
||||
|
||||
def test_non_fatal_client_error_raises_omni_client_error():
|
||||
"""A non-fatal ErrorMessage carrying a 4xx status_code (e.g. a guardrail
|
||||
block) must surface as an OmniClientError with the metadata intact, not a
|
||||
bare RuntimeError. This covers the offline/sync Omni consumer path."""
|
||||
base = _make_base()
|
||||
msg = ErrorMessage(
|
||||
error="Input was blocked by Cosmos3 guardrails.",
|
||||
status_code=400,
|
||||
error_type="BadRequestError",
|
||||
)
|
||||
|
||||
with pytest.raises(OmniClientError) as exc_info:
|
||||
base._handle_output_message(msg)
|
||||
|
||||
assert exc_info.value.status_code == 400
|
||||
assert exc_info.value.error_type == "BadRequestError"
|
||||
assert str(exc_info.value) == "Input was blocked by Cosmos3 guardrails."
|
||||
|
||||
|
||||
_NON_400_CLIENT_ERRORS = [
|
||||
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"),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("status_code", "error_type"), [pytest.param(400, "BadRequestError", id="400")] + _NON_400_CLIENT_ERRORS
|
||||
)
|
||||
def test_client_error_metadata_round_trip_preserves_4xx(status_code: int, error_type: str):
|
||||
original = OmniClientError("blocked", status_code=status_code, error_type=error_type)
|
||||
|
||||
# Outbound: what the broad `except Exception` handlers serialize.
|
||||
carried_status, carried_type = client_error_metadata(original)
|
||||
assert carried_status == status_code
|
||||
assert carried_type == error_type
|
||||
assert is_client_error_status(carried_status)
|
||||
|
||||
# Inbound: reconstruction at the consuming end.
|
||||
rebuilt = client_error_from_metadata("blocked", status_code=carried_status, error_type=carried_type)
|
||||
assert isinstance(rebuilt, OmniClientError)
|
||||
assert rebuilt.status_code == status_code
|
||||
assert rebuilt.error_type == error_type
|
||||
assert str(rebuilt) == "blocked"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS)
|
||||
def test_non_fatal_client_error_preserves_non_400_status(status_code: int, error_type: str):
|
||||
base = _make_base()
|
||||
msg = ErrorMessage(
|
||||
error="client side failure",
|
||||
status_code=status_code,
|
||||
error_type=error_type,
|
||||
)
|
||||
|
||||
with pytest.raises(OmniClientError) as exc_info:
|
||||
base._handle_output_message(msg)
|
||||
|
||||
assert exc_info.value.status_code == status_code
|
||||
assert exc_info.value.error_type == error_type
|
||||
assert str(exc_info.value) == "client side failure"
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"status_code",
|
||||
[
|
||||
pytest.param(None, id="no-status"),
|
||||
pytest.param(399, id="399-below-4xx"),
|
||||
pytest.param(500, id="500-server-error"),
|
||||
pytest.param(503, id="503-service-unavailable"),
|
||||
],
|
||||
)
|
||||
def test_non_fatal_non_4xx_status_raises_runtime(status_code: int | None):
|
||||
base = _make_base()
|
||||
msg = ErrorMessage(
|
||||
error="server side failure",
|
||||
status_code=status_code,
|
||||
error_type="ShouldBeIgnoredForNon4xx",
|
||||
)
|
||||
|
||||
with pytest.raises(RuntimeError) as exc_info:
|
||||
base._handle_output_message(msg)
|
||||
|
||||
assert not isinstance(exc_info.value, OmniClientError)
|
||||
assert str(exc_info.value) == "server side failure"
|
||||
|
||||
|
||||
def _make_enqueue_client_error(
|
||||
status_code: int,
|
||||
error_type: str,
|
||||
error_text: str,
|
||||
) -> Callable[[FakeAsyncOmniEngine, dict[str, Any]], None]:
|
||||
def _enqueue(engine: FakeAsyncOmniEngine, msg: dict[str, Any]) -> None:
|
||||
engine.output_q.put_nowait(
|
||||
ErrorMessage(
|
||||
request_id=msg["request_id"],
|
||||
stage_id=2,
|
||||
error=error_text,
|
||||
status_code=status_code,
|
||||
error_type=error_type,
|
||||
)
|
||||
)
|
||||
|
||||
return _enqueue
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@pytest.mark.parametrize(("status_code", "error_type"), _NON_400_CLIENT_ERRORS)
|
||||
async def test_async_omni_propagates_non_400_client_error_status(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
status_code: int,
|
||||
error_type: str,
|
||||
):
|
||||
error_text = f"blocked with {status_code}"
|
||||
engine = FakeAsyncOmniEngine(
|
||||
stage_metadata=THREE_STAGE_META,
|
||||
on_add_request=_make_enqueue_client_error(status_code, error_type, error_text),
|
||||
)
|
||||
_patch_engine(monkeypatch, engine)
|
||||
|
||||
app = AsyncOmni("dummy-model")
|
||||
try:
|
||||
with pytest.raises(OmniClientError) as exc_info:
|
||||
async for _ in app.generate(prompt="blocked", request_id="req-1"):
|
||||
pass
|
||||
finally:
|
||||
app.shutdown()
|
||||
|
||||
assert exc_info.value.status_code == status_code
|
||||
assert exc_info.value.error_type == error_type
|
||||
assert str(exc_info.value) == error_text
|
||||
|
||||
|
||||
def test_async_omni_errored_property_alive():
|
||||
omni = object.__new__(AsyncOmni)
|
||||
omni.engine = SimpleNamespace(
|
||||
|
||||
+8
-28
@@ -1,6 +1,5 @@
|
||||
"""Synthetic media generation and media/text utilities for tests."""
|
||||
|
||||
import atexit
|
||||
import base64
|
||||
import concurrent.futures
|
||||
import gc
|
||||
@@ -15,6 +14,7 @@ import re
|
||||
import subprocess
|
||||
import tempfile
|
||||
import time
|
||||
import uuid
|
||||
from contextlib import contextmanager
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -588,32 +588,15 @@ def cosine_similarity_text(text1, text2, n: int = 3):
|
||||
return cosine * length_harmony
|
||||
|
||||
|
||||
class _AudioBuffer:
|
||||
"""Minimal replacement for pydub.AudioSegment used by test helpers."""
|
||||
def _merge_base64_audio_to_segment(base64_list: list[str]):
|
||||
from pydub import AudioSegment
|
||||
|
||||
def __init__(self, data: np.ndarray, sample_rate: int):
|
||||
self.data = data
|
||||
self.sample_rate = sample_rate
|
||||
|
||||
def export(self, buf: io.BytesIO, format: str = "wav"):
|
||||
sf.write(buf, self.data, self.sample_rate, format=format.upper())
|
||||
buf.seek(0)
|
||||
|
||||
|
||||
def _merge_base64_audio_to_segment(base64_list: list[str]) -> _AudioBuffer:
|
||||
from vllm.multimodal.media.audio import AudioMediaIO
|
||||
|
||||
io_ = AudioMediaIO()
|
||||
chunks: list[np.ndarray] = []
|
||||
sample_rate: int | None = None
|
||||
merged = None
|
||||
for b64 in base64_list:
|
||||
raw = base64.b64decode(b64.split(",", 1)[-1])
|
||||
waveform, sr = io_.load_bytes(raw)
|
||||
if sample_rate is None:
|
||||
sample_rate = int(sr)
|
||||
chunks.append(waveform)
|
||||
merged = np.concatenate(chunks) if chunks else np.array([], dtype=np.float32)
|
||||
return _AudioBuffer(merged, sample_rate or 16000)
|
||||
seg = AudioSegment.from_file(io.BytesIO(raw))
|
||||
merged = seg if merged is None else merged + seg
|
||||
return merged
|
||||
|
||||
|
||||
@contextmanager
|
||||
@@ -683,10 +666,7 @@ def convert_audio_file_to_text(output_path: str) -> str:
|
||||
|
||||
|
||||
def convert_audio_bytes_to_text(raw_bytes: bytes) -> str:
|
||||
output_fd, output_path = tempfile.mkstemp(prefix="test_", suffix=".wav")
|
||||
os.close(output_fd)
|
||||
if os.environ.get("VLLM_OMNI_KEEP_REQUEST_MEDIA", "").lower() not in ("1", "true", "yes"):
|
||||
atexit.register(Path(output_path).unlink, missing_ok=True)
|
||||
output_path = f"./test_{uuid.uuid4().hex}.wav"
|
||||
data, samplerate = sf.read(io.BytesIO(raw_bytes))
|
||||
sf.write(output_path, data, samplerate, format="WAV", subtype="PCM_16")
|
||||
print(f"audio data is saved: {output_path}")
|
||||
|
||||
+11
-26
@@ -1473,8 +1473,6 @@ class OpenAIClientHandler:
|
||||
|
||||
- ``send_frames``: optional ``str`` or sequence of ``str`` raw WebSocket text frames (omit when the server
|
||||
speaks first, e.g. ``/v1/realtime`` rejection path).
|
||||
- ``ws_skip_types``: optional event ``type`` strings to ignore while waiting for the first matching frame
|
||||
(e.g. ``["session.created"]`` on ``/v1/realtime``).
|
||||
- ``timeout``: seconds to wait for the first inbound text frame (default ``120``).
|
||||
- ``ws_max_size``: passed through as ``max_size`` to :func:`websockets.connect` when the key is present.
|
||||
"""
|
||||
@@ -1488,7 +1486,6 @@ class OpenAIClientHandler:
|
||||
|
||||
timeout = float(cfg.get("timeout", 120.0))
|
||||
uri = self._build_ws_url(path)
|
||||
skip_types = set(cfg.get("ws_skip_types") or [])
|
||||
|
||||
connect_kw: dict[str, Any] = {}
|
||||
if "ws_max_size" in cfg:
|
||||
@@ -1500,19 +1497,16 @@ class OpenAIClientHandler:
|
||||
async with websockets.connect(uri, **connect_kw) as ws:
|
||||
for frame in frames:
|
||||
await ws.send(frame)
|
||||
while True:
|
||||
raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
|
||||
if not isinstance(raw, str):
|
||||
raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
|
||||
if not isinstance(data, dict):
|
||||
raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
|
||||
if skip_types and data.get("type") in skip_types:
|
||||
continue
|
||||
return WebSocketJsonResponse(json_body=data)
|
||||
raw = await asyncio.wait_for(ws.recv(), timeout=timeout)
|
||||
if not isinstance(raw, str):
|
||||
raise AssertionError(f"Expected JSON text frame from {uri}, got {type(raw).__name__}")
|
||||
try:
|
||||
data = json.loads(raw)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise AssertionError(f"Expected JSON text frame from {uri}, body={raw[:500]!r}") from exc
|
||||
if not isinstance(data, dict):
|
||||
raise AssertionError(f"Expected JSON object from {uri}, got {type(data).__name__}")
|
||||
return WebSocketJsonResponse(json_body=data)
|
||||
|
||||
resp = asyncio.run(_recv_first_json_object())
|
||||
_run_ws_expectations_from_request_config(cfg, resp)
|
||||
@@ -1773,16 +1767,7 @@ class OpenAIClientHandler:
|
||||
# Qwen3-TTS custom fields, forwarded via extra_body.
|
||||
extra_body: dict[str, Any] = {}
|
||||
# Keep this list aligned with vllm_omni.entrypoints.openai.protocol.audio params.
|
||||
for key in (
|
||||
"task_type",
|
||||
"ref_text",
|
||||
"ref_audio",
|
||||
"language",
|
||||
"max_new_tokens",
|
||||
"seed",
|
||||
"instructions",
|
||||
"speed",
|
||||
):
|
||||
for key in ("task_type", "ref_text", "ref_audio", "language", "max_new_tokens", "seed"):
|
||||
if key in request_config:
|
||||
extra_body[key] = request_config[key]
|
||||
|
||||
|
||||
@@ -84,11 +84,6 @@ def _make_code2wav_model(
|
||||
token_mel_ratio=2 if with_stride_cfg else 0,
|
||||
)
|
||||
model.code2wav = _DummyCode2Wav(vocab_size=4, num_samples=num_samples, outputs=outputs)
|
||||
# Short-circuit the lazy TensorRT estimator swap: these tests exercise the
|
||||
# forward audio logic, not the TRT path. On a GPU CI runner the swap would
|
||||
# otherwise run and dereference ``self.model_dir`` (only set in __init__,
|
||||
# which this fixture bypasses via object.__new__).
|
||||
model._code2wav_trt_done = True
|
||||
model.source_cache_len = 4
|
||||
model.speech_window = torch.hamming_window(8, periodic=False)
|
||||
model._stream_audio_cache_by_req = {}
|
||||
|
||||
@@ -1,185 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
||||
from vllm_omni.model_executor.models.common.ming.audio_vae import AudioVAEConfig
|
||||
from vllm_omni.model_executor.models.common.ming.fm import Solver
|
||||
from vllm_omni.model_executor.models.ming_tts.constants import (
|
||||
AGGREGATOR_HIDDEN_SIZE,
|
||||
HISTORY_PATCH_SIZE,
|
||||
LATENT_DIM,
|
||||
LLM_HIDDEN_SIZE,
|
||||
LLM_VOCAB_SIZE,
|
||||
PATCH_SIZE,
|
||||
SAMPLE_RATE,
|
||||
VAE_PATCH_SIZE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.ming_tts.flowloss_head import FlowLoss
|
||||
from vllm_omni.model_executor.models.ming_tts.ming_tts_llm import MingLLMModel
|
||||
from vllm_omni.model_executor.models.ming_tts.validation import validate_ming_tts_config
|
||||
from vllm_omni.model_executor.models.output_templates import OmniOutput
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.tts]
|
||||
|
||||
|
||||
def test_ming_tts_audio_vae_uses_common_config():
|
||||
"""AudioVAEConfig is shared by Ming dense and Ming flash modules."""
|
||||
cfg = AudioVAEConfig(sample_rate=16000, patch_size=-1)
|
||||
|
||||
assert cfg.sample_rate == 16000
|
||||
assert cfg.patch_size == -1
|
||||
|
||||
|
||||
def test_ming_tts_cfm_solver_uses_common_implementation():
|
||||
"""Ming dense imports the shared solver implementation directly."""
|
||||
assert Solver.__module__ == "vllm_omni.model_executor.models.common.ming.fm"
|
||||
|
||||
|
||||
def test_ming_tts_flowloss_preserves_checkpoint_prefix():
|
||||
flowloss = FlowLoss(z_channels=4, llm_cond_dim=8, hidden_size=16, depth=1, num_heads=2)
|
||||
|
||||
assert any(name.startswith("cfm.model.") for name in flowloss.state_dict())
|
||||
|
||||
|
||||
def test_ming_dense_validation_rejects_semantic_audio_vae_config():
|
||||
"""Dense 0.5B validation rejects semantic AudioVAE configs."""
|
||||
cfg = SimpleNamespace(
|
||||
audio_dummy_token_id=151705,
|
||||
audio_eos_token_id=151704,
|
||||
text_eos_token_id=151669,
|
||||
audio_tokenizer_config=AudioVAEConfig(
|
||||
sample_rate=SAMPLE_RATE,
|
||||
patch_size=VAE_PATCH_SIZE,
|
||||
semantic_module_kwargs={"whisper_encoder": {}},
|
||||
enc_kwargs={"latent_dim": LATENT_DIM, "input_dim": 882, "hop_size": 882},
|
||||
dec_kwargs={"latent_dim": LATENT_DIM, "output_dim": 882},
|
||||
),
|
||||
latent_dim=LATENT_DIM,
|
||||
patch_size=PATCH_SIZE,
|
||||
history_patch_size=HISTORY_PATCH_SIZE,
|
||||
llm_hidden_size=LLM_HIDDEN_SIZE,
|
||||
llm_vocab_size=LLM_VOCAB_SIZE,
|
||||
sample_rate=SAMPLE_RATE,
|
||||
vae_patch_size=VAE_PATCH_SIZE,
|
||||
llm_config={"hidden_size": LLM_HIDDEN_SIZE},
|
||||
aggregator_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE},
|
||||
ditar_config={"hidden_size": AGGREGATOR_HIDDEN_SIZE},
|
||||
latent_chunk_size=1,
|
||||
latent_left_context=0,
|
||||
max_decode_steps=1,
|
||||
stop_head_threshold=0.5,
|
||||
stop_head_min_steps=0,
|
||||
)
|
||||
|
||||
with pytest.raises(ValueError, match="semantic_module_kwargs"):
|
||||
validate_ming_tts_config(cfg)
|
||||
|
||||
|
||||
def test_ming_instruction_parser_preserves_dense_and_flash_defaults():
|
||||
"""Ming dense and Ming flash keep distinct instruction defaults."""
|
||||
serving = object.__new__(OmniOpenAIServingSpeech)
|
||||
serving.uploaded_speakers = {"uploaded": {}}
|
||||
|
||||
dense_plain = serving._parse_ming_instruction(SimpleNamespace(instructions="calm", language=None, voice=None))
|
||||
assert dense_plain == "calm"
|
||||
|
||||
dense_with_fields = serving._parse_ming_instruction(
|
||||
SimpleNamespace(instructions="calm", language="Auto", voice="灵小甄")
|
||||
)
|
||||
assert dense_with_fields == {"IP": "灵小甄", "风格": "calm"}
|
||||
|
||||
flash_fields = serving._parse_ming_instruction_fields(
|
||||
SimpleNamespace(instructions="calm", language="粤语", voice="灵小甄")
|
||||
)
|
||||
assert flash_fields == {"风格": "calm"}
|
||||
|
||||
|
||||
def _make_ming_logits_model(vocab_size=8):
|
||||
model = object.__new__(MingLLMModel)
|
||||
model.ming_config = SimpleNamespace(
|
||||
llm_vocab_size=vocab_size,
|
||||
max_decode_steps=1,
|
||||
stop_head_min_steps=0,
|
||||
text_eos_token_id=7,
|
||||
)
|
||||
model._last_text_mode = False
|
||||
model._last_ming_next_token_ids = None
|
||||
return model
|
||||
|
||||
|
||||
def test_ming_compute_logits_uses_cached_forced_next_token_ids():
|
||||
model = _make_ming_logits_model()
|
||||
model._last_ming_next_token_ids = [2, 5]
|
||||
|
||||
logits = MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace())
|
||||
|
||||
assert logits.shape == (2, 8)
|
||||
assert logits[0, 2].item() == 0.0
|
||||
assert logits[1, 5].item() == 0.0
|
||||
assert torch.isneginf(logits[0, [0, 1, 3, 4, 5, 6, 7]]).all()
|
||||
assert torch.isneginf(logits[1, [0, 1, 2, 3, 4, 6, 7]]).all()
|
||||
assert model._last_ming_next_token_ids is None
|
||||
|
||||
|
||||
def test_ming_compute_logits_falls_back_to_dummy_token_id():
|
||||
model = _make_ming_logits_model()
|
||||
|
||||
logits = MingLLMModel.compute_logits(model, torch.zeros((1, 4)), SimpleNamespace())
|
||||
assert logits.shape == (1, 8)
|
||||
assert logits[0, 7].item() == 0.0
|
||||
assert torch.isneginf(logits[0, [0, 1, 2, 3, 4, 5, 6]]).all()
|
||||
|
||||
|
||||
def test_ming_forward_non_decode_return_clears_cached_forced_next_token_ids():
|
||||
class FakeBackbone:
|
||||
def __call__(self, **kwargs):
|
||||
return kwargs["inputs_embeds"]
|
||||
|
||||
model = _make_ming_logits_model()
|
||||
model.model = FakeBackbone()
|
||||
model._last_ming_next_token_ids = [2]
|
||||
|
||||
output = MingLLMModel.forward(
|
||||
model,
|
||||
input_ids=torch.tensor([1]),
|
||||
positions=torch.tensor([0]),
|
||||
inputs_embeds=torch.zeros((1, 4)),
|
||||
model_intermediate_buffer=[],
|
||||
)
|
||||
|
||||
assert isinstance(output, OmniOutput)
|
||||
assert output.multimodal_outputs is None
|
||||
assert model._last_ming_next_token_ids is None
|
||||
|
||||
|
||||
def test_ming_compute_logits_rejects_forced_token_batch_mismatch():
|
||||
model = _make_ming_logits_model()
|
||||
model._last_ming_next_token_ids = [2]
|
||||
|
||||
with pytest.raises(RuntimeError, match="batch mismatch"):
|
||||
MingLLMModel.compute_logits(model, torch.zeros((2, 4)), SimpleNamespace())
|
||||
|
||||
|
||||
def test_ming_compute_logits_text_mode_delegates_to_backbone():
|
||||
class FakeBackbone:
|
||||
def __init__(self):
|
||||
self.hidden_states = None
|
||||
|
||||
def compute_logits(self, hidden_states):
|
||||
self.hidden_states = hidden_states
|
||||
return torch.ones((hidden_states.shape[0], 3))
|
||||
|
||||
model = _make_ming_logits_model(vocab_size=3)
|
||||
model._last_text_mode = True
|
||||
model.model = FakeBackbone()
|
||||
hidden_states = torch.zeros((2, 4))
|
||||
|
||||
logits = MingLLMModel.compute_logits(model, hidden_states, SimpleNamespace())
|
||||
|
||||
assert torch.equal(model.model.hidden_states, hidden_states)
|
||||
assert torch.equal(logits, torch.ones((2, 3)))
|
||||
Whitespace-only changes.
@@ -1,361 +0,0 @@
|
||||
"""Tests for MossTTSCUDAGraphCodecWrapper numerical equivalence.
|
||||
|
||||
Verifies that CUDA Graph-accelerated decoding produces results equivalent
|
||||
to eager mode, with special attention to the two-argument _decode interface
|
||||
(codes [NQ, 1, T] + lengths [1]) and the NQ-first input convention.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from vllm_omni.model_executor.models.moss_tts.audio_tokenizer import (
|
||||
MossAudioTokenizerDecoderOutput,
|
||||
)
|
||||
from vllm_omni.model_executor.models.moss_tts.moss_codec_cudagraph import (
|
||||
MossTTSCUDAGraphCodecWrapper,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required")]
|
||||
|
||||
DEVICE = torch.device("cuda:0")
|
||||
NUM_QUANTIZERS = 8
|
||||
DOWNSAMPLE_RATE = 4 # synthetic; real checkpoint uses 1920
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Synthetic codec model
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SyntheticCodecModel(nn.Module):
|
||||
"""Minimal stand-in for MossAudioTokenizerModel.
|
||||
|
||||
Exposes the same two-argument _decode(codes, lengths) interface and
|
||||
returns a MossAudioTokenizerDecoderOutput, mirroring the real model.
|
||||
|
||||
Input:
|
||||
codes: [NQ, B, T] long — RVQ codes
|
||||
lengths: [B] long — valid frame counts
|
||||
Output:
|
||||
MossAudioTokenizerDecoderOutput(audio=[B, 1, T*upsample], audio_lengths=[B])
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_quantizers: int = NUM_QUANTIZERS,
|
||||
downsample_rate: int = DOWNSAMPLE_RATE,
|
||||
):
|
||||
super().__init__()
|
||||
self.downsample_rate = downsample_rate
|
||||
hidden = 32
|
||||
# embed: treat NQ quantizers as input channels for Conv1d
|
||||
self.embed = nn.Conv1d(num_quantizers, hidden, kernel_size=3, padding=1)
|
||||
self.conv = nn.Conv1d(hidden, hidden, kernel_size=3, padding=1)
|
||||
self.upsample = nn.ConvTranspose1d(hidden, 1, kernel_size=downsample_rate, stride=downsample_rate)
|
||||
|
||||
def _decode(self, codes: torch.Tensor, lengths: torch.Tensor) -> MossAudioTokenizerDecoderOutput:
|
||||
"""codes: [NQ, B, T], lengths: [B] → audio [B, 1, T*upsample]."""
|
||||
nq, b, t = codes.shape
|
||||
# treat NQ as channel dim → [B, NQ, T] for Conv1d (expects [B, C, T])
|
||||
x = codes.permute(1, 0, 2).to(dtype=self.embed.weight.dtype) # [B, NQ, T]
|
||||
x = torch.relu(self.embed(x)) # [B, hidden, T]
|
||||
x = torch.relu(self.conv(x)) # [B, hidden, T]
|
||||
audio = self.upsample(x) # [B, 1, T * upsample]
|
||||
audio_lengths = lengths * self.downsample_rate
|
||||
return MossAudioTokenizerDecoderOutput(audio=audio, audio_lengths=audio_lengths)
|
||||
|
||||
def batch_decode(
|
||||
self,
|
||||
codes_list: list[torch.Tensor],
|
||||
num_quantizers: int | None = None,
|
||||
) -> MossAudioTokenizerDecoderOutput:
|
||||
"""Mirrors MossAudioTokenizerModel.batch_decode (fallback path)."""
|
||||
device = codes_list[0].device
|
||||
nq = num_quantizers or codes_list[0].shape[0]
|
||||
max_t = max(c.shape[-1] for c in codes_list)
|
||||
codes = torch.zeros(nq, len(codes_list), max_t, device=device, dtype=torch.long)
|
||||
lengths = torch.zeros(len(codes_list), device=device, dtype=torch.long)
|
||||
for i, c in enumerate(codes_list):
|
||||
codes[:nq, i, : c.shape[-1]] = c[:nq]
|
||||
lengths[i] = c.shape[-1]
|
||||
return self._decode(codes, lengths)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def model():
|
||||
torch.manual_seed(0)
|
||||
return SyntheticCodecModel().to(DEVICE).eval()
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def wrapper(model):
|
||||
w = MossTTSCUDAGraphCodecWrapper(
|
||||
model=model,
|
||||
capture_sizes=[25, 50, 100],
|
||||
num_quantizers=NUM_QUANTIZERS,
|
||||
enabled=True,
|
||||
)
|
||||
w.warmup(DEVICE)
|
||||
return w
|
||||
|
||||
|
||||
def _random_codes(t: int, device: torch.device = DEVICE) -> torch.Tensor:
|
||||
"""Return [NQ, T] long tensor (single-request convention)."""
|
||||
return torch.randint(0, 64, (NUM_QUANTIZERS, t), dtype=torch.long, device=device)
|
||||
|
||||
|
||||
def _eager_decode(model: SyntheticCodecModel, codes_nq_t: torch.Tensor) -> MossAudioTokenizerDecoderOutput:
|
||||
"""Run the model directly in eager mode (reference output)."""
|
||||
t = codes_nq_t.shape[-1]
|
||||
codes_nq_1_t = codes_nq_t.unsqueeze(1) # [NQ, 1, T]
|
||||
lengths = torch.tensor([t], dtype=torch.long, device=codes_nq_t.device)
|
||||
with torch.no_grad():
|
||||
return model._decode(codes_nq_1_t, lengths)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. Exact-size inputs — must be bit-identical
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t", [25, 50, 100])
|
||||
def test_exact_size_bit_identical(model, wrapper, t):
|
||||
codes = _random_codes(t)
|
||||
ref = _eager_decode(model, codes)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. Padded inputs — output trimmed to actual length
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t", [10, 30, 47, 73, 99])
|
||||
def test_padded_output_shape(model, wrapper, t):
|
||||
codes = _random_codes(t)
|
||||
ref = _eager_decode(model, codes)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
expected_len = t * DOWNSAMPLE_RATE
|
||||
assert out.audio.shape[-1] == expected_len, f"got {out.audio.shape[-1]}, want {expected_len}"
|
||||
assert out.audio.shape == ref.audio.shape
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t", [10, 30, 47, 73, 99])
|
||||
def test_padded_interior_positions_close(model, wrapper, t):
|
||||
"""Positions well away from the zero-padding boundary must be numerically close."""
|
||||
codes = _random_codes(t)
|
||||
ref = _eager_decode(model, codes)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
boundary = 2 * DOWNSAMPLE_RATE
|
||||
if ref.audio.shape[-1] > boundary:
|
||||
torch.testing.assert_close(
|
||||
out.audio[..., :-boundary],
|
||||
ref.audio[..., :-boundary],
|
||||
atol=1e-5,
|
||||
rtol=1e-5,
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. Fallback to eager (T > all capture sizes)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t", [101, 150, 200])
|
||||
def test_fallback_exact_match(model, wrapper, t):
|
||||
codes = _random_codes(t)
|
||||
ref = _eager_decode(model, codes)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 4. Static buffer not aliased across calls
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_output_not_aliased_after_later_replay(model, wrapper):
|
||||
"""clone() in decode() must prevent later replays from overwriting the result."""
|
||||
codes_a = _random_codes(50)
|
||||
codes_b = _random_codes(50)
|
||||
with torch.no_grad():
|
||||
out_a = wrapper.decode(codes_a)
|
||||
saved = out_a.audio.clone()
|
||||
_ = wrapper.decode(codes_b) # replay overwrites static buffer
|
||||
torch.testing.assert_close(out_a.audio, saved, atol=0, rtol=0)
|
||||
|
||||
|
||||
def test_deterministic_across_calls(model, wrapper):
|
||||
codes = _random_codes(30)
|
||||
with torch.no_grad():
|
||||
out1 = wrapper.decode(codes)
|
||||
out2 = wrapper.decode(codes)
|
||||
torch.testing.assert_close(out1.audio, out2.audio, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 5. Disabled wrapper falls back to eager
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_disabled_falls_back_to_eager(model, wrapper):
|
||||
codes = _random_codes(30)
|
||||
ref = _eager_decode(model, codes)
|
||||
wrapper.enabled = False
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
wrapper.enabled = True
|
||||
torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 6. audio_lengths is consistent with audio shape
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t", [25, 30, 50, 99, 100])
|
||||
def test_audio_lengths_consistent(model, wrapper, t):
|
||||
codes = _random_codes(t)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
assert out.audio_lengths is not None
|
||||
assert int(out.audio_lengths[0].item()) == out.audio.shape[-1]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 7. NQ-first input layout: codes_nq_t is [NQ, T], not [T, NQ]
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_nq_first_layout_matches_eager(model, wrapper):
|
||||
"""Verify that wrapper correctly interprets [NQ, T] (NQ-first) layout."""
|
||||
t = 25
|
||||
codes = _random_codes(t) # [NQ, T]
|
||||
ref = _eager_decode(model, codes)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
torch.testing.assert_close(out.audio, ref.audio, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 8. Non-contiguous input — copy_() must handle strided tensors correctly
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t", [20, 47])
|
||||
def test_noncontiguous_input_matches_contiguous(model, wrapper, t):
|
||||
"""Padded path uses copy_() on a slice; must produce the same result
|
||||
whether codes_nq_t is contiguous or a non-contiguous view (e.g. a column
|
||||
slice of a larger buffer, which is the realistic production layout when the
|
||||
caller slices out one request from a batched tensor)."""
|
||||
codes_contiguous = _random_codes(t)
|
||||
# Build a non-contiguous view: embed in a larger [NQ, 2, T] buffer and
|
||||
# take column 0. The resulting tensor has stride (2*T, 2, 1) → non-contiguous.
|
||||
buf = torch.zeros(NUM_QUANTIZERS, 2, t, dtype=torch.long, device=DEVICE)
|
||||
buf[:, 0, :] = codes_contiguous
|
||||
codes_noncontig = buf[:, 0, :] # [NQ, T], non-contiguous
|
||||
assert not codes_noncontig.is_contiguous()
|
||||
|
||||
with torch.no_grad():
|
||||
out_c = wrapper.decode(codes_contiguous)
|
||||
out_nc = wrapper.decode(codes_noncontig)
|
||||
torch.testing.assert_close(out_c.audio, out_nc.audio, atol=0, rtol=0)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 9. Zero-padding uses token 0 — must be a valid codebook index (no OOB)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_zero_padding_no_cuda_error(model, wrapper):
|
||||
"""Padded inputs fill unused frames with code 0. Verify that the zero-
|
||||
padded region does not corrupt the output shape.
|
||||
|
||||
Note: SyntheticCodecModel uses Conv1d (not nn.Embedding), so it cannot
|
||||
trigger a real CUDA index-out-of-bounds on large token ids. This test
|
||||
validates shape correctness only; OOB safety on the real codec requires
|
||||
an integration test with MossAudioTokenizerModel."""
|
||||
t = 7 # well below the 25-frame bucket → large zero-padded region
|
||||
codes = _random_codes(t)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
torch.accelerator.synchronize() # surface any deferred CUDA errors
|
||||
assert out.audio.shape[-1] == t * DOWNSAMPLE_RATE
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 10. Shape alignment: graph path and eager fallback return identical shapes
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("t", [25, 30, 101])
|
||||
def test_graph_and_eager_shape_identical(model, wrapper, t):
|
||||
"""Graph replay (t<=100) and eager fallback (t>100) must return tensors
|
||||
with the same number of dimensions and the same batch/channel layout so
|
||||
downstream consumers can treat both paths uniformly."""
|
||||
codes = _random_codes(t)
|
||||
ref = _eager_decode(model, codes)
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
assert out.audio.ndim == ref.audio.ndim, f"ndim mismatch: graph={out.audio.ndim} eager={ref.audio.ndim}"
|
||||
assert out.audio.shape[:-1] == ref.audio.shape[:-1], (
|
||||
f"batch/channel shape mismatch: graph={out.audio.shape} eager={ref.audio.shape}"
|
||||
)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 11. Half-precision weights: graph and eager must agree under fp16/bf16
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
|
||||
def test_half_precision_graph_eager_agree(dtype):
|
||||
"""CUDA Graph must produce outputs consistent with eager when model weights
|
||||
are in half precision. The wrapper itself stores codes as long and lengths
|
||||
as long; only the model's conv weights change dtype."""
|
||||
torch.manual_seed(0)
|
||||
model_hp = SyntheticCodecModel().to(device=DEVICE, dtype=dtype).eval()
|
||||
wrapper_hp = MossTTSCUDAGraphCodecWrapper(
|
||||
model=model_hp,
|
||||
capture_sizes=[25, 50],
|
||||
num_quantizers=NUM_QUANTIZERS,
|
||||
enabled=True,
|
||||
)
|
||||
wrapper_hp.warmup(DEVICE)
|
||||
|
||||
codes = _random_codes(25)
|
||||
ref = _eager_decode(model_hp, codes)
|
||||
with torch.no_grad():
|
||||
out = wrapper_hp.decode(codes)
|
||||
# Half-precision accumulation introduces small numerical error; use
|
||||
# tolerances that match typical fp16/bf16 rounding.
|
||||
atol = 1e-2 if dtype == torch.float16 else 5e-2
|
||||
torch.testing.assert_close(out.audio.float(), ref.audio.float(), atol=atol, rtol=1e-3)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 12. Empty input (T=0): wrapper must not crash and return zero-length audio
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def test_empty_input_t0(model, wrapper):
|
||||
"""T=0 input must be handled gracefully: no exception, audio length == 0.
|
||||
|
||||
The outer MossTTSCodecDecoder guards against empty segments, but the
|
||||
wrapper should also handle T=0 without crashing so it is safe to call
|
||||
directly."""
|
||||
codes = _random_codes(0) # [NQ, 0]
|
||||
with torch.no_grad():
|
||||
out = wrapper.decode(codes)
|
||||
assert out.audio.shape[-1] == 0, f"expected empty audio, got shape {out.audio.shape}"
|
||||
@@ -921,7 +921,6 @@ def test_cosyvoice3_full_payload_replace_keys_present() -> None:
|
||||
def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
|
||||
"""Smoke: ming_flash_omni token-only carries voice metadata."""
|
||||
from vllm_omni.model_executor.stage_input_processors.ming_flash_omni import (
|
||||
thinker2talker,
|
||||
thinker2talker_token_only,
|
||||
)
|
||||
|
||||
@@ -939,15 +938,14 @@ def test_ming_flash_omni_thinker2talker_token_only_smoke() -> None:
|
||||
|
||||
src = [_Wrap("hello world")]
|
||||
prompt = _Prompt({"voice_name": "ZH_FEMALE", "prompt_text": "ref text"})
|
||||
for func in (thinker2talker, thinker2talker_token_only):
|
||||
out = func(src, prompt=prompt)
|
||||
assert len(out) == 1
|
||||
assert out[0]["prompt_token_ids"] == [0]
|
||||
info = out[0]["additional_information"]
|
||||
assert info["text"] == "hello world"
|
||||
assert info["voice_name"] == "ZH_FEMALE"
|
||||
assert info["prompt_text"] == "ref text"
|
||||
assert info["ming_task"] == "omni"
|
||||
out = thinker2talker_token_only(src, prompt=prompt)
|
||||
assert len(out) == 1
|
||||
assert out[0]["prompt_token_ids"] == [0] # talker self-tokenizes; dummy id
|
||||
info = out[0]["additional_information"]
|
||||
assert info["text"] == "hello world"
|
||||
assert info["voice_name"] == "ZH_FEMALE"
|
||||
assert info["prompt_text"] == "ref text"
|
||||
assert info["ming_task"] == "omni"
|
||||
|
||||
|
||||
def test_qwen2_5_omni_thinker2talker_token_only_smoke() -> None:
|
||||
|
||||
Whitespace-only changes.
@@ -1,918 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project
|
||||
|
||||
"""Tests for StepAudioReasoningParser.
|
||||
|
||||
Step-Audio supports two representations of thinking markers:
|
||||
|
||||
1. **Special tokens**: ``<|THINK_START|>`` and ``<|THINK_END|>``
|
||||
(single-token IDs in vocab).
|
||||
|
||||
2. **Text markers**: ``<think>`` and ``</think>``
|
||||
(multi-token sequences for </think>, e.g. [522, 26865]).
|
||||
|
||||
These tests use standalone mock objects so they can run without
|
||||
installing the full vLLM/torch stack.
|
||||
"""
|
||||
|
||||
import sys
|
||||
from types import ModuleType
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Bootstrap minimal vLLM stubs so the parser can be imported
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _install_vllm_stubs() -> dict[str, ModuleType | None]:
|
||||
"""Temporarily install minimal stub modules for the vllm imports the
|
||||
parser needs, so it can be imported without the full vLLM/torch stack.
|
||||
|
||||
Returns a snapshot of the original ``sys.modules`` entries that were
|
||||
replaced. The caller MUST pass it to ``_restore_vllm_stubs`` once the
|
||||
parser is imported. These stubs are process-global: if they leak, they
|
||||
shadow the real vllm modules for every other test in the same worker
|
||||
(and any forked server subprocess), e.g. replacing
|
||||
``vllm.entrypoints.openai.engine.protocol`` with a stub that lacks
|
||||
``ErrorResponse`` and breaking unrelated server-startup tests.
|
||||
"""
|
||||
stubs: dict[str, ModuleType] = {}
|
||||
|
||||
# DeltaMessage stub
|
||||
delta_mod = ModuleType("vllm.entrypoints.openai.engine.protocol")
|
||||
|
||||
class DeltaMessage:
|
||||
__slots__ = ("reasoning", "content")
|
||||
|
||||
def __init__(self, reasoning=None, content=None):
|
||||
self.reasoning = reasoning
|
||||
self.content = content
|
||||
|
||||
delta_mod.DeltaMessage = DeltaMessage
|
||||
stubs["vllm.entrypoints.openai.engine.protocol"] = delta_mod
|
||||
|
||||
# ReasoningParser stub
|
||||
reasoning_mod = ModuleType("vllm.reasoning")
|
||||
|
||||
class ReasoningParser:
|
||||
def __init__(self, tokenizer, *args, **kwargs):
|
||||
self.model_tokenizer = tokenizer
|
||||
|
||||
@property
|
||||
def vocab(self):
|
||||
if not self.model_tokenizer:
|
||||
return {}
|
||||
return self.model_tokenizer.get_vocab()
|
||||
|
||||
class _ReasoningParserManager:
|
||||
reasoning_parsers = {}
|
||||
lazy_parsers = {}
|
||||
|
||||
@classmethod
|
||||
def register_lazy_module(cls, name, module_path, class_name):
|
||||
cls.lazy_parsers[name] = (module_path, class_name)
|
||||
|
||||
reasoning_mod.ReasoningParser = ReasoningParser
|
||||
reasoning_mod.ReasoningParserManager = _ReasoningParserManager
|
||||
stubs["vllm.reasoning"] = reasoning_mod
|
||||
stubs["vllm.reasoning.abs_reasoning_parsers"] = reasoning_mod
|
||||
|
||||
# TokenizerLike stub
|
||||
tokenizers_mod = ModuleType("vllm.tokenizers")
|
||||
tokenizers_mod.TokenizerLike = object
|
||||
stubs["vllm.tokenizers"] = tokenizers_mod
|
||||
|
||||
# ChatCompletionRequest / ResponsesRequest stubs
|
||||
chat_mod = ModuleType("vllm.entrypoints.openai.chat_completion.protocol")
|
||||
chat_mod.ChatCompletionRequest = MagicMock
|
||||
stubs["vllm.entrypoints.openai.chat_completion.protocol"] = chat_mod
|
||||
responses_mod = ModuleType("vllm.entrypoints.openai.responses.protocol")
|
||||
responses_mod.ResponsesRequest = MagicMock
|
||||
stubs["vllm.entrypoints.openai.responses.protocol"] = responses_mod
|
||||
|
||||
# Parent packages: only stub the ones not already importable, so we never
|
||||
# clobber a real (installed) vllm package tree.
|
||||
for mod_path in [
|
||||
"vllm",
|
||||
"vllm.entrypoints",
|
||||
"vllm.entrypoints.openai",
|
||||
"vllm.entrypoints.openai.engine",
|
||||
"vllm.entrypoints.openai.chat_completion",
|
||||
"vllm.entrypoints.openai.responses",
|
||||
]:
|
||||
if mod_path not in sys.modules and mod_path not in stubs:
|
||||
stubs[mod_path] = ModuleType(mod_path)
|
||||
|
||||
saved = {name: sys.modules.get(name) for name in stubs}
|
||||
sys.modules.update(stubs)
|
||||
return saved
|
||||
|
||||
|
||||
def _restore_vllm_stubs(saved: dict[str, ModuleType | None]) -> None:
|
||||
"""Undo ``_install_vllm_stubs`` so ``sys.modules`` is left untouched."""
|
||||
for name, original in saved.items():
|
||||
if original is None:
|
||||
sys.modules.pop(name, None)
|
||||
else:
|
||||
sys.modules[name] = original
|
||||
|
||||
|
||||
# Install stubs only for the duration of the import. The parser binds these
|
||||
# stub classes into its own module namespace at import time, so it keeps using
|
||||
# them at runtime even after sys.modules is restored below.
|
||||
_saved_modules = _install_vllm_stubs()
|
||||
try:
|
||||
from vllm_omni.reasoning.step_audio_reasoning_parser import ( # noqa: E402
|
||||
StepAudioReasoningParser,
|
||||
)
|
||||
finally:
|
||||
_restore_vllm_stubs(_saved_modules)
|
||||
|
||||
# Short aliases for marker constants
|
||||
TS = StepAudioReasoningParser.THINK_START_TEXT # <think>
|
||||
TE = StepAudioReasoningParser.THINK_END_TEXT # </think>
|
||||
TSS = StepAudioReasoningParser.THINK_START_SPECIAL # <|THINK_START|>
|
||||
TES = StepAudioReasoningParser.THINK_END_SPECIAL # <|THINK_END|>
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Mock tokenizers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SingleTokenTokenizer:
|
||||
"""Tokenizer where both ``<think>`` and ``</think>`` are single tokens.
|
||||
|
||||
This exercises the fast (token-ID) path.
|
||||
Also includes special-token forms ``<|THINK_START|>`` and ``<|THINK_END|>``.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._vocab = {
|
||||
TS: 100, # <think>
|
||||
TE: 101, # </think>
|
||||
TSS: 151669, # <|THINK_START|>
|
||||
TES: 151670, # <|THINK_END|>
|
||||
"Hello": 1,
|
||||
" world": 2,
|
||||
"reasoning": 3,
|
||||
" content": 4,
|
||||
"\n": 5,
|
||||
"</": 6,
|
||||
"think>": 7,
|
||||
"some text": 8,
|
||||
}
|
||||
self._reverse = {v: k for k, v in self._vocab.items()}
|
||||
|
||||
def get_vocab(self):
|
||||
return dict(self._vocab)
|
||||
|
||||
def decode(self, ids, skip_special_tokens=False):
|
||||
return "".join(self._reverse.get(i, "") for i in ids)
|
||||
|
||||
def tokenize(self, text):
|
||||
if not text:
|
||||
return []
|
||||
return [text]
|
||||
|
||||
|
||||
class MultiTokenTokenizer:
|
||||
"""Tokenizer where ``</think>`` is NOT a single token.
|
||||
|
||||
``<think>`` is in the vocab (ID 100), but ``</think>`` is NOT.
|
||||
Instead, ``</think>`` encodes as ["</", "think>"] → [522, 26865].
|
||||
|
||||
Both special tokens ``<|THINK_START|>`` and ``<|THINK_END|>`` ARE
|
||||
present as single tokens.
|
||||
|
||||
This exercises the text-based path (with special-token fallback).
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._vocab = {
|
||||
TS: 100, # <think> - single token
|
||||
# </think> is intentionally NOT in the vocab
|
||||
"</": 522,
|
||||
"think>": 26865,
|
||||
TSS: 151669, # <|THINK_START|>
|
||||
TES: 151670, # <|THINK_END|>
|
||||
"Hello": 1,
|
||||
" world": 2,
|
||||
"reasoning": 3,
|
||||
" content": 4,
|
||||
"\n": 5,
|
||||
"some text": 8,
|
||||
}
|
||||
self._reverse = {v: k for k, v in self._vocab.items()}
|
||||
|
||||
def get_vocab(self):
|
||||
return dict(self._vocab)
|
||||
|
||||
def decode(self, ids, skip_special_tokens=False):
|
||||
return "".join(self._reverse.get(i, "") for i in ids)
|
||||
|
||||
def tokenize(self, text):
|
||||
if not text:
|
||||
return []
|
||||
return [text]
|
||||
|
||||
|
||||
class TextOnlyTokenizer:
|
||||
"""Tokenizer with only text-form markers (no special tokens at all).
|
||||
|
||||
Both ``<think>`` and ``</think>`` are single tokens.
|
||||
``<|THINK_START|>`` and ``<|THINK_END|>`` are NOT in the vocab.
|
||||
|
||||
This tests the pure text-based fallback paths.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
self._vocab = {
|
||||
TS: 100, # <think>
|
||||
TE: 101, # </think>
|
||||
"Hello": 1,
|
||||
" world": 2,
|
||||
"reasoning": 3,
|
||||
" content": 4,
|
||||
"\n": 5,
|
||||
"</": 6,
|
||||
"think>": 7,
|
||||
}
|
||||
self._reverse = {v: k for k, v in self._vocab.items()}
|
||||
|
||||
def get_vocab(self):
|
||||
return dict(self._vocab)
|
||||
|
||||
def decode(self, ids, skip_special_tokens=False):
|
||||
return "".join(self._reverse.get(i, "") for i in ids)
|
||||
|
||||
def tokenize(self, text):
|
||||
if not text:
|
||||
return []
|
||||
return [text]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _stream_extract(parser, deltas_and_ids):
|
||||
"""Run streaming extraction over a list of (delta_text, delta_ids) tuples.
|
||||
Returns (reasoning_text, content_text)."""
|
||||
prev_text = ""
|
||||
prev_ids = []
|
||||
all_r = []
|
||||
all_c = []
|
||||
|
||||
for delta_text, delta_ids in deltas_and_ids:
|
||||
cur_text = prev_text + delta_text
|
||||
cur_ids = prev_ids + delta_ids
|
||||
result = parser.extract_reasoning_streaming(prev_text, cur_text, delta_text, prev_ids, cur_ids, delta_ids)
|
||||
if result is not None:
|
||||
if result.reasoning:
|
||||
all_r.append(result.reasoning)
|
||||
if result.content:
|
||||
all_c.append(result.content)
|
||||
prev_text = cur_text
|
||||
prev_ids = cur_ids
|
||||
|
||||
return "".join(all_r), "".join(all_c)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixtures
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def single_token_parser():
|
||||
return StepAudioReasoningParser(SingleTokenTokenizer())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def multi_token_parser():
|
||||
return StepAudioReasoningParser(MultiTokenTokenizer())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def text_only_parser():
|
||||
return StepAudioReasoningParser(TextOnlyTokenizer())
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def request_obj():
|
||||
return MagicMock(messages=[], model="test-model")
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Initialization tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestInit:
|
||||
def test_single_token_ids(self, single_token_parser):
|
||||
"""Text-form markers are in vocab as single tokens."""
|
||||
assert single_token_parser.think_start_text_id == 100
|
||||
assert single_token_parser.think_end_text_id == 101
|
||||
|
||||
def test_single_token_special_ids(self, single_token_parser):
|
||||
"""Special-token markers are also in vocab."""
|
||||
assert single_token_parser.think_start_special_id == 151669
|
||||
assert single_token_parser.think_end_special_id == 151670
|
||||
|
||||
def test_single_token_has_ids(self, single_token_parser):
|
||||
assert single_token_parser._has_start_token_id is True
|
||||
assert single_token_parser._has_end_token_id is True
|
||||
|
||||
def test_multi_token_text_ids(self, multi_token_parser):
|
||||
"""<think> is a single token; </think> is NOT."""
|
||||
assert multi_token_parser.think_start_text_id == 100
|
||||
assert multi_token_parser.think_end_text_id == -1
|
||||
|
||||
def test_multi_token_special_ids(self, multi_token_parser):
|
||||
"""Special tokens are still in vocab."""
|
||||
assert multi_token_parser.think_start_special_id == 151669
|
||||
assert multi_token_parser.think_end_special_id == 151670
|
||||
|
||||
def test_multi_token_has_ids(self, multi_token_parser):
|
||||
"""Even without </think> as single token, special tokens provide IDs."""
|
||||
assert multi_token_parser._has_start_token_id is True
|
||||
assert multi_token_parser._has_end_token_id is True
|
||||
|
||||
def test_text_only_ids(self, text_only_parser):
|
||||
"""Only text-form markers, no special tokens."""
|
||||
assert text_only_parser.think_start_text_id == 100
|
||||
assert text_only_parser.think_end_text_id == 101
|
||||
assert text_only_parser.think_start_special_id == -1
|
||||
assert text_only_parser.think_end_special_id == -1
|
||||
assert text_only_parser._has_start_token_id is True
|
||||
assert text_only_parser._has_end_token_id is True
|
||||
|
||||
def test_missing_tokenizer_raises(self):
|
||||
with pytest.raises(ValueError, match="model tokenizer must be passed"):
|
||||
StepAudioReasoningParser(None)
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# extract_reasoning (non-streaming)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestExtractReasoning:
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
"text_only_parser",
|
||||
],
|
||||
)
|
||||
def test_basic_extraction(self, request, parser_name, request_obj):
|
||||
parser = request.getfixturevalue(parser_name)
|
||||
output = f"some reasoning{TE}Hello world"
|
||||
reasoning, content = parser.extract_reasoning(output, request_obj)
|
||||
assert reasoning == "some reasoning"
|
||||
assert content == "Hello world"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
"text_only_parser",
|
||||
],
|
||||
)
|
||||
def test_no_end_token(self, request, parser_name, request_obj):
|
||||
parser = request.getfixturevalue(parser_name)
|
||||
output = "still thinking..."
|
||||
reasoning, content = parser.extract_reasoning(output, request_obj)
|
||||
assert reasoning == "still thinking..."
|
||||
assert content is None
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
"text_only_parser",
|
||||
],
|
||||
)
|
||||
def test_with_start_token_prefix(self, request, parser_name, request_obj):
|
||||
parser = request.getfixturevalue(parser_name)
|
||||
output = f"{TS}my reasoning{TE}answer"
|
||||
reasoning, content = parser.extract_reasoning(output, request_obj)
|
||||
assert reasoning == "my reasoning"
|
||||
assert content == "answer"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
"text_only_parser",
|
||||
],
|
||||
)
|
||||
def test_strips_newline_after_end(self, request, parser_name, request_obj):
|
||||
parser = request.getfixturevalue(parser_name)
|
||||
output = f"reasoning{TE}\nHello"
|
||||
reasoning, content = parser.extract_reasoning(output, request_obj)
|
||||
assert reasoning == "reasoning"
|
||||
assert content == "Hello"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
"text_only_parser",
|
||||
],
|
||||
)
|
||||
def test_empty_reasoning(self, request, parser_name, request_obj):
|
||||
parser = request.getfixturevalue(parser_name)
|
||||
output = f"{TE}Hello"
|
||||
reasoning, content = parser.extract_reasoning(output, request_obj)
|
||||
assert reasoning is None
|
||||
assert content == "Hello"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
],
|
||||
)
|
||||
def test_special_end_token(self, request, parser_name, request_obj):
|
||||
"""Test with <|THINK_END|> special token form."""
|
||||
parser = request.getfixturevalue(parser_name)
|
||||
output = f"reasoning{TES}answer"
|
||||
reasoning, content = parser.extract_reasoning(output, request_obj)
|
||||
assert reasoning == "reasoning"
|
||||
assert content == "answer"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
],
|
||||
)
|
||||
def test_special_start_token(self, request, parser_name, request_obj):
|
||||
"""Test with <|THINK_START|> special token form."""
|
||||
parser = request.getfixturevalue(parser_name)
|
||||
output = f"{TSS}reasoning{TE}answer"
|
||||
reasoning, content = parser.extract_reasoning(output, request_obj)
|
||||
assert reasoning == "reasoning"
|
||||
assert content == "answer"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_reasoning_end
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIsReasoningEnd:
|
||||
def test_single_token_text_end(self, single_token_parser):
|
||||
assert single_token_parser.is_reasoning_end([1, 2, 101]) is True
|
||||
|
||||
def test_single_token_special_end(self, single_token_parser):
|
||||
assert single_token_parser.is_reasoning_end([1, 2, 151670]) is True
|
||||
|
||||
def test_single_token_no_end(self, single_token_parser):
|
||||
assert single_token_parser.is_reasoning_end([1, 2, 3]) is False
|
||||
|
||||
def test_multi_token_text_end_decoded(self, multi_token_parser):
|
||||
"""</think> = [522, 26865] detected via decode."""
|
||||
assert multi_token_parser.is_reasoning_end([1, 522, 26865]) is True
|
||||
|
||||
def test_multi_token_special_end(self, multi_token_parser):
|
||||
"""<|THINK_END|> detected as single token ID."""
|
||||
assert multi_token_parser.is_reasoning_end([1, 151670]) is True
|
||||
|
||||
def test_multi_token_no_end(self, multi_token_parser):
|
||||
assert multi_token_parser.is_reasoning_end([1, 2, 3]) is False
|
||||
|
||||
def test_multi_token_partial_marker(self, multi_token_parser):
|
||||
"""Only </ without think> should NOT count."""
|
||||
assert multi_token_parser.is_reasoning_end([1, 522, 3]) is False
|
||||
|
||||
def test_multi_token_empty_input(self, multi_token_parser):
|
||||
assert multi_token_parser.is_reasoning_end([]) is False
|
||||
|
||||
def test_text_only_end(self, text_only_parser):
|
||||
assert text_only_parser.is_reasoning_end([1, 2, 101]) is True
|
||||
|
||||
def test_text_only_no_end(self, text_only_parser):
|
||||
assert text_only_parser.is_reasoning_end([1, 2, 3]) is False
|
||||
|
||||
def test_start_after_end_returns_false(self, single_token_parser):
|
||||
"""Multi-turn prompt: start marker AFTER end marker → reasoning
|
||||
is still active (the generation prompt's start marker)."""
|
||||
# IDs: [start=100, end=101, start=100] → last marker is start
|
||||
assert single_token_parser.is_reasoning_end([100, 101, 100]) is False
|
||||
|
||||
def test_end_is_last_marker_returns_true(self, single_token_parser):
|
||||
"""End marker is the last marker → reasoning has ended."""
|
||||
# IDs: [start=100, end=101] → last marker is end
|
||||
assert single_token_parser.is_reasoning_end([100, 101]) is True
|
||||
|
||||
def test_start_only_no_end(self, single_token_parser):
|
||||
"""Only a start marker, no end marker → reasoning is active."""
|
||||
assert single_token_parser.is_reasoning_end([100, 1, 2]) is False
|
||||
|
||||
def test_multi_turn_prompt_start_after_end(self, multi_token_parser):
|
||||
"""Simulated multi-turn prompt with previous assistant response.
|
||||
|
||||
Prompt contains: ... </think ... <think ...
|
||||
The last marker is <think (start), so is_reasoning_end
|
||||
should return False.
|
||||
"""
|
||||
# Decode of these IDs: "reasoning" + "</ " + "think>" + " " + "<think"
|
||||
# The last marker is the start marker.
|
||||
# Using: [3=reasoning, 522="</", 26865="think>", 100=TS]
|
||||
ids = [3, 522, 26865, 100]
|
||||
assert multi_token_parser.is_reasoning_end(ids) is False
|
||||
|
||||
def test_multi_turn_prompt_end_is_last(self, multi_token_parser):
|
||||
"""Previous turn ends with </think, no new start → ended."""
|
||||
# [3=reasoning, 522="</", 26865="think>"]
|
||||
ids = [3, 522, 26865]
|
||||
assert multi_token_parser.is_reasoning_end(ids) is True
|
||||
|
||||
def test_special_start_after_end(self, multi_token_parser):
|
||||
"""Special token start after end marker → reasoning still active."""
|
||||
# [151670=<|THINK_END|>, 151669=<|THINK_START|>]
|
||||
ids = [151670, 151669]
|
||||
assert multi_token_parser.is_reasoning_end(ids) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# is_reasoning_end_streaming
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestIsReasoningEndStreaming:
|
||||
def test_single_token_delta_has_end(self, single_token_parser):
|
||||
assert single_token_parser.is_reasoning_end_streaming([1, 2], [101]) is True
|
||||
|
||||
def test_single_token_delta_no_end(self, single_token_parser):
|
||||
assert single_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False
|
||||
|
||||
def test_multi_token_delta_has_both_parts(self, multi_token_parser):
|
||||
"""Full </think> in delta → [522, 26865]."""
|
||||
assert multi_token_parser.is_reasoning_end_streaming([1], [522, 26865]) is True
|
||||
|
||||
def test_multi_token_marker_spans_boundary(self, multi_token_parser):
|
||||
"""</ in previous, think> in delta."""
|
||||
assert multi_token_parser.is_reasoning_end_streaming([1, 522], [26865]) is True
|
||||
|
||||
def test_multi_token_no_end(self, multi_token_parser):
|
||||
assert multi_token_parser.is_reasoning_end_streaming([1, 2], [3]) is False
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# extract_reasoning_streaming
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestExtractReasoningStreaming:
|
||||
"""Test streaming extraction for both single-token and multi-token modes."""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
"text_only_parser",
|
||||
],
|
||||
)
|
||||
def test_simple_streaming(self, request_, parser_name):
|
||||
parser = request_.getfixturevalue(parser_name)
|
||||
# Simulate: <think> reasoning </think> final
|
||||
if parser_name == "multi_token_parser":
|
||||
deltas = [TS, "reasoning", TE, " content"]
|
||||
token_ids_per_delta = [
|
||||
[100], # <think>
|
||||
[3], # reasoning
|
||||
[522, 26865], # </think>
|
||||
[4], # content
|
||||
]
|
||||
else:
|
||||
deltas = [TS, "reasoning", TE, " content"]
|
||||
token_ids_per_delta = [
|
||||
[100], # <think>
|
||||
[3], # reasoning
|
||||
[101], # </think>
|
||||
[4], # content
|
||||
]
|
||||
|
||||
r, c = _stream_extract(parser, list(zip(deltas, token_ids_per_delta)))
|
||||
assert r == "reasoning"
|
||||
assert c == " content"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
"text_only_parser",
|
||||
],
|
||||
)
|
||||
def test_no_end_token_streaming(self, request_, parser_name):
|
||||
parser = request_.getfixturevalue(parser_name)
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
("thinking", [3]),
|
||||
(" more", [2]),
|
||||
]
|
||||
r, c = _stream_extract(parser, deltas_and_ids)
|
||||
assert r == "thinking more"
|
||||
assert c == ""
|
||||
|
||||
def test_multi_token_end_split_across_chunks(self, multi_token_parser):
|
||||
"""The key test: </think> is split across streaming chunks.
|
||||
|
||||
Previous chunk ends with "reasoning</" and delta starts with
|
||||
"think>\\nfinal answer". This is the real-world scenario when
|
||||
``</think>`` is tokenized as two separate tokens.
|
||||
"""
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
("</", [522]),
|
||||
("think>\nHello", [26865, 1]),
|
||||
]
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "Hello"
|
||||
|
||||
def test_multi_token_start_token_in_delta(self, multi_token_parser):
|
||||
"""Test that <think> start token is skipped in streaming output."""
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
("reasoning", [3]),
|
||||
(TE, [522, 26865]),
|
||||
("Hello", [1]),
|
||||
]
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "Hello"
|
||||
|
||||
def test_multi_token_three_way_split(self, multi_token_parser):
|
||||
"""</think> split across 3 chunks: </ + thi + nk>."""
|
||||
deltas_and_ids = [
|
||||
("some reasoning", [3]),
|
||||
("</", [522]),
|
||||
("thi", [8]),
|
||||
("nk>", [8]),
|
||||
("\nHello", [1]),
|
||||
]
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == "some reasoning"
|
||||
assert c == "Hello"
|
||||
|
||||
def test_special_end_token_in_stream(self, multi_token_parser):
|
||||
"""<|THINK_END|> special token in streaming (decoded as text)."""
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
(TES, [151670]),
|
||||
("answer", [1]),
|
||||
]
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "answer"
|
||||
|
||||
def test_no_start_token_in_stream(self, multi_token_parser):
|
||||
"""Model starts directly with reasoning (no <think> prefix)."""
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
(TE, [522, 26865]),
|
||||
("answer", [1]),
|
||||
]
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "answer"
|
||||
|
||||
def test_empty_reasoning(self, single_token_parser):
|
||||
"""Start immediately followed by end → empty reasoning."""
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
(TE, [101]),
|
||||
("Hello", [1]),
|
||||
]
|
||||
r, c = _stream_extract(single_token_parser, deltas_and_ids)
|
||||
assert r in ("", None)
|
||||
assert c == "Hello"
|
||||
|
||||
def test_end_marker_no_content_after(self, multi_token_parser):
|
||||
"""End marker with no content following it."""
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
(TE, [522, 26865]),
|
||||
]
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == ""
|
||||
|
||||
def test_delta_is_just_start_token(self, multi_token_parser):
|
||||
"""Delta consisting only of the start token should be skipped."""
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
]
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == ""
|
||||
assert c == ""
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
],
|
||||
)
|
||||
def test_only_first_newline_stripped_after_end(self, request_, parser_name):
|
||||
"""Only the first newline after </think> should be stripped."""
|
||||
parser = request_.getfixturevalue(parser_name)
|
||||
if parser_name == "multi_token_parser":
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
(TE, [522, 26865]),
|
||||
("\n\nHello", [5, 1]),
|
||||
]
|
||||
else:
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
(TE, [101]),
|
||||
("\n\nHello", [5, 1]),
|
||||
]
|
||||
r, c = _stream_extract(parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "\nHello"
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"parser_name",
|
||||
[
|
||||
"single_token_parser",
|
||||
"multi_token_parser",
|
||||
],
|
||||
)
|
||||
def test_content_after_end_multiple_deltas(self, request_, parser_name):
|
||||
"""Multiple content deltas after end marker."""
|
||||
parser = request_.getfixturevalue(parser_name)
|
||||
if parser_name == "multi_token_parser":
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
("reasoning", [3]),
|
||||
(TE, [522, 26865]),
|
||||
("Hello", [1]),
|
||||
(" world", [2]),
|
||||
]
|
||||
else:
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
("reasoning", [3]),
|
||||
(TE, [101]),
|
||||
("Hello", [1]),
|
||||
(" world", [2]),
|
||||
]
|
||||
r, c = _stream_extract(parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "Hello world"
|
||||
|
||||
def test_pending_buffer_resolved_on_non_marker(self, multi_token_parser):
|
||||
"""When buffered text turns out not to be a marker prefix, it is
|
||||
emitted as reasoning on the next call."""
|
||||
# "reason<" is buffered because "<" could start a marker.
|
||||
# Next delta ">" doesn't form a marker, so the buffer is
|
||||
# flushed as reasoning.
|
||||
parser = multi_token_parser
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
("reason", [3]),
|
||||
("<", [8]), # "<" is buffered (prefix of markers)
|
||||
("> more", [8]), # ">" doesn't complete any marker
|
||||
]
|
||||
r, c = _stream_extract(parser, deltas_and_ids)
|
||||
assert r == "reason> more"
|
||||
assert c == ""
|
||||
|
||||
def test_multi_token_end_with_reasoning_before_buffer(self, multi_token_parser):
|
||||
"""Reasoning text before a buffered prefix should be emitted."""
|
||||
# "reasoning<" → "reasoning" emitted, "<" buffered
|
||||
# "</think>\nanswer" → buffer resolves as end marker
|
||||
parser = multi_token_parser
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
("<", [8]),
|
||||
(f"/{TE[2:]}\nanswer", [522, 26865, 1]),
|
||||
]
|
||||
r, c = _stream_extract(parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "answer"
|
||||
|
||||
def test_multi_token_five_way_split(self, multi_token_parser):
|
||||
"""</think> split across 5 tiny chunks."""
|
||||
# Split: each chunk is one character of </think>
|
||||
chunks = list(TE)
|
||||
deltas_and_ids = [
|
||||
("reasoning", [3]),
|
||||
]
|
||||
for ch in chunks:
|
||||
deltas_and_ids.append((ch, [8]))
|
||||
deltas_and_ids.append(("\nHello", [1]))
|
||||
r, c = _stream_extract(multi_token_parser, deltas_and_ids)
|
||||
assert r == "reasoning"
|
||||
assert c == "Hello"
|
||||
|
||||
def test_single_token_parser_with_pending(self, single_token_parser):
|
||||
"""Single-token parser: ambiguous '<' prefix buffered then resolved."""
|
||||
# Even with single-token markers, the text-based path buffers
|
||||
# '<' because it could be a prefix of a marker.
|
||||
parser = single_token_parser
|
||||
deltas_and_ids = [
|
||||
(TS, [100]),
|
||||
("reason", [3]),
|
||||
("<", [8]),
|
||||
("> not a marker", [8]),
|
||||
]
|
||||
r, c = _stream_extract(parser, deltas_and_ids)
|
||||
assert r == "reason> not a marker"
|
||||
assert c == ""
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestCountReasoningTokens:
|
||||
def test_single_token_basic(self, single_token_parser):
|
||||
# IDs: [100=<think>, 3=reasoning, 101=</think>, 4=content]
|
||||
assert single_token_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1
|
||||
|
||||
def test_single_token_no_end(self, single_token_parser):
|
||||
# All tokens after start are reasoning
|
||||
assert single_token_parser.count_reasoning_tokens([100, 3, 2]) == 2
|
||||
|
||||
def test_multi_token_basic(self, multi_token_parser):
|
||||
# </think> encodes as [522, 26865]
|
||||
# Full sequence: <think> reasoning </think> content
|
||||
# IDs: [100, 3, 522, 26865, 4]
|
||||
result = multi_token_parser.count_reasoning_tokens([100, 3, 522, 26865, 4])
|
||||
assert result > 0
|
||||
|
||||
def test_multi_token_no_end(self, multi_token_parser):
|
||||
# No end token → all tokens inside <think> are counted.
|
||||
result = multi_token_parser.count_reasoning_tokens([100, 3, 2])
|
||||
assert result > 0
|
||||
|
||||
def test_text_only_basic(self, text_only_parser):
|
||||
assert text_only_parser.count_reasoning_tokens([100, 3, 101, 4]) == 1
|
||||
|
||||
def test_text_only_no_end(self, text_only_parser):
|
||||
assert text_only_parser.count_reasoning_tokens([100, 3, 2]) == 2
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# extract_content_ids
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestExtractContentIds:
|
||||
def test_single_token_basic(self, single_token_parser):
|
||||
# [1, 101=</think>, 2, 3] → content starts after ID 101
|
||||
result = single_token_parser.extract_content_ids([1, 101, 2, 3])
|
||||
assert result == [2, 3]
|
||||
|
||||
def test_single_token_no_end(self, single_token_parser):
|
||||
result = single_token_parser.extract_content_ids([1, 2, 3])
|
||||
assert result == []
|
||||
|
||||
def test_multi_token_basic(self, multi_token_parser):
|
||||
# </think> = [522, 26865]
|
||||
# [3, 522, 26865, 4] → content after </think>
|
||||
result = multi_token_parser.extract_content_ids([3, 522, 26865, 4])
|
||||
assert 4 in result
|
||||
|
||||
def test_multi_token_no_end(self, multi_token_parser):
|
||||
result = multi_token_parser.extract_content_ids([1, 2, 3])
|
||||
assert result == []
|
||||
|
||||
def test_text_only_basic(self, text_only_parser):
|
||||
result = text_only_parser.extract_content_ids([1, 101, 2, 3])
|
||||
assert result == [2, 3]
|
||||
|
||||
def test_text_only_no_end(self, text_only_parser):
|
||||
result = text_only_parser.extract_content_ids([1, 2, 3])
|
||||
assert result == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Fixture request workaround (pytest "request" fixture shadowed by parametrize)
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def request_(request):
|
||||
"""Proxy for the pytest ``request`` fixture to avoid name collisions
|
||||
with the ``request_obj`` fixture used in extract_reasoning tests."""
|
||||
return request
|
||||
@@ -1,182 +0,0 @@
|
||||
"""Functional tests for pydub removal.
|
||||
|
||||
Verifies that the new av/numpy/torchaudio replacements produce
|
||||
identical results to the old pydub-based code for all three paths:
|
||||
1. WAV PCM → duration & frame count (pure arithmetic)
|
||||
2. Non-WAV (mp3/ogg) → duration & frame count (av.open)
|
||||
3. TTS WER PCM export: resample to 24000Hz mono s16
|
||||
"""
|
||||
|
||||
import io
|
||||
import wave
|
||||
|
||||
import av
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
|
||||
def _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=1.0):
|
||||
"""Generate a WAV file with a sine wave, return (wav_bytes, pcm_bytes, params)."""
|
||||
n_samples = int(sample_rate * duration_sec)
|
||||
t = np.linspace(0, duration_sec, n_samples, endpoint=False)
|
||||
samples = (np.sin(2 * np.pi * 440 * t) * 32767 * 0.5).astype(np.int16)
|
||||
if channels == 2:
|
||||
samples = np.column_stack([samples, samples]).flatten()
|
||||
|
||||
buf = io.BytesIO()
|
||||
with wave.open(buf, "wb") as wf:
|
||||
wf.setnchannels(channels)
|
||||
wf.setsampwidth(sample_width)
|
||||
wf.setframerate(sample_rate)
|
||||
wf.writeframes(samples.tobytes())
|
||||
wav_bytes = buf.getvalue()
|
||||
|
||||
# Extract raw PCM (skip WAV header)
|
||||
with wave.open(io.BytesIO(wav_bytes), "rb") as wr:
|
||||
pcm_bytes = wr.readframes(wr.getnframes())
|
||||
params = (wr.getnchannels(), wr.getsampwidth(), wr.getframerate())
|
||||
|
||||
return wav_bytes, pcm_bytes, params
|
||||
|
||||
|
||||
def _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0):
|
||||
"""Generate an OGG Vorbis file with a sine wave, return bytes."""
|
||||
n_samples = int(sample_rate * duration_sec)
|
||||
t = np.linspace(0, duration_sec, n_samples, endpoint=False)
|
||||
samples = (np.sin(2 * np.pi * 440 * t) * 0.5).astype(np.float32)
|
||||
buf = io.BytesIO()
|
||||
sf.write(buf, samples, sample_rate, format="OGG", subtype="VORBIS")
|
||||
return buf.getvalue()
|
||||
|
||||
|
||||
def test_wav_path_duration_and_frames():
|
||||
"""Path 1: WAV PCM buffer → duration & frames via arithmetic."""
|
||||
for sr in [16000, 24000, 44100]:
|
||||
for ch in [1, 2]:
|
||||
for dur in [0.5, 1.0, 2.5]:
|
||||
_, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes(
|
||||
sample_rate=sr, channels=ch, duration_sec=dur
|
||||
)
|
||||
|
||||
# New code logic (from patch.py)
|
||||
frame_width = sample_width * channels
|
||||
audio_frames = len(pcm_bytes) // frame_width
|
||||
audio_duration_sec = audio_frames / frame_rate
|
||||
|
||||
assert abs(audio_duration_sec - dur) < 0.001, f"sr={sr} ch={ch} dur={dur}: got {audio_duration_sec}"
|
||||
expected_frames = int(sr * dur)
|
||||
assert audio_frames == expected_frames, (
|
||||
f"sr={sr} ch={ch} dur={dur}: frames {audio_frames} != {expected_frames}"
|
||||
)
|
||||
|
||||
print("[PASS] test_wav_path_duration_and_frames")
|
||||
|
||||
|
||||
def test_non_wav_path_duration_av():
|
||||
"""Path 2: non-WAV audio → duration & frames via av."""
|
||||
for dur in [0.5, 1.0, 2.0]:
|
||||
ogg_bytes = _make_ogg_bytes(sample_rate=16000, duration_sec=dur)
|
||||
|
||||
container = av.open(io.BytesIO(ogg_bytes))
|
||||
stream = container.streams.audio[0]
|
||||
audio_frames = sum(f.samples for f in container.decode(stream))
|
||||
audio_duration_sec = audio_frames / stream.rate if stream.rate else 0.0
|
||||
container.close()
|
||||
|
||||
# OGG encoding may add padding, allow small tolerance
|
||||
assert abs(audio_duration_sec - dur) < 0.05, f"dur={dur}: got {audio_duration_sec}"
|
||||
|
||||
print("[PASS] test_non_wav_path_duration_av")
|
||||
|
||||
|
||||
def test_tts_pcm_export_wav_path():
|
||||
"""Path 3a: WAV PCM → resample to 24000Hz mono s16."""
|
||||
import torch
|
||||
import torchaudio
|
||||
|
||||
_, pcm_bytes, (channels, sample_width, frame_rate) = _make_wav_bytes(
|
||||
sample_rate=16000, channels=2, sample_width=2, duration_sec=1.0
|
||||
)
|
||||
|
||||
# New code logic
|
||||
pcm_array = np.frombuffer(pcm_bytes, dtype=np.int16)
|
||||
if channels > 1:
|
||||
pcm_array = pcm_array.reshape(-1, channels).mean(axis=1).astype(np.int16)
|
||||
if frame_rate != 24000:
|
||||
t = torch.from_numpy(pcm_array.astype(np.float32)).unsqueeze(0)
|
||||
t = torchaudio.functional.resample(t, frame_rate, 24000)
|
||||
pcm_array = t.squeeze(0).to(torch.int16).numpy()
|
||||
|
||||
result_bytes = pcm_array.tobytes()
|
||||
# Should be mono, 24000Hz, 2 bytes per sample
|
||||
expected_samples = int(24000 * 1.0)
|
||||
actual_samples = len(result_bytes) // 2
|
||||
assert abs(actual_samples - expected_samples) < 10, f"expected ~{expected_samples} samples, got {actual_samples}"
|
||||
assert len(pcm_array.shape) == 1, "should be mono (1D array)"
|
||||
|
||||
print("[PASS] test_tts_pcm_export_wav_path")
|
||||
|
||||
|
||||
def test_tts_pcm_export_non_wav_path():
|
||||
"""Path 3b: non-WAV → resample to 24000Hz mono s16 via av.AudioResampler."""
|
||||
ogg_bytes = _make_ogg_bytes(sample_rate=16000, channels=1, duration_sec=1.0)
|
||||
|
||||
container = av.open(io.BytesIO(ogg_bytes))
|
||||
resampler = av.AudioResampler(format="s16", layout="mono", rate=24000)
|
||||
frames = []
|
||||
for frame in container.decode(audio=0):
|
||||
for rf in resampler.resample(frame):
|
||||
frames.append(rf.to_ndarray().flatten())
|
||||
container.close()
|
||||
|
||||
result = np.concatenate(frames)
|
||||
result_bytes = result.tobytes()
|
||||
expected_samples = int(24000 * 1.0)
|
||||
actual_samples = len(result_bytes) // 2
|
||||
assert abs(actual_samples - expected_samples) < 500, f"expected ~{expected_samples} samples, got {actual_samples}"
|
||||
|
||||
print("[PASS] test_tts_pcm_export_non_wav_path")
|
||||
|
||||
|
||||
def test_media_merge_and_export():
|
||||
"""media.py: _merge_base64_audio_to_segment decode + concat + export WAV."""
|
||||
import base64 as b64mod
|
||||
import sys
|
||||
|
||||
sys.path.insert(0, ".")
|
||||
from tests.helpers.media import _merge_base64_audio_to_segment
|
||||
|
||||
# Create two WAV chunks as base64
|
||||
chunks_b64 = []
|
||||
for dur in [0.5, 0.3]:
|
||||
wav_bytes, _, _ = _make_wav_bytes(sample_rate=16000, channels=1, sample_width=2, duration_sec=dur)
|
||||
chunks_b64.append(b64mod.b64encode(wav_bytes).decode())
|
||||
|
||||
merged = _merge_base64_audio_to_segment(chunks_b64)
|
||||
|
||||
# Test export to WAV
|
||||
buf = io.BytesIO()
|
||||
merged.export(buf, format="wav")
|
||||
wav_out = buf.getvalue()
|
||||
assert len(wav_out) > 44, f"WAV output too small: {len(wav_out)} bytes"
|
||||
|
||||
# Verify merged duration (~0.8s)
|
||||
out_data, out_sr = sf.read(io.BytesIO(wav_out))
|
||||
actual_dur = len(out_data) / out_sr
|
||||
assert abs(actual_dur - 0.8) < 0.05, f"expected ~0.8s merged duration, got {actual_dur}"
|
||||
|
||||
print(f"[PASS] test_media_merge_and_export (duration={actual_dur:.3f}s, wav_size={len(wav_out)} bytes)")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_wav_path_duration_and_frames()
|
||||
test_non_wav_path_duration_av()
|
||||
import importlib.util
|
||||
|
||||
if importlib.util.find_spec("torch") and importlib.util.find_spec("torchaudio"):
|
||||
test_tts_pcm_export_wav_path()
|
||||
else:
|
||||
print("[SKIP] test_tts_pcm_export_wav_path (torch/torchaudio not available)")
|
||||
test_tts_pcm_export_non_wav_path()
|
||||
test_media_merge_and_export()
|
||||
print("\nAll tests passed.")
|
||||
@@ -6,10 +6,6 @@ AC-5 (delay pattern), AC-7 (stage processor), and AC-10 (registry)
|
||||
without requiring the actual checkpoint or GPU.
|
||||
"""
|
||||
|
||||
import asyncio
|
||||
import time
|
||||
from collections import OrderedDict
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
@@ -236,7 +232,7 @@ class TestDelayPatternBehavior:
|
||||
|
||||
|
||||
class TestSamplerMethods:
|
||||
"""Test the actual sampler and batched delay masking methods."""
|
||||
"""Test the actual _sample_audio_codes and _apply_delay_pattern_masking methods."""
|
||||
|
||||
def _make_minimal_talker(self):
|
||||
"""Create a minimal talker-like object with sampler/masking methods."""
|
||||
@@ -245,65 +241,18 @@ class TestSamplerMethods:
|
||||
class FakeTalker:
|
||||
num_codebooks = 8
|
||||
codebook_size = 1026
|
||||
_audio_state = {}
|
||||
|
||||
t = FakeTalker()
|
||||
t._sample_audio_codes = mod.HiggsAudioV3TalkerForConditionalGeneration._sample_audio_codes.__get__(t)
|
||||
return t
|
||||
|
||||
def _make_batched_sampler_talker(self, num_rows=4):
|
||||
"""Create a fake talker with GPU-resident sampler state helpers."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
|
||||
|
||||
class FakeTalker:
|
||||
num_codebooks = 8
|
||||
codebook_size = 1026
|
||||
|
||||
t = FakeTalker()
|
||||
t._decode_last_codes = torch.arange(num_rows * t.num_codebooks, dtype=torch.long).view(
|
||||
num_rows, t.num_codebooks
|
||||
t._apply_delay_pattern_masking = (
|
||||
mod.HiggsAudioV3TalkerForConditionalGeneration._apply_delay_pattern_masking.__get__(t)
|
||||
)
|
||||
t._decode_has_codes = torch.tensor([True, False, True, True])[:num_rows].clone()
|
||||
t._decode_delay_count = torch.tensor([0, 2, 7, 4], dtype=torch.int32)[:num_rows].clone()
|
||||
t._decode_eoc_countdown = torch.tensor([-1, -1, 3, 1], dtype=torch.int32)[:num_rows].clone()
|
||||
t._decode_generation_done = torch.tensor([False, False, False, True])[:num_rows].clone()
|
||||
t._decode_active_audio_count = 0
|
||||
t._codebook_index_cache = {}
|
||||
t._row_index_cache = {}
|
||||
t._boc_frame_cache = {}
|
||||
t._last_audio_codes_buffer = None
|
||||
t._last_audio_host_staging = None
|
||||
t._last_audio_gpu_staging = None
|
||||
t._last_audio_staging_event = None
|
||||
t._audio_staging_events = []
|
||||
t._audio_staging_event_cursor = 0
|
||||
t._last_audio_codes = None
|
||||
t._last_audio_code_valid = []
|
||||
t._postprocess_cursor = 0
|
||||
t._postprocess_audio_rows = 0
|
||||
t._postprocess_audio_active_rows = 0
|
||||
|
||||
cls = mod.HiggsAudioV3TalkerForConditionalGeneration
|
||||
for name in (
|
||||
"_sample_audio_codes",
|
||||
"_ensure_decode_state_capacity",
|
||||
"_get_row_indices",
|
||||
"_get_codebook_indices",
|
||||
"_get_boc_frames",
|
||||
"_get_audio_codes_buffer",
|
||||
"_get_audio_gpu_staging_buffer",
|
||||
"_get_audio_host_staging_buffer",
|
||||
"_apply_delay_pattern_masking_batched",
|
||||
"_update_delay_state_batched",
|
||||
"_prefill_row_mask",
|
||||
"_audio_seed_mask_from_step_input",
|
||||
):
|
||||
setattr(t, name, getattr(cls, name).__get__(t))
|
||||
t._device_cache_key = cls._device_cache_key
|
||||
return t
|
||||
|
||||
def test_sample_respects_mask(self):
|
||||
"""Tokens masked to -inf must never be sampled."""
|
||||
t = self._make_batched_sampler_talker(num_rows=1)
|
||||
t = self._make_minimal_talker()
|
||||
logits = torch.full((10, 1026), float("-inf"))
|
||||
# Only token 500 is allowed for each row
|
||||
logits[:, 500] = 0.0
|
||||
@@ -313,7 +262,7 @@ class TestSamplerMethods:
|
||||
|
||||
def test_sample_all_masked_falls_back_to_argmax(self):
|
||||
"""All-masked row should fall back to argmax (least-negative logit)."""
|
||||
t = self._make_batched_sampler_talker(num_rows=1)
|
||||
t = self._make_minimal_talker()
|
||||
logits = torch.full((2, 1026), float("-inf"))
|
||||
# Row 0: all masked
|
||||
# Row 1: only token 42 allowed
|
||||
@@ -326,13 +275,11 @@ class TestSamplerMethods:
|
||||
"""During delay phase, codebooks beyond delay_count must have only BOC allowed."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
|
||||
|
||||
t = self._make_batched_sampler_talker(num_rows=1)
|
||||
t = self._make_minimal_talker()
|
||||
# Simulate delay_count=2 for batch row 0
|
||||
t._ensure_decode_state_capacity(1, torch.device("cpu"))
|
||||
t._decode_delay_count[0] = 2
|
||||
t._decode_eoc_countdown[0] = -1
|
||||
t._audio_state[0] = {"num_delay": 2, "num_remaining_delays": None}
|
||||
cb_logits = torch.zeros(1, 8, 1026) # [1 audio row, 8 codebooks, 1026 vocab]
|
||||
t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
|
||||
t._apply_delay_pattern_masking(cb_logits, [0])
|
||||
# CBs 3-7 should have only BOC allowed (everything else -inf)
|
||||
for q in range(3, 8):
|
||||
row = cb_logits[0, q]
|
||||
@@ -345,12 +292,10 @@ class TestSamplerMethods:
|
||||
"""Active codebooks during delay should have BOC disallowed."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import BOC_ID
|
||||
|
||||
t = self._make_batched_sampler_talker(num_rows=1)
|
||||
t._ensure_decode_state_capacity(1, torch.device("cpu"))
|
||||
t._decode_delay_count[0] = 3
|
||||
t._decode_eoc_countdown[0] = -1
|
||||
t = self._make_minimal_talker()
|
||||
t._audio_state[0] = {"num_delay": 3, "num_remaining_delays": None}
|
||||
cb_logits = torch.zeros(1, 8, 1026)
|
||||
t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
|
||||
t._apply_delay_pattern_masking(cb_logits, [0])
|
||||
# CBs 0-3 are active: BOC should be -inf
|
||||
for q in range(4):
|
||||
assert cb_logits[0, q, BOC_ID].item() == float("-inf")
|
||||
@@ -359,12 +304,10 @@ class TestSamplerMethods:
|
||||
"""Only codebook 0 should allow EOC during normal generation."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
|
||||
|
||||
t = self._make_batched_sampler_talker(num_rows=1)
|
||||
t._ensure_decode_state_capacity(1, torch.device("cpu"))
|
||||
t._decode_delay_count[0] = 8
|
||||
t._decode_eoc_countdown[0] = -1
|
||||
t = self._make_minimal_talker()
|
||||
t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": None}
|
||||
cb_logits = torch.zeros(1, 8, 1026)
|
||||
t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
|
||||
t._apply_delay_pattern_masking(cb_logits, [0])
|
||||
# CB0 should keep EOC
|
||||
assert cb_logits[0, 0, EOC_ID].item() != float("-inf")
|
||||
# CB1-7 should have EOC masked
|
||||
@@ -378,13 +321,11 @@ class TestSamplerMethods:
|
||||
EOC_ID,
|
||||
)
|
||||
|
||||
t = self._make_batched_sampler_talker(num_rows=1)
|
||||
t = self._make_minimal_talker()
|
||||
# Ramp-down with 4 remaining delays: lock CBs 0-3 to EOC
|
||||
t._ensure_decode_state_capacity(1, torch.device("cpu"))
|
||||
t._decode_delay_count[0] = 8
|
||||
t._decode_eoc_countdown[0] = 4
|
||||
t._audio_state[0] = {"num_delay": 8, "num_remaining_delays": 4}
|
||||
cb_logits = torch.zeros(1, 8, 1026)
|
||||
t._apply_delay_pattern_masking_batched(cb_logits, torch.tensor([0], dtype=torch.long))
|
||||
t._apply_delay_pattern_masking(cb_logits, [0])
|
||||
# CBs 0-3 locked: only EOC allowed
|
||||
for q in range(4):
|
||||
row = cb_logits[0, q]
|
||||
@@ -396,195 +337,6 @@ class TestSamplerMethods:
|
||||
assert cb_logits[0, q, BOC_ID].item() == float("-inf")
|
||||
assert cb_logits[0, q, EOC_ID].item() == float("-inf")
|
||||
|
||||
def test_batched_delay_masking_all_rows_matches_sparse_path(self):
|
||||
"""The all-row fast path must preserve sparse path masking semantics."""
|
||||
torch.manual_seed(0)
|
||||
sparse = self._make_batched_sampler_talker()
|
||||
all_rows = self._make_batched_sampler_talker()
|
||||
rows = torch.arange(4, dtype=torch.long)
|
||||
logits_sparse = torch.randn(4, 8, 1026)
|
||||
logits_all_rows = logits_sparse.clone()
|
||||
|
||||
sparse._apply_delay_pattern_masking_batched(logits_sparse, rows, all_rows=False)
|
||||
all_rows._apply_delay_pattern_masking_batched(logits_all_rows, rows, all_rows=True)
|
||||
|
||||
assert torch.equal(logits_all_rows, logits_sparse)
|
||||
|
||||
def test_delay_state_update_all_rows_matches_sparse_path(self):
|
||||
"""The all-row fast path must update state and staging like index_copy_."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
|
||||
EOC_ID,
|
||||
)
|
||||
|
||||
sparse = self._make_batched_sampler_talker()
|
||||
all_rows = self._make_batched_sampler_talker()
|
||||
rows = torch.arange(4, dtype=torch.long)
|
||||
codes = torch.tensor(
|
||||
[
|
||||
[10, 11, 12, 13, 14, 15, 16, 17],
|
||||
[20, 21, 22, 23, 24, 25, 26, 27],
|
||||
[30, 31, EOC_ID, 33, 34, 35, 36, 37],
|
||||
[40, 41, 42, 43, 44, 45, 46, 47],
|
||||
],
|
||||
dtype=torch.long,
|
||||
)
|
||||
code_row_mask = torch.tensor([True, False, True, True])
|
||||
|
||||
sparse._update_delay_state_batched(
|
||||
codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=False
|
||||
)
|
||||
all_rows._update_delay_state_batched(
|
||||
codes.clone(), rows, 4, torch.device("cpu"), code_row_mask=code_row_mask, all_rows=True
|
||||
)
|
||||
|
||||
assert torch.equal(all_rows._decode_delay_count, sparse._decode_delay_count)
|
||||
assert torch.equal(all_rows._decode_eoc_countdown, sparse._decode_eoc_countdown)
|
||||
assert torch.equal(all_rows._decode_generation_done, sparse._decode_generation_done)
|
||||
assert torch.equal(all_rows._decode_last_codes, sparse._decode_last_codes)
|
||||
assert torch.equal(all_rows._decode_has_codes, sparse._decode_has_codes)
|
||||
assert all_rows._last_audio_codes is None
|
||||
assert sparse._last_audio_codes is None
|
||||
assert torch.equal(all_rows._last_audio_host_staging, sparse._last_audio_host_staging)
|
||||
|
||||
def test_direct_audio_sampler_does_not_override_done_tokens(self):
|
||||
"""Done rows are represented directly in GPU sampled tokens."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
|
||||
|
||||
t = self._make_batched_sampler_talker(num_rows=2)
|
||||
t._resolve_token_ids = lambda: None
|
||||
t._audio_continuation_id = 99999
|
||||
t._eos_token_id = 151671
|
||||
t._last_logits_hidden = torch.zeros(2, 16)
|
||||
t._last_step_input_ids = torch.tensor([99999, 99999])
|
||||
t._last_step_query_start_loc = None
|
||||
t._decode_has_codes = torch.tensor([True, True])
|
||||
t._decode_generation_done = torch.tensor([False, True])
|
||||
t._decode_delay_count = torch.zeros(2, dtype=torch.long)
|
||||
t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long)
|
||||
t._fast_audio_direct_rows = 2
|
||||
t._last_audio_done_flags = None
|
||||
t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: None
|
||||
t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026)
|
||||
t._apply_delay_pattern_masking_batched = lambda cb_logits, audio_rows, all_rows=False: None
|
||||
t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long)
|
||||
t._update_delay_state_batched = lambda *args, **kwargs: None
|
||||
t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t)
|
||||
|
||||
sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object())
|
||||
|
||||
assert not getattr(
|
||||
mod.HiggsAudioV3TalkerForConditionalGeneration, "supports_sampled_token_ids_cpu_override", False
|
||||
)
|
||||
assert sampler_output.sampled_token_ids.tolist() == [[99999], [151671]]
|
||||
|
||||
def test_direct_audio_sampler_uses_gpu_tokens_for_active_rows(self):
|
||||
"""Direct audio sampling leaves scheduler-visible tokens to the GPU tensor."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
|
||||
|
||||
t = self._make_batched_sampler_talker(num_rows=2)
|
||||
t._resolve_token_ids = lambda: None
|
||||
t._audio_continuation_id = 99999
|
||||
t._eos_token_id = 151671
|
||||
t._last_logits_hidden = torch.zeros(2, 16)
|
||||
t._last_step_input_ids = torch.tensor([99999, 99999])
|
||||
t._last_step_query_start_loc = None
|
||||
t._decode_has_codes = torch.tensor([True, True])
|
||||
t._decode_generation_done = torch.tensor([False, False])
|
||||
t._decode_delay_count = torch.zeros(2, dtype=torch.long)
|
||||
t._decode_eoc_countdown = torch.full((2,), -1, dtype=torch.long)
|
||||
t._fast_audio_direct_rows = 2
|
||||
t._last_audio_done_flags = [0, 0]
|
||||
t._fast_audio_sampler_gpu_fallback_reason = lambda **kwargs: None
|
||||
t._audio_codebook_logits_from_rows = lambda hidden, rows, all_rows=False: torch.zeros(2, 8, 1026)
|
||||
t._apply_delay_pattern_masking_batched = lambda cb_logits, audio_rows, all_rows=False: None
|
||||
t._sample_audio_codes = lambda logits_2d: torch.zeros(int(logits_2d.shape[0]), dtype=torch.long)
|
||||
t._update_delay_state_batched = lambda *args, **kwargs: None
|
||||
t.sample = mod.HiggsAudioV3TalkerForConditionalGeneration.sample.__get__(t)
|
||||
|
||||
sampler_output = t.sample(torch.zeros(2, 200000), sampling_metadata=object())
|
||||
|
||||
assert sampler_output.sampled_token_ids.tolist() == [[99999], [99999]]
|
||||
|
||||
def test_terminal_rampdown_frame_is_emitted_before_done(self):
|
||||
"""The final ramp-down frame must be emitted, not replaced by a done marker."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
|
||||
|
||||
t = self._make_batched_sampler_talker(num_rows=1)
|
||||
t._decode_delay_count[0] = 8
|
||||
t._decode_eoc_countdown[0] = 1
|
||||
t._decode_generation_done[0] = False
|
||||
t._decode_has_codes[0] = True
|
||||
codes = torch.tensor([[10, 11, 12, 13, 14, 15, 16, 17]], dtype=torch.long)
|
||||
|
||||
t._update_delay_state_batched(
|
||||
codes,
|
||||
torch.tensor([0], dtype=torch.long),
|
||||
1,
|
||||
torch.device("cpu"),
|
||||
code_row_mask=torch.tensor([True]),
|
||||
all_rows=True,
|
||||
)
|
||||
|
||||
staged = t._last_audio_host_staging[0]
|
||||
assert staged[t.num_codebooks].item() == 1 # valid
|
||||
assert staged[t.num_codebooks + 1].item() == 1 # done
|
||||
assert not (staged[: t.num_codebooks] == -1).any()
|
||||
assert staged[: t.num_codebooks - 1].eq(EOC_ID).all()
|
||||
assert t._decode_generation_done[0].item() is True
|
||||
assert t._decode_has_codes[0].item() is False
|
||||
|
||||
def test_decode_state_growth_preserves_inflight_rows(self):
|
||||
"""Growing decode buffers must copy old rows and initialize new rows inactive."""
|
||||
t = self._make_batched_sampler_talker(num_rows=2)
|
||||
old_codes = t._decode_last_codes.clone()
|
||||
old_has = t._decode_has_codes.clone()
|
||||
old_delay = t._decode_delay_count.clone()
|
||||
old_rem = t._decode_eoc_countdown.clone()
|
||||
old_done = t._decode_generation_done.clone()
|
||||
|
||||
t._ensure_decode_state_capacity(5, torch.device("cpu"))
|
||||
|
||||
assert t._decode_last_codes.shape[0] >= 5
|
||||
assert torch.equal(t._decode_last_codes[:2], old_codes)
|
||||
assert torch.equal(t._decode_has_codes[:2], old_has)
|
||||
assert torch.equal(t._decode_delay_count[:2], old_delay)
|
||||
assert torch.equal(t._decode_eoc_countdown[:2], old_rem)
|
||||
assert torch.equal(t._decode_generation_done[:2], old_done)
|
||||
assert not t._decode_has_codes[2:5].any()
|
||||
assert not t._decode_generation_done[2:5].any()
|
||||
assert t._decode_eoc_countdown[2:5].eq(-1).all()
|
||||
|
||||
def test_mixed_batch_prefill_mask_targets_request_rows(self):
|
||||
"""Mixed prefill/decode must reset only prefill request rows."""
|
||||
t = self._make_batched_sampler_talker(num_rows=3)
|
||||
t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
|
||||
|
||||
mask = t._prefill_row_mask(3, torch.device("cpu"))
|
||||
|
||||
assert torch.equal(mask, torch.tensor([False, True, False]))
|
||||
|
||||
def test_mixed_batch_seed_mask_maps_token_tails_to_request_rows(self):
|
||||
"""Seed detection should return one flag per request, not per token."""
|
||||
t = self._make_batched_sampler_talker(num_rows=3)
|
||||
t._audio_continuation_id = 99999
|
||||
t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
|
||||
t._last_step_input_ids = torch.tensor([99999, 11, 12, 13, 99999])
|
||||
|
||||
mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu"))
|
||||
|
||||
assert torch.equal(mask, torch.tensor([True, False, True]))
|
||||
|
||||
def test_mixed_batch_seed_mask_allows_prefill_tail_audio_token(self):
|
||||
"""A prefill span ending in <|audio|> starts audio-code generation."""
|
||||
t = self._make_batched_sampler_talker(num_rows=3)
|
||||
t._audio_continuation_id = 99999
|
||||
t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
|
||||
t._last_step_input_ids = torch.tensor([99999, 11, 12, 99999, 99999])
|
||||
|
||||
mask = t._audio_seed_mask_from_step_input(3, torch.device("cpu"))
|
||||
|
||||
assert torch.equal(mask, torch.tensor([True, True, True]))
|
||||
|
||||
|
||||
# ---- AC-6: Feedback Method Tests ----
|
||||
|
||||
@@ -596,14 +348,10 @@ class TestFeedbackMethods:
|
||||
|
||||
class FakeTalker:
|
||||
_last_audio_codes = torch.tensor([[100, 200, 300, 400, 500, 600, 700, 800]])
|
||||
_last_audio_code_valid = [True]
|
||||
_last_audio_host_staging = None
|
||||
_last_audio_staging_event = None
|
||||
_postprocess_cursor = 0
|
||||
|
||||
t = FakeTalker()
|
||||
t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
|
||||
t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t)
|
||||
result = t.postprocess(torch.zeros(1, 64))
|
||||
assert "codes" in result
|
||||
assert "audio" in result["codes"]
|
||||
@@ -615,14 +363,10 @@ class TestFeedbackMethods:
|
||||
|
||||
class FakeTalker:
|
||||
_last_audio_codes = torch.tensor([[-1, -1, -1, -1, -1, -1, -1, -1]])
|
||||
_last_audio_code_valid = [False]
|
||||
_last_audio_host_staging = None
|
||||
_last_audio_staging_event = None
|
||||
_postprocess_cursor = 0
|
||||
|
||||
t = FakeTalker()
|
||||
t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
|
||||
t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t)
|
||||
result = t.postprocess(torch.zeros(1, 64))
|
||||
assert result == {}
|
||||
|
||||
@@ -634,14 +378,10 @@ class TestFeedbackMethods:
|
||||
_last_audio_codes = torch.tensor(
|
||||
[[100, 200, 300, 400, 500, 600, 700, 800], [-1, -1, -1, -1, -1, -1, -1, -1]]
|
||||
)
|
||||
_last_audio_code_valid = [True, False]
|
||||
_last_audio_host_staging = None
|
||||
_last_audio_staging_event = None
|
||||
_postprocess_cursor = 0
|
||||
|
||||
t = FakeTalker()
|
||||
t.postprocess = mod.HiggsAudioV3TalkerForConditionalGeneration.postprocess.__get__(t)
|
||||
t._postprocess_impl = mod.HiggsAudioV3TalkerForConditionalGeneration._postprocess_impl.__get__(t)
|
||||
r1 = t.postprocess(torch.zeros(1, 64))
|
||||
assert "codes" in r1
|
||||
r2 = t.postprocess(torch.zeros(1, 64))
|
||||
@@ -665,23 +405,14 @@ class TestAudioFeedback:
|
||||
codebook_size = 1026
|
||||
_audio_continuation_id = 99999 # fake audio token
|
||||
_last_step_query_start_loc = None
|
||||
_decode_active_audio_count = 0
|
||||
_decode_last_codes = torch.zeros(4, 8, dtype=torch.long)
|
||||
_decode_has_codes = torch.zeros(4, dtype=torch.bool)
|
||||
_use_external_decode_cudagraph = False
|
||||
_audio_state = {}
|
||||
multimodal_embedding = embed
|
||||
model = type("M", (), {"embed_tokens": lambda self, ids: torch.zeros(ids.shape[0], 16)})()
|
||||
|
||||
t = FakeTalker()
|
||||
t._decode_last_codes = torch.zeros(256, 8, dtype=torch.long)
|
||||
t._decode_has_codes = torch.zeros(256, dtype=torch.bool)
|
||||
t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
|
||||
t._is_single_token_decode_step = (
|
||||
mod.HiggsAudioV3TalkerForConditionalGeneration._is_single_token_decode_step.__get__(t)
|
||||
)
|
||||
t._step_request_count = mod.HiggsAudioV3TalkerForConditionalGeneration._step_request_count.__get__(t)
|
||||
t._decode_request_token_positions = (
|
||||
mod.HiggsAudioV3TalkerForConditionalGeneration._decode_request_token_positions.__get__(t)
|
||||
)
|
||||
t._ensure_decode_state_capacity = lambda min_bs, device: None
|
||||
return t
|
||||
|
||||
def test_text_positions_unchanged(self):
|
||||
@@ -697,11 +428,9 @@ class TestAudioFeedback:
|
||||
"""Audio position with state should have its embedding replaced."""
|
||||
t = self._make_feedback_talker()
|
||||
audio_id = t._audio_continuation_id
|
||||
t._audio_state[1] = {"last_codes": torch.zeros(8, dtype=torch.long)}
|
||||
t._decode_last_codes[1] = torch.zeros(8, dtype=torch.long)
|
||||
t._decode_has_codes[1] = True
|
||||
t._decode_active_audio_count = 1
|
||||
t._ensure_decode_state_capacity = lambda min_bs, device: None
|
||||
t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3])
|
||||
input_ids = torch.tensor([1, audio_id, 3])
|
||||
hidden = torch.zeros(3, 16)
|
||||
result = t._apply_audio_feedback(hidden, input_ids)
|
||||
@@ -715,80 +444,12 @@ class TestAudioFeedback:
|
||||
"""Audio position without state should not be replaced."""
|
||||
t = self._make_feedback_talker()
|
||||
audio_id = t._audio_continuation_id
|
||||
# No state for row 0
|
||||
# No state for row 0 — _decode_has_codes[0] defaults to False
|
||||
input_ids = torch.tensor([audio_id])
|
||||
hidden = torch.zeros(1, 16)
|
||||
result = t._apply_audio_feedback(hidden, input_ids)
|
||||
assert torch.equal(result, hidden)
|
||||
|
||||
def test_prefill_span_does_not_receive_audio_feedback(self):
|
||||
"""A prefill span can have numel == hidden rows but must not be treated as decode."""
|
||||
t = self._make_feedback_talker()
|
||||
t._last_step_input_ids = torch.tensor([1, 2, 3, 4])
|
||||
t._last_step_query_start_loc = torch.tensor([0, 4])
|
||||
t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long)
|
||||
t._decode_has_codes[0] = True
|
||||
t._decode_active_audio_count = 1
|
||||
hidden = torch.zeros(4, 16)
|
||||
|
||||
result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
|
||||
|
||||
assert torch.equal(result, hidden)
|
||||
|
||||
def test_multi_request_decode_receives_audio_feedback(self):
|
||||
"""Concurrent decode has numel > 1 but each request span is exactly one token."""
|
||||
t = self._make_feedback_talker()
|
||||
t._last_step_input_ids = torch.tensor([10, 11, 12, 13])
|
||||
t._last_step_query_start_loc = torch.tensor([0, 1, 2, 3, 4])
|
||||
t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long)
|
||||
t._decode_has_codes[2] = True
|
||||
t._decode_active_audio_count = 1
|
||||
hidden = torch.zeros(4, 16)
|
||||
|
||||
result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
|
||||
|
||||
assert result[2].abs().sum() > 0
|
||||
assert result[0].abs().sum() == 0
|
||||
assert result[1].abs().sum() == 0
|
||||
assert result[3].abs().sum() == 0
|
||||
|
||||
def test_external_decode_graph_capture_keeps_feedback_ops_without_query_start_loc(self):
|
||||
"""FULL_DECODE capture has no request qsl but still needs dense decode feedback ops."""
|
||||
t = self._make_feedback_talker()
|
||||
t._use_external_decode_cudagraph = True
|
||||
t._last_step_input_ids = torch.tensor([10, 11, 12, 13])
|
||||
t._last_step_query_start_loc = None
|
||||
t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long)
|
||||
t._decode_has_codes[2] = True
|
||||
hidden = torch.zeros(4, 16)
|
||||
|
||||
result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
|
||||
|
||||
assert result[2].abs().sum() > 0
|
||||
assert result[0].abs().sum() == 0
|
||||
assert result[1].abs().sum() == 0
|
||||
assert result[3].abs().sum() == 0
|
||||
|
||||
def test_mixed_prefill_decode_receives_feedback_only_on_decode_rows(self):
|
||||
"""Mixed batches should skip prefill spans but keep audio feedback for decode spans."""
|
||||
t = self._make_feedback_talker()
|
||||
t._last_step_input_ids = torch.tensor([99999, 21, 22, 23, 99999])
|
||||
t._last_step_query_start_loc = torch.tensor([0, 1, 4, 5])
|
||||
t._decode_last_codes[0] = torch.zeros(8, dtype=torch.long)
|
||||
t._decode_last_codes[2] = torch.zeros(8, dtype=torch.long)
|
||||
t._decode_has_codes[0] = True
|
||||
t._decode_has_codes[2] = True
|
||||
t._decode_active_audio_count = 2
|
||||
hidden = torch.zeros(5, 16)
|
||||
|
||||
result = t._apply_audio_feedback(hidden, t._last_step_input_ids)
|
||||
|
||||
assert result[0].abs().sum() > 0
|
||||
assert result[4].abs().sum() > 0
|
||||
assert result[1].abs().sum() == 0
|
||||
assert result[2].abs().sum() == 0
|
||||
assert result[3].abs().sum() == 0
|
||||
|
||||
|
||||
# ---- AC-8: Codec Strictness ----
|
||||
|
||||
@@ -1114,23 +775,6 @@ class TestStageInputProcessor:
|
||||
with pytest.raises(ValueError, match="Not enough frames"):
|
||||
_revert_delay_pattern(codes)
|
||||
|
||||
def test_talker2code2wav_skips_too_few_frames_without_crashing(self):
|
||||
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
|
||||
talker2code2wav,
|
||||
)
|
||||
|
||||
class Output:
|
||||
multimodal_output = {"codes": {"audio": torch.zeros(5, 8, dtype=torch.long)}}
|
||||
|
||||
class TalkerOutput:
|
||||
finished = True
|
||||
outputs = [Output()]
|
||||
|
||||
result = talker2code2wav([TalkerOutput()])
|
||||
|
||||
assert len(result) == 1
|
||||
assert result[0]["prompt_token_ids"] == []
|
||||
|
||||
def test_filter_real_code_frames(self):
|
||||
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
|
||||
_filter_real_code_frames,
|
||||
@@ -1182,14 +826,8 @@ class TestRegistry:
|
||||
|
||||
# __file__ = tests/unit/higgs_audio_v3/test_*.py → 4 dirnames to repo root
|
||||
repo_root = os.path.dirname(os.path.dirname(os.path.dirname(os.path.dirname(__file__))))
|
||||
deploy_dir = os.path.join(repo_root, "vllm_omni", "deploy")
|
||||
for name in (
|
||||
"higgs_multimodal_qwen3.yaml",
|
||||
"higgs_multimodal_qwen3_high_throughput.yaml",
|
||||
"higgs_multimodal_qwen3_low_latency.yaml",
|
||||
):
|
||||
yaml_path = os.path.join(deploy_dir, name)
|
||||
assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
|
||||
yaml_path = os.path.join(repo_root, "vllm_omni", "deploy", "higgs_multimodal_qwen3.yaml")
|
||||
assert os.path.isfile(yaml_path), f"Deploy YAML not found at {yaml_path}"
|
||||
|
||||
|
||||
# ---- AC-3: Prompt Builder ----
|
||||
@@ -1276,89 +914,298 @@ class TestPromptBuilder:
|
||||
assert 151704 not in ids # <|ref_text|>
|
||||
|
||||
|
||||
class TestVoiceCloneReferenceCache:
|
||||
def test_audio_tokenizer_dir_env_accepts_parent_or_subdir(self, tmp_path, monkeypatch):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v2 import higgs_audio_v2_tokenizer as tok
|
||||
# ---- Regression: ramp-down done-before-emit off-by-one ----
|
||||
|
||||
parent = tmp_path / "OmniVoice"
|
||||
subdir = parent / "audio_tokenizer"
|
||||
subdir.mkdir(parents=True)
|
||||
(subdir / "config.json").write_text("{}", encoding="utf-8")
|
||||
|
||||
monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(parent))
|
||||
assert tok._resolve_audio_tokenizer_dir() == str(subdir)
|
||||
def _simulate_delay_state_machine(
|
||||
sampled_codes_per_step: list[torch.Tensor],
|
||||
) -> list[torch.Tensor]:
|
||||
"""Reproduce the delay-pattern state machine from sample().
|
||||
|
||||
monkeypatch.setenv("HIGGS_AUDIO_TOKENIZER_PATH", str(subdir))
|
||||
assert tok._resolve_audio_tokenizer_dir() == str(subdir)
|
||||
Feeds ``sampled_codes_per_step`` (each [Q]) through the exact same
|
||||
logic as HiggsAudioV3TalkerForConditionalGeneration.sample() and
|
||||
returns the list of emitted frames (non -1 codes).
|
||||
"""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import (
|
||||
BOC_ID,
|
||||
EOC_ID,
|
||||
)
|
||||
|
||||
def test_higgs_v3_ref_code_cache_returns_clone(self):
|
||||
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
||||
num_codebooks = 8
|
||||
bos = BOC_ID
|
||||
eos_stream = EOC_ID
|
||||
|
||||
serving = object.__new__(OmniOpenAIServingSpeech)
|
||||
serving._higgs_audio_v3_ref_code_cache = OrderedDict()
|
||||
serving._higgs_audio_v3_ref_code_cache_bytes = 0
|
||||
serving._higgs_audio_v3_ref_code_inflight = {}
|
||||
state: dict = {
|
||||
"num_delay": 0,
|
||||
"num_remaining_delays": None,
|
||||
"audio_out_ids": None,
|
||||
"last_codes": torch.full((num_codebooks,), bos, dtype=torch.long),
|
||||
"should_terminate": False,
|
||||
}
|
||||
|
||||
codes = torch.arange(16, dtype=torch.long).reshape(2, 8)
|
||||
serving._put_higgs_audio_v3_ref_codes("ref-a", codes)
|
||||
emitted: list[torch.Tensor] = []
|
||||
for step_codes in sampled_codes_per_step:
|
||||
if state["should_terminate"]:
|
||||
break
|
||||
|
||||
cached = serving._get_higgs_audio_v3_ref_codes("ref-a")
|
||||
assert cached is not None
|
||||
cached.fill_(0)
|
||||
num_delay = state["num_delay"]
|
||||
num_remaining_delays = state["num_remaining_delays"]
|
||||
this_codes = step_codes.clone()
|
||||
|
||||
cached_again = serving._get_higgs_audio_v3_ref_codes("ref-a")
|
||||
assert cached_again is not None
|
||||
assert torch.equal(cached_again, codes)
|
||||
if num_delay + 1 < num_codebooks:
|
||||
this_codes[num_delay + 1 :] = bos
|
||||
num_delay += 1
|
||||
|
||||
def test_higgs_v3_ref_code_inflight_deduplicates_concurrent_encode(self):
|
||||
from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech
|
||||
if num_remaining_delays is not None:
|
||||
this_codes[: num_codebooks - num_remaining_delays] = eos_stream
|
||||
num_remaining_delays -= 1
|
||||
else:
|
||||
eos_positions = (this_codes == eos_stream).nonzero(as_tuple=False).reshape(-1)
|
||||
if eos_positions.numel() > 0:
|
||||
last_eos_idx = int(eos_positions[-1].item())
|
||||
this_codes[: last_eos_idx + 1] = eos_stream
|
||||
num_remaining_delays = num_codebooks - last_eos_idx - 1
|
||||
|
||||
serving = object.__new__(OmniOpenAIServingSpeech)
|
||||
serving._higgs_audio_v3_ref_code_cache = OrderedDict()
|
||||
serving._higgs_audio_v3_ref_code_cache_bytes = 0
|
||||
serving._higgs_audio_v3_ref_code_inflight = {}
|
||||
calls = 0
|
||||
if num_remaining_delays is not None and num_remaining_delays <= 0:
|
||||
state["num_delay"] = 0
|
||||
state["num_remaining_delays"] = None
|
||||
state["should_terminate"] = True
|
||||
continue
|
||||
|
||||
def encode_reference_audio(_wav, _sr):
|
||||
nonlocal calls
|
||||
calls += 1
|
||||
time.sleep(0.05)
|
||||
return torch.arange(16, dtype=torch.long).reshape(2, 8)
|
||||
state["num_delay"] = num_delay
|
||||
state["num_remaining_delays"] = num_remaining_delays
|
||||
state["last_codes"] = this_codes.clone()
|
||||
emitted.append(this_codes)
|
||||
|
||||
def apply_delay_pattern(codes):
|
||||
return codes + 1
|
||||
return emitted
|
||||
|
||||
async def run():
|
||||
return await asyncio.gather(
|
||||
*[
|
||||
serving._resolve_higgs_audio_v3_ref_codes(
|
||||
"ref-a",
|
||||
object(),
|
||||
24000,
|
||||
encode_reference_audio,
|
||||
apply_delay_pattern,
|
||||
)
|
||||
for _ in range(3)
|
||||
]
|
||||
)
|
||||
|
||||
results = asyncio.run(run())
|
||||
assert calls == 1
|
||||
assert sum(int(inflight_wait) for _, _, inflight_wait in results) == 2
|
||||
for codes, cache_hit, _ in results:
|
||||
assert cache_hit is False
|
||||
assert torch.equal(codes, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1)
|
||||
class TestRampDownOffByOne:
|
||||
"""Regression tests for the done-before-emit off-by-one bug.
|
||||
|
||||
cached, cache_hit, inflight_wait = asyncio.run(
|
||||
serving._resolve_higgs_audio_v3_ref_codes(
|
||||
"ref-a",
|
||||
object(),
|
||||
24000,
|
||||
encode_reference_audio,
|
||||
apply_delay_pattern,
|
||||
)
|
||||
When EOC fires early, the ramp-down must emit enough frames so that
|
||||
Stage1's _revert_delay_pattern receives T >= Q = 8.
|
||||
"""
|
||||
|
||||
def test_eoc_at_step0_terminates(self):
|
||||
"""EOC on cb0 at step 0 terminates; T < Q is handled by Stage1 defensive try/except."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
|
||||
|
||||
Q = 8
|
||||
steps = []
|
||||
for s in range(20):
|
||||
codes = torch.randint(0, 1024, (Q,))
|
||||
if s == 0:
|
||||
codes[0] = EOC_ID
|
||||
steps.append(codes)
|
||||
|
||||
emitted = _simulate_delay_state_machine(steps)
|
||||
assert len(emitted) == Q - 1, f"EOC at step 0: expected {Q - 1} emitted frames, got {len(emitted)}"
|
||||
|
||||
def test_eoc_at_step1_emits_enough_frames(self):
|
||||
"""EOC on cb0 at step 1 must produce >= Q emitted frames."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
|
||||
|
||||
Q = 8
|
||||
steps = []
|
||||
for s in range(20):
|
||||
codes = torch.randint(0, 1024, (Q,))
|
||||
if s == 1:
|
||||
codes[0] = EOC_ID
|
||||
steps.append(codes)
|
||||
|
||||
emitted = _simulate_delay_state_machine(steps)
|
||||
assert len(emitted) >= Q, f"EOC at step 1: expected >= {Q} emitted frames, got {len(emitted)}"
|
||||
|
||||
def test_eoc_after_full_rampup_emits_enough(self):
|
||||
"""EOC after all codebooks are active produces T >> Q."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
|
||||
|
||||
Q = 8
|
||||
eoc_step = 10
|
||||
steps = []
|
||||
for s in range(30):
|
||||
codes = torch.randint(0, 1024, (Q,))
|
||||
if s == eoc_step:
|
||||
codes[0] = EOC_ID
|
||||
steps.append(codes)
|
||||
|
||||
emitted = _simulate_delay_state_machine(steps)
|
||||
assert len(emitted) >= Q, f"EOC at step {eoc_step}: expected >= {Q} emitted frames, got {len(emitted)}"
|
||||
|
||||
def test_last_rampdown_frame_not_all_negative(self):
|
||||
"""The last emitted ramp-down frame must contain real codes, not -1."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
|
||||
|
||||
Q = 8
|
||||
steps = []
|
||||
for s in range(20):
|
||||
codes = torch.randint(0, 1024, (Q,))
|
||||
if s == 1:
|
||||
codes[0] = EOC_ID
|
||||
steps.append(codes)
|
||||
|
||||
emitted = _simulate_delay_state_machine(steps)
|
||||
last = emitted[-1]
|
||||
assert (last >= 0).all(), f"Last ramp-down frame contains negative values: {last.tolist()}"
|
||||
|
||||
def test_emitted_frames_revertible_by_stage1(self):
|
||||
"""Emitted frames from EOC-at-step-1 must be revertible by _revert_delay_pattern."""
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_talker import EOC_ID
|
||||
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
|
||||
_revert_delay_pattern,
|
||||
)
|
||||
assert calls == 1
|
||||
assert cache_hit is True
|
||||
assert inflight_wait is False
|
||||
assert torch.equal(cached, torch.arange(16, dtype=torch.long).reshape(2, 8) + 1)
|
||||
|
||||
Q = 8
|
||||
steps = []
|
||||
for s in range(20):
|
||||
codes = torch.randint(0, 1024, (Q,))
|
||||
if s == 1:
|
||||
codes[0] = EOC_ID
|
||||
steps.append(codes)
|
||||
|
||||
emitted = _simulate_delay_state_machine(steps)
|
||||
# Stack to [T, Q], transpose to [Q, T]
|
||||
codes_qt = torch.stack(emitted, dim=0).t().contiguous()
|
||||
assert codes_qt.shape[0] == Q
|
||||
assert codes_qt.shape[1] >= Q
|
||||
|
||||
# Must not raise
|
||||
reverted = _revert_delay_pattern(codes_qt)
|
||||
assert reverted.shape[0] == Q
|
||||
assert reverted.shape[1] == codes_qt.shape[1] - Q + 1
|
||||
|
||||
|
||||
class TestStage1DefensiveHandling:
|
||||
"""Stage1 must not crash the server when receiving T < Q frames."""
|
||||
|
||||
def test_revert_delay_pattern_raises_on_insufficient_frames(self):
|
||||
"""_revert_delay_pattern correctly raises ValueError for T < Q."""
|
||||
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
|
||||
_revert_delay_pattern,
|
||||
)
|
||||
|
||||
codes = torch.zeros(8, 7) # T=7 < Q=8
|
||||
with pytest.raises(ValueError, match="Not enough frames"):
|
||||
_revert_delay_pattern(codes)
|
||||
|
||||
def test_talker2code2wav_survives_insufficient_frames(self):
|
||||
"""talker2code2wav must not crash when audio_codes has T < Q frames."""
|
||||
from vllm_omni.model_executor.stage_input_processors.higgs_audio_v3 import (
|
||||
talker2code2wav,
|
||||
)
|
||||
|
||||
# Build a fake talker output with only 7 frames (T=7 < Q=8)
|
||||
class FakeOutput:
|
||||
def __init__(self):
|
||||
self.multimodal_output = {
|
||||
"codes": {"audio": torch.zeros(7, 8, dtype=torch.long)},
|
||||
}
|
||||
|
||||
class FakeTalkerOutput:
|
||||
finished = True
|
||||
outputs = [FakeOutput()]
|
||||
|
||||
# Must not raise — should return gracefully
|
||||
result = talker2code2wav([FakeTalkerOutput()])
|
||||
assert len(result) == 1
|
||||
# Result is an OmniTokensPrompt with empty prompt_token_ids
|
||||
out = result[0]
|
||||
token_ids = out.prompt_token_ids if hasattr(out, "prompt_token_ids") else out.get("prompt_token_ids", [])
|
||||
assert token_ids == []
|
||||
|
||||
|
||||
# ---- Regression: buffer growth + prefill guard ----
|
||||
|
||||
|
||||
class TestBufferGrowthPreservesState:
|
||||
"""Bug #1: buffer growth in _apply_audio_feedback must not wipe existing state."""
|
||||
|
||||
def _make_feedback_talker_with_buffers(self):
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3 import higgs_audio_v3_talker as mod
|
||||
|
||||
embed = mod.HiggsFusedMultiTextEmbedding(num_codebooks=8, vocab_size=1026, hidden_size=16)
|
||||
torch.nn.init.ones_(embed.weight)
|
||||
|
||||
class FakeTalker:
|
||||
num_codebooks = 8
|
||||
codebook_size = 1026
|
||||
multimodal_embedding = embed
|
||||
|
||||
t = FakeTalker()
|
||||
t._decode_last_codes = torch.zeros(4, 8, dtype=torch.long)
|
||||
t._decode_has_codes = torch.zeros(4, dtype=torch.bool)
|
||||
t._apply_audio_feedback = mod.HiggsAudioV3TalkerForConditionalGeneration._apply_audio_feedback.__get__(t)
|
||||
return t
|
||||
|
||||
def test_buffer_growth_preserves_existing_codes(self):
|
||||
"""Growing the buffer must copy old _decode_last_codes data."""
|
||||
t = self._make_feedback_talker_with_buffers()
|
||||
# Set state for slot 2
|
||||
t._decode_last_codes[2] = torch.tensor([100, 200, 300, 400, 500, 600, 700, 800])
|
||||
t._decode_has_codes[2] = True
|
||||
|
||||
# Trigger growth by calling with bs=8 > current capacity 4
|
||||
hidden = torch.randn(8, 16)
|
||||
t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long))
|
||||
|
||||
# After growth, slot 2 data must still be present
|
||||
assert t._decode_has_codes[2].item() is True
|
||||
assert t._decode_last_codes[2, 0].item() == 100
|
||||
|
||||
def test_buffer_growth_new_slots_are_inactive(self):
|
||||
"""New slots from buffer growth must have has_codes=False."""
|
||||
t = self._make_feedback_talker_with_buffers()
|
||||
t._decode_has_codes[1] = True
|
||||
|
||||
hidden = torch.randn(8, 16)
|
||||
t._apply_audio_feedback(hidden, torch.zeros(8, dtype=torch.long))
|
||||
|
||||
# Slots 4-7 (new) must be inactive
|
||||
assert not t._decode_has_codes[4].item()
|
||||
assert not t._decode_has_codes[7].item()
|
||||
# Slot 1 (old) must still be active
|
||||
assert t._decode_has_codes[1].item()
|
||||
|
||||
|
||||
class TestPrefillGuard:
|
||||
"""Bug #2: _apply_audio_feedback must NOT run during prefill.
|
||||
|
||||
The prefill detection uses attn_metadata.max_query_len (== 1 for decode,
|
||||
> 1 for prefill). This correctly handles concurrent decode (N requests,
|
||||
each contributing 1 token → numel=N but max_query_len=1).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _is_prefill_via_max_query_len(max_query_len: int | None, input_ids: torch.Tensor) -> bool:
|
||||
"""Mirror the actual detection logic in forward()."""
|
||||
if max_query_len is not None:
|
||||
return int(max_query_len) > 1
|
||||
return input_ids is not None and int(input_ids.numel()) > 1
|
||||
|
||||
def test_prefill_detected_via_max_query_len(self):
|
||||
"""Prefill: max_query_len > 1 → is_prefill=True."""
|
||||
input_ids = torch.tensor([1, 2, 3, 4, 5])
|
||||
assert self._is_prefill_via_max_query_len(5, input_ids) is True
|
||||
|
||||
def test_single_decode_detected_via_max_query_len(self):
|
||||
"""Single-request decode: max_query_len=1 → is_prefill=False."""
|
||||
input_ids = torch.tensor([42])
|
||||
assert self._is_prefill_via_max_query_len(1, input_ids) is False
|
||||
|
||||
def test_concurrent_decode_not_misclassified(self):
|
||||
"""Multi-request decode: numel=4 but max_query_len=1 → is_prefill=False.
|
||||
|
||||
This is the regression caught by reviewer: with c=4 concurrent decode
|
||||
requests, input_ids.numel()==4 but each request contributes exactly 1
|
||||
token, so max_query_len==1 and audio feedback must still apply.
|
||||
"""
|
||||
input_ids = torch.tensor([42, 43, 44, 45]) # 4 concurrent decode
|
||||
assert self._is_prefill_via_max_query_len(1, input_ids) is False
|
||||
|
||||
def test_fallback_when_max_query_len_unavailable(self):
|
||||
"""When attn_metadata is unavailable, fall back to numel heuristic."""
|
||||
prefill_ids = torch.tensor([1, 2, 3])
|
||||
assert self._is_prefill_via_max_query_len(None, prefill_ids) is True
|
||||
decode_ids = torch.tensor([42])
|
||||
assert self._is_prefill_via_max_query_len(None, decode_ids) is False
|
||||
@@ -22,7 +22,6 @@ class DummyInputBatch:
|
||||
|
||||
def __init__(self, req_ids):
|
||||
self.req_ids = req_ids
|
||||
self.req_id_to_index = {r: i for i, r in enumerate(req_ids)}
|
||||
|
||||
|
||||
class DummyReqState:
|
||||
|
||||
@@ -16,8 +16,8 @@ from datetime import datetime
|
||||
from typing import Any, Literal
|
||||
|
||||
import aiohttp
|
||||
import numpy as np
|
||||
import pybase64 as base64
|
||||
from pydub import AudioSegment
|
||||
from tqdm.asyncio import tqdm
|
||||
from vllm.benchmarks import datasets
|
||||
from vllm.benchmarks.datasets import SampleRequest
|
||||
@@ -653,8 +653,9 @@ async def async_request_openai_chat_omni_completions(
|
||||
# Reset per-attempt state so that retries do not mix partial
|
||||
# outputs or metrics from previous attempts.
|
||||
generated_text = ""
|
||||
generated_audio = None
|
||||
# For wav responses, accumulate decoded PCM bytes per chunk
|
||||
# to avoid repeated decode/concat.
|
||||
# to avoid repeated AudioSegment decode/concat.
|
||||
wav_pcm_buffer = bytearray()
|
||||
wav_audio_params: tuple[int, int, int] | None = None
|
||||
wav_inconsistent_chunk_count = 0
|
||||
@@ -794,33 +795,30 @@ async def async_request_openai_chat_omni_completions(
|
||||
|
||||
output.latency = timestamp - st
|
||||
output.generated_text = generated_text
|
||||
audio_duration_sec = 0.0
|
||||
audio_frames = 0
|
||||
if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None:
|
||||
channels, sample_width, frame_rate = wav_audio_params
|
||||
frame_width = sample_width * channels
|
||||
if frame_width > 0:
|
||||
audio_frames = len(wav_pcm_buffer) // frame_width
|
||||
audio_duration_sec = audio_frames / frame_rate
|
||||
else:
|
||||
logger.warning("Audio frame width is zero")
|
||||
generated_audio = AudioSegment(
|
||||
data=bytes(wav_pcm_buffer),
|
||||
sample_width=sample_width,
|
||||
frame_rate=frame_rate,
|
||||
channels=channels,
|
||||
)
|
||||
elif audio_bytes_buffer:
|
||||
try:
|
||||
from vllm.multimodal.audio import get_audio_duration
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
waveform, sr = load_audio(
|
||||
generated_audio = AudioSegment.from_file(
|
||||
io.BytesIO(bytes(audio_bytes_buffer)),
|
||||
sr=None,
|
||||
mono=False,
|
||||
format=response_format,
|
||||
)
|
||||
audio_duration_sec = get_audio_duration(y=waveform, sr=sr)
|
||||
audio_frames = int(audio_duration_sec * sr)
|
||||
except Exception as ex:
|
||||
logger.warning("Failed to decode accumulated audio bytes: %s", ex)
|
||||
if audio_duration_sec > 0 or audio_frames > 0:
|
||||
output.audio_duration = audio_duration_sec
|
||||
output.audio_frames = audio_frames
|
||||
if generated_audio is not None:
|
||||
output.audio_duration = len(generated_audio) / 1000.0
|
||||
frame_width = generated_audio.frame_width
|
||||
if frame_width > 0:
|
||||
output.audio_frames = len(generated_audio.raw_data) // frame_width
|
||||
else:
|
||||
output.audio_frames = 0
|
||||
logger.warning("Audio frame width is zero")
|
||||
audio_duration = output.audio_duration
|
||||
if audio_duration > 0:
|
||||
output.audio_rtf = audio_generate_time / output.audio_duration
|
||||
@@ -829,29 +827,8 @@ async def async_request_openai_chat_omni_completions(
|
||||
logger.warning("Audio duration is zero")
|
||||
if _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False):
|
||||
try:
|
||||
if response_format == "wav" and wav_pcm_buffer and wav_audio_params is not None:
|
||||
from vllm.multimodal.audio import AudioResampler
|
||||
|
||||
pcm_channels, pcm_sw, pcm_rate = wav_audio_params
|
||||
pcm = np.frombuffer(
|
||||
bytes(wav_pcm_buffer), dtype=np.int16 if pcm_sw == 2 else np.float32
|
||||
)
|
||||
if pcm_channels > 1:
|
||||
pcm = pcm.reshape(-1, pcm_channels).mean(axis=1).astype(pcm.dtype)
|
||||
pcm_f32 = pcm.astype(np.float32) / 32767.0 if pcm.dtype == np.int16 else pcm
|
||||
if pcm_rate != 24000:
|
||||
resampler = AudioResampler(target_sr=24000)
|
||||
pcm_f32 = resampler.resample(pcm_f32, orig_sr=pcm_rate)
|
||||
output.tts_output_pcm_bytes = (pcm_f32 * 32767).astype(np.int16).tobytes()
|
||||
elif audio_bytes_buffer:
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
|
||||
waveform, _ = load_audio(
|
||||
io.BytesIO(bytes(audio_bytes_buffer)),
|
||||
sr=24000,
|
||||
mono=True,
|
||||
)
|
||||
output.tts_output_pcm_bytes = (waveform * 32767).astype(np.int16).tobytes()
|
||||
seg = generated_audio.set_frame_rate(24000).set_channels(1).set_sample_width(2)
|
||||
output.tts_output_pcm_bytes = bytes(seg.raw_data)
|
||||
except Exception as ex:
|
||||
logger.warning("seed_tts WER PCM export failed: %s", ex)
|
||||
output.success = True
|
||||
|
||||
@@ -101,10 +101,6 @@ _OMNI_PIPELINES: dict[str, tuple[str, str]] = {
|
||||
"vllm_omni.model_executor.models.mimo_audio.pipeline",
|
||||
"MIMO_AUDIO_PIPELINE",
|
||||
),
|
||||
"ming_tts": (
|
||||
"vllm_omni.model_executor.models.ming_tts.pipeline",
|
||||
"MING_TTS_PIPELINE",
|
||||
),
|
||||
"voxtral_tts": (
|
||||
"vllm_omni.model_executor.models.voxtral_tts.pipeline",
|
||||
"VOXTRAL_TTS_PIPELINE",
|
||||
|
||||
@@ -1128,19 +1128,6 @@ class StageConfigFactory:
|
||||
continue
|
||||
return cls._create_from_registry(registered.model_type, cli_overrides, deploy_config_path)
|
||||
|
||||
# An explicit deploy config can select a pipeline when the model's HF
|
||||
# metadata is too generic for registry auto-detection.
|
||||
if deploy_config_path is not None:
|
||||
deploy_path = Path(deploy_config_path)
|
||||
if deploy_path.exists():
|
||||
deploy_config = load_deploy_config(deploy_path)
|
||||
if deploy_config.pipeline in _PIPELINE_REGISTRY:
|
||||
return cls._create_from_registry(
|
||||
deploy_config.pipeline,
|
||||
cli_overrides,
|
||||
deploy_config_path,
|
||||
)
|
||||
|
||||
# Not in the pipeline registry — let the caller fall back to the
|
||||
# legacy ``stage_configs/*.yaml`` path (resolve_model_config_path).
|
||||
return None
|
||||
|
||||
@@ -218,9 +218,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
self._process_pending_input_timeouts()
|
||||
|
||||
if self.chunk_transfer_adapter:
|
||||
self.chunk_transfer_adapter.process_pending_chunks(
|
||||
self.waiting, self.running, scheduler_requests=self.requests
|
||||
)
|
||||
self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
|
||||
|
||||
try:
|
||||
scheduler_output = super().schedule()
|
||||
@@ -628,34 +626,7 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
if self.chunk_transfer_adapter:
|
||||
self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
|
||||
|
||||
# Realign stale ``request.status`` (chunk-transfer-adapter's
|
||||
# ``requests_origin_status`` table doesn't follow the
|
||||
# ``waiting → running`` admit transition; without this, an abort
|
||||
# arriving between admit and the next deque round-trip leaves
|
||||
# the request in ``self.running`` with ``status=WAITING`` and
|
||||
# upstream ``Scheduler.finish_requests`` silently fails to
|
||||
# release the worker's ``input_batch`` slot -- after
|
||||
# ``max_num_seqs`` such aborts new requests hang at
|
||||
# ``chunks=0``). Only the ``async_chunk`` path triggers the
|
||||
# staleness; with ``async_chunk`` disabled this is a cheap O(n)
|
||||
# no-op over an already-aligned set, kept unconditional so the
|
||||
# abort path stays uniform across configurations. See
|
||||
# ``OmniSchedulerMixin._realign_request_status_to_queues`` and
|
||||
# #3774 discussion.
|
||||
self._realign_request_status_to_queues(request_ids)
|
||||
|
||||
finished = super().finish_requests(request_ids, finished_status)
|
||||
|
||||
# Defensive post-finish purge: belt-and-suspenders to the
|
||||
# realignment above. Even after realign + ``super()``, corner
|
||||
# cases (mid-transition status, connector cleanups that pop
|
||||
# from ``self.requests`` without unwinding ``self.running``)
|
||||
# can leave already-finished or untracked entries in
|
||||
# ``self.running``. Sweep them now so the worker's
|
||||
# ``input_batch`` slot never pins a freed request and starves
|
||||
# new admissions. See ``OmniSchedulerMixin._purge_finished_from_running``.
|
||||
self._purge_finished_from_running()
|
||||
|
||||
input_coordinator = getattr(self, "input_coordinator", None)
|
||||
if input_coordinator is not None:
|
||||
for request_id, _ in finished:
|
||||
|
||||
@@ -87,9 +87,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
self._consume_pending_connector_output(model_mode="generation")
|
||||
self._process_pending_input_timeouts()
|
||||
if self.chunk_transfer_adapter:
|
||||
self.chunk_transfer_adapter.process_pending_chunks(
|
||||
self.waiting, self.running, scheduler_requests=self.requests
|
||||
)
|
||||
self.chunk_transfer_adapter.process_pending_chunks(self.waiting, self.running)
|
||||
|
||||
# OMNI: Track requests that are already finished (e.g., marked by connector)
|
||||
# These should be removed from running and not scheduled
|
||||
@@ -366,23 +364,7 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler):
|
||||
if self.chunk_transfer_adapter:
|
||||
self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests)
|
||||
|
||||
# See ``OmniSchedulerMixin._realign_request_status_to_queues`` --
|
||||
# closes the residual hang on the
|
||||
# ``waiting → running → abort-before-next-deque-round-trip`` race
|
||||
# surfaced by #3774's reproduction matrix. Only the
|
||||
# ``async_chunk`` path actually triggers the staleness; with
|
||||
# ``async_chunk`` disabled this is a cheap O(n) no-op, kept
|
||||
# unconditional so the abort path stays uniform.
|
||||
self._realign_request_status_to_queues(request_ids)
|
||||
|
||||
finished = super().finish_requests(request_ids, finished_status)
|
||||
|
||||
# See ``OmniSchedulerMixin._purge_finished_from_running`` --
|
||||
# defensive belt-and-suspenders sweep paired with the realign
|
||||
# above. Closes residual ``self.running`` slot leaks if any
|
||||
# corner case slips past upstream's status-driven removal.
|
||||
self._purge_finished_from_running()
|
||||
|
||||
if self.input_coordinator is not None:
|
||||
for request_id, _ in finished:
|
||||
self._free_input_coordinator_request(request_id)
|
||||
|
||||
Loaded 100 of 211 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user