Compare commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
5a7cda40c3 | ||
|
|
1be1a71156 | ||
|
|
3fe3aa359b | ||
|
|
541f06c60d | ||
|
|
a3834e32cf | ||
|
|
152343d5fb | ||
|
|
03cc888485 | ||
|
|
e27fd46ab9 | ||
|
|
eb964dcec5 | ||
|
|
5414f78e5f | ||
|
|
f3d40b5402 | ||
|
|
3529b8280d | ||
|
|
e26f5cb918 | ||
|
|
0d3c259a76 | ||
|
|
74b5ae5ebe | ||
|
|
1fe267d6c9 | ||
|
|
14879d6619 | ||
|
|
18e083e512 | ||
|
|
a1a53e5f94 | ||
|
|
1b2ab23586 | ||
|
|
57ea047acf | ||
|
|
0342827dd7 | ||
|
|
362b3ab013 | ||
|
|
c2b4564a67 | ||
|
|
aa92a3a35d | ||
|
|
2fa4c7a4a9 | ||
|
|
39ae6e12f8 | ||
|
|
2d9e5e0515 | ||
|
|
d50bf1aa42 | ||
|
|
f1bf544208 | ||
|
|
57c8f3abfc | ||
|
|
6ae355b9a9 | ||
|
|
73df8326d7 | ||
|
|
95f633f830 | ||
|
|
e1b0b83238 | ||
|
|
57227dc7f3 | ||
|
|
ca75623ff3 | ||
|
|
3753b21c3d | ||
|
|
e0fde88a5f | ||
|
|
bf03413c36 | ||
|
|
f174ee54b0 | ||
|
|
6ddaf6188d | ||
|
|
02eb7a4ecb | ||
|
|
cb1c5a0c39 | ||
|
|
1b85df0510 | ||
|
|
5de3b11822 | ||
|
|
7dcbeaea03 | ||
|
|
b8f68174ad | ||
|
|
b5352a5869 | ||
|
|
2f25195750 | ||
|
|
9c9d4ed4a8 | ||
|
|
13f9ecaef8 | ||
|
|
72cc496fb6 | ||
|
|
fc7393e485 | ||
|
|
5f1a876cbd | ||
|
|
eb1b782dd1 | ||
|
|
29b18ea298 | ||
|
|
176b00538c | ||
|
|
4d71a162e6 | ||
|
|
923c1d9402 | ||
|
|
82cce633a3 | ||
|
|
fce0f51089 | ||
|
|
82e74ce972 | ||
|
|
963ba1ab7a | ||
|
|
49c102589b | ||
|
|
930c317b2b |
No files matched your search
@@ -1,14 +1,14 @@
|
||||
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
|
||||
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
|
||||
# exception in upload_pipeline_with_skip_ci.sh (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# exception in upload_pipeline.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
|
||||
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline_with_skip_ci.sh (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
|
||||
# Document 2: appended after `---`; same file, read by upload_pipeline.py (not evaluated as a second pipeline by Buildkite). Child uploads include test-ready.yml, test-merge.yml, test-nightly.yml, and test-weekly.yml.
|
||||
steps:
|
||||
- label: ":github: Resolve skip-ci & upload pipeline"
|
||||
key: upload-ci-pipeline
|
||||
commands:
|
||||
- "bash .buildkite/scripts/upload_pipeline_with_skip_ci.sh"
|
||||
- "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -31,7 +31,7 @@ steps:
|
||||
key: upload-ready-pipeline
|
||||
if: __UPLOAD_READY_IF__
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -41,7 +41,7 @@ steps:
|
||||
key: upload-merge-pipeline
|
||||
if: __UPLOAD_MERGE_IF__
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-ready.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -23,7 +23,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-merge.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -0,0 +1,308 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""Render and optionally upload Buildkite pipeline YAML with diff-aware logic.
|
||||
|
||||
Bootstrap mode (pipeline.yml with __IMAGE_BUILD_IF__ placeholders):
|
||||
- Detect docs-only PR/main changes and substitute skip-ci ``if`` expressions.
|
||||
|
||||
Test pipeline mode (e.g. test-merge.yml):
|
||||
- Drop steps or groups whose ``source_file_dependencies`` do not match changed files.
|
||||
- ``source_file_dependencies`` may be set on a leaf step or on a group step.
|
||||
- Always strip ``source_file_dependencies`` from uploaded YAML (Buildkite does not
|
||||
recognize this uploader-only key).
|
||||
|
||||
Usage:
|
||||
python3 upload_pipeline.py [--upload] <pipeline.yml>
|
||||
|
||||
# Bootstrap (replaces upload_pipeline_with_skip_ci.sh):
|
||||
python3 upload_pipeline.py --upload .buildkite/pipeline.yml
|
||||
|
||||
# Test pipeline (replaces upload_test_pipeline_with_diff_skip.py):
|
||||
python3 upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import yaml
|
||||
except ImportError:
|
||||
subprocess.check_call(
|
||||
[sys.executable, "-m", "pip", "install", "-q", "pyyaml"],
|
||||
)
|
||||
import yaml
|
||||
|
||||
LOG = "upload_pipeline"
|
||||
ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
DOC_SEP = "\n---\n"
|
||||
BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
|
||||
|
||||
|
||||
def _log(message: str) -> None:
|
||||
print(f"{LOG}: {message}", file=sys.stderr)
|
||||
|
||||
|
||||
def _git(*args: str) -> subprocess.CompletedProcess[str]:
|
||||
return subprocess.run(
|
||||
["git", *args],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
text=True,
|
||||
)
|
||||
|
||||
|
||||
def resolve_changed_files() -> list[str] | None:
|
||||
"""Return changed file paths, or None when diff cannot be resolved."""
|
||||
is_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false" and os.environ.get(
|
||||
"BUILDKITE_PULL_REQUEST", ""
|
||||
)
|
||||
commit = os.environ.get("BUILDKITE_COMMIT", "")
|
||||
|
||||
if is_pr:
|
||||
base_branch = os.environ.get("BUILDKITE_PULL_REQUEST_BASE_BRANCH", "main")
|
||||
base_ref = f"origin/{base_branch}"
|
||||
if _git("rev-parse", "--verify", base_ref).returncode != 0:
|
||||
_log(f"origin/{base_branch} not found locally; trying fetch")
|
||||
_git("fetch", "--depth=200", "origin", base_branch)
|
||||
if _git("rev-parse", "--verify", base_ref).returncode != 0:
|
||||
if _git("rev-parse", "--verify", base_branch).returncode == 0:
|
||||
base_ref = base_branch
|
||||
else:
|
||||
_log(f"cannot resolve PR base {base_branch}; using safe defaults")
|
||||
return None
|
||||
diff_range = f"{base_ref}...{commit}"
|
||||
elif os.environ.get("BUILDKITE_BRANCH", "") == "main":
|
||||
if _git("rev-parse", "--verify", f"{commit}^").returncode != 0:
|
||||
_log("main commit has no parent; using safe defaults")
|
||||
return None
|
||||
diff_range = f"{commit}^..{commit}"
|
||||
else:
|
||||
_log("not PR/main build; using safe defaults")
|
||||
return None
|
||||
|
||||
result = _git("diff", "--name-only", diff_range)
|
||||
if result.returncode != 0:
|
||||
_log(f"git diff failed ({diff_range}); using safe defaults")
|
||||
return None
|
||||
|
||||
files = [line for line in result.stdout.splitlines() if line.strip()]
|
||||
_log(f"{len(files)} changed file(s)")
|
||||
return files
|
||||
|
||||
|
||||
def is_docs_only_change(changed_files: list[str]) -> bool:
|
||||
has_any = False
|
||||
for file_path in changed_files:
|
||||
if not file_path:
|
||||
continue
|
||||
has_any = True
|
||||
if file_path.startswith("docs/"):
|
||||
continue
|
||||
if file_path.endswith(".md"):
|
||||
continue
|
||||
if file_path == "mkdocs.yml":
|
||||
continue
|
||||
return False
|
||||
return has_any
|
||||
|
||||
|
||||
def resolve_skip_ci(changed_files: list[str] | None) -> bool:
|
||||
if changed_files is None:
|
||||
_log("skip-ci=0 (could not resolve changed files)")
|
||||
return False
|
||||
if is_docs_only_change(changed_files):
|
||||
_log("docs-only change detected; skip-ci=1")
|
||||
return True
|
||||
_log("non-doc changes detected; skip-ci=0")
|
||||
return False
|
||||
|
||||
|
||||
def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
|
||||
if DOC_SEP in text:
|
||||
_, continuation = text.split(DOC_SEP, 1)
|
||||
else:
|
||||
continuation = text
|
||||
|
||||
nightly_only = (
|
||||
'(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
|
||||
)
|
||||
if skip_ci:
|
||||
image_if = f"'{nightly_only}'"
|
||||
ready_if = "'false'"
|
||||
merge_if = "'false'"
|
||||
else:
|
||||
image_if = "'true'"
|
||||
ready_if = '\'build.branch != "main" && build.pull_request.labels includes "ready"\''
|
||||
merge_if = (
|
||||
'\'(build.branch == "main" && build.env("NIGHTLY") != "1" '
|
||||
'&& build.env("WEEKLY") != "1") || '
|
||||
'(build.branch != "main" && build.pull_request.labels includes "merge-test")\''
|
||||
)
|
||||
|
||||
return (
|
||||
continuation.replace("__IMAGE_BUILD_IF__", image_if)
|
||||
.replace("__UPLOAD_READY_IF__", ready_if)
|
||||
.replace("__UPLOAD_MERGE_IF__", merge_if)
|
||||
)
|
||||
|
||||
|
||||
def _matches_dependencies(changed_files: list[str], prefixes: list[str]) -> bool:
|
||||
for path in changed_files:
|
||||
for prefix in prefixes:
|
||||
normalized = prefix.rstrip("/")
|
||||
if path == normalized or path.startswith(f"{normalized}/"):
|
||||
return True
|
||||
return False
|
||||
|
||||
|
||||
def _strip_source_file_dependencies(step: dict[str, Any]) -> dict[str, Any]:
|
||||
return {key: value for key, value in step.items() if key != "source_file_dependencies"}
|
||||
|
||||
|
||||
def _step_label(step: dict[str, Any]) -> str:
|
||||
return str(step.get("group") or step.get("label") or "<step>")
|
||||
|
||||
|
||||
def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
|
||||
filtered: list[Any] = []
|
||||
for step in steps:
|
||||
if not isinstance(step, dict):
|
||||
filtered.append(step)
|
||||
continue
|
||||
|
||||
deps = step.get("source_file_dependencies")
|
||||
if deps is not None and not isinstance(deps, list):
|
||||
raise ValueError(
|
||||
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
|
||||
)
|
||||
if deps is not None and not _matches_dependencies(changed_files, deps):
|
||||
_log(
|
||||
f"skip {_step_label(step)!r} (no changes under {deps})",
|
||||
)
|
||||
continue
|
||||
|
||||
nested = step.get("steps")
|
||||
if nested is not None:
|
||||
kept_nested = _filter_steps(nested, changed_files)
|
||||
if not kept_nested:
|
||||
_log(f"omit empty group {_step_label(step)!r}")
|
||||
continue
|
||||
new_step = _strip_source_file_dependencies(step)
|
||||
new_step["steps"] = kept_nested
|
||||
filtered.append(new_step)
|
||||
continue
|
||||
|
||||
if deps is not None:
|
||||
filtered.append(_strip_source_file_dependencies(step))
|
||||
else:
|
||||
filtered.append(step)
|
||||
|
||||
return filtered
|
||||
|
||||
|
||||
def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
|
||||
"""Remove uploader-only keys while keeping all steps (no diff filtering)."""
|
||||
stripped: list[Any] = []
|
||||
for step in steps:
|
||||
if not isinstance(step, dict):
|
||||
stripped.append(step)
|
||||
continue
|
||||
|
||||
deps = step.get("source_file_dependencies")
|
||||
if deps is not None and not isinstance(deps, list):
|
||||
raise ValueError(
|
||||
f"source_file_dependencies must be a list in step {_step_label(step)!r}",
|
||||
)
|
||||
|
||||
nested = step.get("steps")
|
||||
new_step = _strip_source_file_dependencies(step)
|
||||
if nested is not None:
|
||||
new_step["steps"] = _strip_uploader_metadata_from_steps(nested)
|
||||
stripped.append(new_step)
|
||||
|
||||
return stripped
|
||||
|
||||
|
||||
def render_test_pipeline(
|
||||
doc: dict[str, Any],
|
||||
changed_files: list[str] | None,
|
||||
) -> dict[str, Any]:
|
||||
steps = doc.get("steps")
|
||||
if not isinstance(steps, list):
|
||||
return doc
|
||||
if changed_files is not None:
|
||||
steps = _filter_steps(steps, changed_files)
|
||||
else:
|
||||
steps = _strip_uploader_metadata_from_steps(steps)
|
||||
return {**doc, "steps": steps}
|
||||
|
||||
|
||||
def resolve_pipeline_path(arg: str) -> Path:
|
||||
path = Path(arg)
|
||||
if path.is_absolute():
|
||||
return path
|
||||
return ROOT / path
|
||||
|
||||
|
||||
def render_pipeline(path: Path) -> str:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
changed_files = resolve_changed_files()
|
||||
|
||||
if BOOTSTRAP_MARKER in text:
|
||||
rendered = render_bootstrap_pipeline(text, skip_ci=resolve_skip_ci(changed_files))
|
||||
return rendered
|
||||
|
||||
doc = yaml.safe_load(text)
|
||||
if not isinstance(doc, dict):
|
||||
raise ValueError(f"invalid pipeline YAML: {path}")
|
||||
|
||||
doc = render_test_pipeline(doc, changed_files)
|
||||
|
||||
return yaml.safe_dump(doc, sort_keys=False)
|
||||
|
||||
|
||||
def upload_to_buildkite(content: str) -> None:
|
||||
subprocess.run(
|
||||
["buildkite-agent", "pipeline", "upload"],
|
||||
input=content,
|
||||
text=True,
|
||||
check=True,
|
||||
)
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument(
|
||||
"pipeline",
|
||||
nargs="?",
|
||||
default=".buildkite/pipeline.yml",
|
||||
help="Pipeline YAML path (default: .buildkite/pipeline.yml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--upload",
|
||||
action="store_true",
|
||||
help="Pipe rendered YAML to buildkite-agent pipeline upload",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
path = resolve_pipeline_path(args.pipeline)
|
||||
if not path.is_file():
|
||||
_log(f"missing pipeline file: {path}")
|
||||
return 1
|
||||
|
||||
rendered = render_pipeline(path)
|
||||
if args.upload:
|
||||
upload_to_buildkite(rendered)
|
||||
else:
|
||||
sys.stdout.write(rendered)
|
||||
return 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
raise SystemExit(main())
|
||||
@@ -1,146 +0,0 @@
|
||||
#!/usr/bin/env bash
|
||||
# Evaluate docs-only skip-ci and upload continuation steps from the given pipeline YAML
|
||||
# (YAML document after the first `---`). Buildkite `if` is evaluated at upload time.
|
||||
#
|
||||
# Usage: upload_pipeline_with_skip_ci.sh [pipeline_yaml]
|
||||
# pipeline_yaml: path relative to repo root or absolute (default: .buildkite/pipeline.yml)
|
||||
set -euo pipefail
|
||||
|
||||
ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
|
||||
PIPELINE_ARG="${1:-.buildkite/pipeline.yml}"
|
||||
if [[ "${PIPELINE_ARG}" = /* ]]; then
|
||||
PIPELINE_YML="${PIPELINE_ARG}"
|
||||
else
|
||||
PIPELINE_YML="${ROOT}/${PIPELINE_ARG}"
|
||||
fi
|
||||
|
||||
# Prints a single digit to stdout: 1 = skip image CI, 0 = run. Logs go to stderr.
|
||||
is_docs_only_change() {
|
||||
local file_path
|
||||
local has_any=0
|
||||
|
||||
while IFS= read -r file_path; do
|
||||
[[ -z "${file_path}" ]] && continue
|
||||
has_any=1
|
||||
|
||||
if [[ "${file_path}" == docs/* ]]; then
|
||||
continue
|
||||
fi
|
||||
if [[ "${file_path}" == *.md ]]; then
|
||||
continue
|
||||
fi
|
||||
if [[ "${file_path}" == "mkdocs.yml" ]]; then
|
||||
continue
|
||||
fi
|
||||
return 1
|
||||
done
|
||||
|
||||
[[ "${has_any}" -eq 1 ]]
|
||||
}
|
||||
|
||||
resolve_skip_ci() {
|
||||
local is_pr_build=0
|
||||
local files
|
||||
local base_branch base_ref
|
||||
|
||||
if [[ "${BUILDKITE_PULL_REQUEST:-false}" != "false" && -n "${BUILDKITE_PULL_REQUEST:-}" ]]; then
|
||||
is_pr_build=1
|
||||
fi
|
||||
|
||||
if [[ "${is_pr_build}" -eq 1 ]]; then
|
||||
base_branch="${BUILDKITE_PULL_REQUEST_BASE_BRANCH:-main}"
|
||||
if ! git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
|
||||
echo "resolve_skip_ci: origin/${base_branch} not found locally; trying fetch" >&2
|
||||
git fetch --depth=200 origin "${base_branch}" >/dev/null 2>&1 || true
|
||||
fi
|
||||
|
||||
base_ref=""
|
||||
if git rev-parse --verify "origin/${base_branch}" >/dev/null 2>&1; then
|
||||
base_ref="origin/${base_branch}"
|
||||
elif git rev-parse --verify "${base_branch}" >/dev/null 2>&1; then
|
||||
base_ref="${base_branch}"
|
||||
else
|
||||
echo "resolve_skip_ci: cannot resolve PR base ${base_branch}; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
|
||||
if ! files="$(git diff --name-only "${base_ref}...${BUILDKITE_COMMIT}" 2>/dev/null)"; then
|
||||
echo "resolve_skip_ci: failed to compute PR changed files; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
elif [[ "${BUILDKITE_BRANCH:-}" == "main" ]]; then
|
||||
if ! git rev-parse --verify "${BUILDKITE_COMMIT}^" >/dev/null 2>&1; then
|
||||
echo "resolve_skip_ci: commit has no parent on main; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
if ! files="$(git diff --name-only "${BUILDKITE_COMMIT}^..${BUILDKITE_COMMIT}" 2>/dev/null)"; then
|
||||
echo "resolve_skip_ci: failed to compute main changed files; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
else
|
||||
echo "resolve_skip_ci: not PR/main build; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
return 0
|
||||
fi
|
||||
|
||||
if is_docs_only_change <<< "${files}"; then
|
||||
echo "resolve_skip_ci: docs-only change detected; skip-ci=1" >&2
|
||||
echo -n 1
|
||||
return 0
|
||||
fi
|
||||
|
||||
echo "resolve_skip_ci: non-doc changes detected; skip-ci=0" >&2
|
||||
echo -n 0
|
||||
}
|
||||
|
||||
SKIP_CI="$(resolve_skip_ci)"
|
||||
|
||||
if [[ ! -f "${PIPELINE_YML}" ]]; then
|
||||
echo "upload_pipeline_with_skip_ci: missing ${PIPELINE_YML}" >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
export ROOT SKIP_CI PIPELINE_YML
|
||||
python3 <<'PY' | buildkite-agent pipeline upload
|
||||
import os
|
||||
import pathlib
|
||||
|
||||
path = pathlib.Path(os.environ["PIPELINE_YML"])
|
||||
text = path.read_text(encoding="utf-8")
|
||||
sep = "\n---\n"
|
||||
# Two supported layouts:
|
||||
# - multi-doc: doc 1 = bootstrap (loaded by Buildkite), doc 2 (after `---`) = uploaded steps.
|
||||
# - single-doc: whole file is the uploaded steps; caller (UI init step) is the bootstrap.
|
||||
if sep in text:
|
||||
_, continuation = text.split(sep, 1)
|
||||
else:
|
||||
continuation = text
|
||||
|
||||
skip = os.environ.get("SKIP_CI") == "1"
|
||||
# When docs-only skip-ci: skip default CI image, but still build for L4 nightly (PR label nightly-test or
|
||||
# main NIGHTLY=1), otherwise upload-nightly (depends_on image-build) would be skipped too.
|
||||
nightly_only = (
|
||||
'(build.pull_request.labels includes "nightly-test") '
|
||||
'|| (build.branch == "main" && build.env("NIGHTLY") == "1")'
|
||||
)
|
||||
# Placeholder in pipeline.yml is `if: __IMAGE_BUILD_IF__` (valid YAML); replace value only.
|
||||
if skip:
|
||||
rep = f"'{nightly_only}'"
|
||||
ready_rep = "'false'"
|
||||
merge_rep = "'false'"
|
||||
else:
|
||||
rep = "'true'"
|
||||
ready_rep = "'build.branch != \"main\" && build.pull_request.labels includes \"ready\"'"
|
||||
merge_rep = "'(build.branch == \"main\" && build.env(\"NIGHTLY\") != \"1\" && build.env(\"WEEKLY\") != \"1\") || (build.branch != \"main\" && build.pull_request.labels includes \"merge-test\")'"
|
||||
rendered = (
|
||||
continuation
|
||||
.replace("__IMAGE_BUILD_IF__", rep)
|
||||
.replace("__UPLOAD_READY_IF__", ready_rep)
|
||||
.replace("__UPLOAD_MERGE_IF__", merge_rep)
|
||||
)
|
||||
print(rendered, end="")
|
||||
PY
|
||||
+159
-148
@@ -1,54 +1,111 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
timeout_in_minutes: 50
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Diffusion Images API LoRA E2E"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
steps:
|
||||
- label: "Diffusion · LoRA Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/lora/test_images_generations_lora.py
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
"
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
@@ -62,42 +119,6 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test (Need 4 GPUs)"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_ulysses_uaa_perf.py
|
||||
|
||||
# merge-only tests
|
||||
- label: "Diffusion Tensor Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- pytest -s -v tests/diffusion/distributed/test_zimage_parallelism.py
|
||||
|
||||
- label: "Diffusion Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 30m pytest -s -v tests/e2e/offline_inference/test_z_image.py
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Engine Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
@@ -111,50 +132,65 @@ steps:
|
||||
# This model seems to be important, so I am trying to run it here.
|
||||
# Since this is from test-nightly.yml,
|
||||
# I keep this non-blocking.
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
steps:
|
||||
- label: "Omni Model Test Qwen2-5-Omni"
|
||||
agent_pool: mi325_4
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
commands:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- pytest -s -v tests/e2e/*/test_qwen2_5_omni_expansion.py --run-level "core_model"
|
||||
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
- label: "Omni Model Test Qwen3-Omni"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
- label: "Qwen3-TTS CustomVoice E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py tests/e2e/offline_inference/test_qwen3_tts_customvoice.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "Qwen3-TTS Base E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
- label: "Qwen3-TTS Base E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
@@ -195,28 +231,3 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "advanced_model" --run-level "advanced_model" -k "rocm"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/offline_inference/test_wan22_t2v.py tests/e2e/online_serving/test_wan22_t2v.py -m 'advanced_model and cuda and diffusion' --run-level 'advanced_model'
|
||||
"
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
+156
-144
@@ -1,42 +1,33 @@
|
||||
env:
|
||||
VLLM_WORKER_MULTIPROC_METHOD: spawn
|
||||
HF_HUB_DOWNLOAD_TIMEOUT: 300
|
||||
HF_HUB_ETAG_TIMEOUT: 60
|
||||
|
||||
steps:
|
||||
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
- group: ":card_index_dividers: Simple Test"
|
||||
steps:
|
||||
- label: "Simple · Diffusion & Model Executor Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
||||
|
||||
- label: "Diffusion · Batch Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
- label: "Simple · Other Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_ROCM_USE_AITER=0
|
||||
# ignore test_teacache_extractors.py because it use rocm gemm kernel from vLLM
|
||||
# that is not supported on CPU
|
||||
- "timeout 40m pytest -sv tests/diffusion tests/model_executor -m 'core_model and cpu' --ignore=tests/diffusion/cache/test_teacache_extractors.py --cov=vllm_omni --cov-report=term-missing:skip-covered --cov-report=xml"
|
||||
|
||||
- label: "Custom Pipeline Test"
|
||||
agent_pool: mi325_1
|
||||
@@ -46,23 +37,92 @@ steps:
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/e2e/offline_inference/custom_pipeline/ -m "core_model"
|
||||
|
||||
- label: "Diffusion Model CPU offloading Test"
|
||||
|
||||
- label: "Voxtral TTS CUDA Unit Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
- timeout 10m pytest -s -v tests/model_executor/models/voxtral_tts/test_cuda_graph_acoustic_transformer.py
|
||||
|
||||
- group: ":card_index_dividers: Diffusion Test"
|
||||
steps:
|
||||
- label: "Diffusion · Batch Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -sv tests/diffusion/batching -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Offloader Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
set +e
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_cpu_offload.py
|
||||
EXIT1=\$?
|
||||
pytest -s -v tests/diffusion/offloader/test_diffusion_layerwise_offload.py
|
||||
EXIT2=\$?
|
||||
exit \$((EXIT1 | EXIT2))
|
||||
'
|
||||
|
||||
- label: "Diffusion · Cache Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
|
||||
## ISSUE depends on `diffusers` package: https://github.com/huggingface/diffusers/issues/13274
|
||||
# - label: "Audio Generation Model Test"
|
||||
@@ -76,56 +136,6 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - timeout 20m pytest -s -v tests/e2e/offline_inference/test_stable_audio_expansion.py -m "advanced_model and diffusion and L4" --run-level advanced_model
|
||||
|
||||
- label: "Diffusion Cache Backend Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export GPU_ARCHS=gfx942
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 15m pytest -s -v -m "core_model and cache and diffusion and not distributed_cuda and L4"
|
||||
|
||||
- label: "Diffusion Sequence Parallelism Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- timeout 20m pytest -s -v tests/diffusion/distributed/test_sequence_parallel.py -m core_model
|
||||
|
||||
- label: "Diffusion GPU Worker Test"
|
||||
agent_pool: mi325_2
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 20m pytest -s -v tests/diffusion/test_diffusion_worker.py
|
||||
|
||||
- label: "Diffusion · Model Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- timeout 15m pytest -sv tests/diffusion/models/ -m 'core_model and cuda' --run-level "core_model"
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
export VLLM_TEST_CLEAN_GPU_MEMORY=1
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_wan22_t2v.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
|
||||
- label: "Engine Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
@@ -149,43 +159,56 @@ steps:
|
||||
# - timeout 10m pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
# - timeout 20m pytest -s -v tests/e2e/online_serving/test_qwen3_omni.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "MiMo-Audio E2E Test with H100"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
- group: ":card_index_dividers: E2E Test"
|
||||
steps:
|
||||
- label: "MiMo-Audio E2E Test with H100"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m "core_model" --run-level "core_model"
|
||||
'
|
||||
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
- label: "Qwen3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- export VLLM_LOGGING_LEVEL=DEBUG
|
||||
- export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
|
||||
'
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
# Need to debug before reneable numerical changes across large PRs
|
||||
@@ -225,14 +248,3 @@ steps:
|
||||
# - export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
# - export VLLM_ROCM_USE_AITER_RMSNORM=0
|
||||
# - timeout 40m pytest -s -v tests/e2e/online_serving/test_bagel.py -m "core_model" --run-level "core_model" -k "rocm"
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
pytest -s -v tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py -m "full_model" --run-level "core_model"
|
||||
'
|
||||
+78
-25
@@ -206,7 +206,34 @@ 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:
|
||||
- |
|
||||
@@ -249,6 +276,12 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · GLM-TTS Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_glm_tts.py
|
||||
- tests/e2e/offline_inference/test_glm_tts.py
|
||||
- vllm_omni/model_executor/models/glm_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/glm_tts.py
|
||||
- vllm_omni/deploy/glm_tts.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
@@ -291,6 +324,12 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
- tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- vllm_omni/model_executor/models/qwen3_omni/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_qwen3_omni.py tests/e2e/online_serving/test_qwen3_omni.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -329,6 +368,11 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Mimo Audio Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_mimo_audio.py
|
||||
- vllm_omni/model_executor/models/mimo_audio/
|
||||
- vllm_omni/model_executor/stage_input_processors/mimo_audio.py
|
||||
- vllm_omni/deploy/mimo_audio.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_mimo_audio.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -366,37 +410,14 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · VoxCPM2 Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · VoxCPM2 Online Test"
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
@@ -438,6 +459,9 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_qwen_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -476,6 +500,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Longcat Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_longcat_image_edit.py
|
||||
- vllm_omni/diffusion/models/longcat_image/
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_longcat_image_edit.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -514,6 +541,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Longcat Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_longcat_image.py
|
||||
- vllm_omni/diffusion/models/longcat_image/
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_longcat_image.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -552,6 +582,13 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Bagel Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_bagel.py
|
||||
- tests/e2e/online_serving/test_bagel.py
|
||||
- vllm_omni/model_executor/models/bagel/
|
||||
- vllm_omni/diffusion/models/bagel/
|
||||
- vllm_omni/model_executor/stage_input_processors/bagel.py
|
||||
- vllm_omni/deploy/bagel.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/offline_inference/test_bagel.py tests/e2e/online_serving/test_bagel.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
@@ -591,6 +628,9 @@ steps:
|
||||
|
||||
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c '
|
||||
@@ -632,6 +672,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Qwen Image Layered Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_layered.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -673,6 +716,10 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_wan22_t2v.py
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
- vllm_omni/diffusion/models/wan2_2/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -714,6 +761,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · HunyuanVideo 1.5 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_hunyuan_video_15.py
|
||||
- vllm_omni/diffusion/models/hunyuan_video/
|
||||
timeout_in_minutes: 40
|
||||
commands:
|
||||
- |
|
||||
@@ -789,6 +839,9 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Diffusion · SD3 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_sd3.py
|
||||
- vllm_omni/diffusion/models/sd3/
|
||||
timeout_in_minutes: 45
|
||||
commands:
|
||||
- |
|
||||
|
||||
@@ -568,6 +568,55 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: ":full_moon: TTS · Higgs-Audio-V3 · Perf Test"
|
||||
key: nightly-tts-performance-higgs-audio-v3
|
||||
depends_on: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test"
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export VLLM_USE_DEEP_GEMM="0"
|
||||
- export VLLM_MOE_USE_DEEP_GEMM="0"
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_higgs_audio_v3.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/*.json"
|
||||
exit $$EXIT
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
# Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below.
|
||||
- group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test"
|
||||
key: nightly-diffusion-x2iat-group
|
||||
@@ -591,6 +640,7 @@ steps:
|
||||
tests/e2e/online_serving/test_longcat_image_edit_expansion.py
|
||||
tests/e2e/online_serving/test_flux_2_dev_expansion.py
|
||||
tests/e2e/online_serving/test_bagel_expansion.py
|
||||
tests/e2e/offline_inference/test_sensenova_u1_expansion.py
|
||||
-m "full_model and diffusion and H100 and not distributed_cuda"
|
||||
--run-level "full_model"
|
||||
agents:
|
||||
@@ -1462,6 +1512,7 @@ steps:
|
||||
- nightly-omni-performance-vllm-text
|
||||
- nightly-omni-performance-multi-replicas
|
||||
- nightly-tts-performance
|
||||
- nightly-tts-performance-higgs-audio-v3
|
||||
- nightly-diffusion-x2iat-performance-qwen-image
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit
|
||||
- nightly-diffusion-x2iat-performance-qwen-image-edit-2511
|
||||
@@ -1475,6 +1526,7 @@ steps:
|
||||
- export DEFAULT_INPUT_DIR=tests/dfx/perf/results
|
||||
- export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance-higgs-audio-v3
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-no-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-async-chunk
|
||||
- buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance-vllm-text
|
||||
|
||||
+85
-61
@@ -215,6 +215,43 @@ 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:
|
||||
@@ -239,6 +276,11 @@ steps:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_omni.py
|
||||
- vllm_omni/model_executor/models/qwen3_omni/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_omni.py
|
||||
- vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
@@ -273,73 +315,14 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · Qwen3-Omni Realtime WebSocket"
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py::TestQwen3OmniRealtimeWebSocket::test_streaming_audio_input_pcm_output_async_chunk_smoke -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 2
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Omni · VoxCPM2 Test"
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · VoxCPM2 Online Test"
|
||||
- label: "TTS · VoxCPM2 Test"
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model and tts' --run-level 'core_model'
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxcpm2_tts.py tests/e2e/online_serving/test_voxcpm2_tts.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_1_queue"
|
||||
@@ -399,7 +382,33 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · Higgs-Audio-V3 Test"
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
export VLLM_USE_DEEP_GEMM=0
|
||||
export VLLM_MOE_USE_DEEP_GEMM=0
|
||||
pytest -s -v tests/e2e/online_serving/test_higgs_audio_v3.py -m 'core_model' --run-level 'core_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
volumes:
|
||||
- "/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 "
|
||||
@@ -435,6 +444,12 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Bagel Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_bagel.py
|
||||
- vllm_omni/model_executor/models/bagel/
|
||||
- vllm_omni/diffusion/models/bagel/
|
||||
- vllm_omni/model_executor/stage_input_processors/bagel.py
|
||||
- vllm_omni/deploy/bagel.yaml
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -476,6 +491,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Wan22 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_wan22_t2v.py
|
||||
- vllm_omni/diffusion/models/wan2_2/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -518,6 +536,9 @@ steps:
|
||||
|
||||
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
@@ -559,6 +580,9 @@ steps:
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Diffusion · Qwen Image Layered Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_layered.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
commands:
|
||||
- |
|
||||
timeout 40m bash -c "
|
||||
|
||||
@@ -6,6 +6,13 @@
|
||||
{% set docker_image_amd = "rocm/vllm-omni:$BUILDKITE_COMMIT" %}
|
||||
{% set default_working_dir = "/app/vllm-omni" %}
|
||||
|
||||
{% if env %}
|
||||
env:
|
||||
{% for name in env %}
|
||||
{{ name }}: "{{ env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
steps:
|
||||
- group: "AMD Tests"
|
||||
depends_on: ~
|
||||
steps:
|
||||
@@ -31,8 +38,15 @@
|
||||
agents:
|
||||
queue: amd-cpu
|
||||
|
||||
{% for step in steps %}
|
||||
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
|
||||
{% for step_group in steps %}
|
||||
{% if step_group.group %}
|
||||
- group: "{{ step_group.group }}"
|
||||
{% if step_group.depends_on %}
|
||||
depends_on: {{ step_group.depends_on }}
|
||||
{% endif %}
|
||||
steps:
|
||||
{% for step in step_group.steps %}
|
||||
{% if step.mirror_hardwares and mirror_hw in step.mirror_hardwares %}
|
||||
- label: "{{ step.agent_pool }}: {{ step.label }}"
|
||||
depends_on: amd-build
|
||||
agents:
|
||||
@@ -46,8 +60,23 @@
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
{{ name }}: "{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
TEST_COMMAND: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{% if env %}
|
||||
{% for name in env %}
|
||||
export {{ name }}="{{ env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
export {{ name }}="{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if "mi250" in step.agent_pool %}
|
||||
python3 -m pip uninstall -y amd-aiter
|
||||
{% endif %}
|
||||
@@ -61,5 +90,52 @@
|
||||
{% else %}
|
||||
soft_fail: true
|
||||
{% endif%}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
{% elif step_group.mirror_hardwares and mirror_hw in step_group.mirror_hardwares %}
|
||||
{% set step = step_group %}
|
||||
- label: "{{ step.agent_pool }}: {{ step.label }}"
|
||||
depends_on: amd-build
|
||||
agents:
|
||||
{% if step.agent_pool %}
|
||||
queue: amd_{{ step.agent_pool }}
|
||||
{% else %}
|
||||
queue: amd_mi325_1
|
||||
{% endif %}
|
||||
{% set cmd_body = (step.command or (step.commands | join("\n"))) | trim %}
|
||||
{% set indented_cmd = cmd_body | replace("\n", "\n ") %}
|
||||
command: bash .buildkite/scripts/hardware_ci/run-amd-test.sh
|
||||
env:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
{{ name }}: "{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
TEST_COMMAND: |-
|
||||
(command rocm-smi || true) && cd {{ (step.working_dir or default_working_dir) | safe }}
|
||||
{% if env %}
|
||||
{% for name in env %}
|
||||
export {{ name }}="{{ env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if step.env %}
|
||||
{% for name in step.env %}
|
||||
export {{ name }}="{{ step.env[name] }}"
|
||||
{% endfor %}
|
||||
{% endif %}
|
||||
{% if "mi250" in step.agent_pool %}
|
||||
python3 -m pip uninstall -y amd-aiter
|
||||
{% endif %}
|
||||
{{ indented_cmd | safe }}
|
||||
priority: 100
|
||||
{% if step.timeout_in_minutes %}
|
||||
timeout_in_minutes: {{ step.timeout_in_minutes }}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
{% if step.grade and step.grade == "Blocking" %}
|
||||
soft_fail: false
|
||||
{% else %}
|
||||
soft_fail: true
|
||||
{% endif%}
|
||||
{% endif %}
|
||||
{% endfor %}
|
||||
@@ -0,0 +1,97 @@
|
||||
---
|
||||
name: precheck-pr
|
||||
description: Self-check your branch before creating a PR — catch dead code, verify accuracy/perf claims, validate PR title format, and confirm merge readiness. Use when the user says "precheck", "self review", "pre-submit check", or "check my PR before I open it." Never posts to GitHub.
|
||||
---
|
||||
|
||||
# PR Pre-Check
|
||||
|
||||
Self-review your branch before creating a PR against `vllm-project/vllm-omni`. Two modes: **quick** catches showstoppers, **full** does a thorough maintainer-grade review. Never posts to GitHub; the report is for the contributor's terminal only.
|
||||
|
||||
## Mode Selection
|
||||
|
||||
| Mode | When | Time |
|
||||
|------|------|------|
|
||||
| **Quick** | About to push, final sanity check | ~3 min |
|
||||
| **Full** | Ready for review, want maintainer-level scan | ~10 min |
|
||||
|
||||
Default to quick if unsure. Run full before marking a PR "ready for review."
|
||||
|
||||
## Workflow
|
||||
|
||||
### Step 1: Detect Base Branch
|
||||
|
||||
```bash
|
||||
BASE_SHA=$(git merge-base HEAD origin/main 2>/dev/null \
|
||||
|| git merge-base HEAD main 2>/dev/null \
|
||||
|| echo origin/main)
|
||||
echo "diffing against ${BASE_SHA}"
|
||||
git diff --name-only ${BASE_SHA}...HEAD
|
||||
```
|
||||
|
||||
### Step 2: Validate PR Title
|
||||
|
||||
Check the most recent commit message (or branch name if no commit yet) against the project convention. Valid prefixes:
|
||||
|
||||
| Prefix | Applies to |
|
||||
|--------|-----------|
|
||||
| `[Bugfix]` | Bug fixes |
|
||||
| `[CI/Build]` | Build or CI improvements |
|
||||
| `[Doc]` | Documentation changes |
|
||||
| `[Model]` | New/improved models (include model name) |
|
||||
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
|
||||
| `[Kernel]` | CUDA/kernel changes |
|
||||
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
|
||||
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
|
||||
| `[Misc]` | Other changes (use sparingly) |
|
||||
|
||||
✗ if: missing prefix, wrong case (`[bugfix]`), or WIP/Draft in title.
|
||||
⚠ if: `[Model]` prefix without the model identifier (e.g., `[Model] Add new model` — should be `[Model] Add <ModelName> ...`).
|
||||
|
||||
### Step 3: Categorize the PR
|
||||
|
||||
| Diff contains | PR type |
|
||||
|---------------|---------|
|
||||
| New files under `vllm_omni/model_executor/models/<name>/` | **New Model** |
|
||||
| Changes to `vllm_omni/diffusion/` | **Diffusion Model** |
|
||||
| `[Bugfix]` prefix or single-file fix | **Bug Fix** |
|
||||
| Perf/benchmark/throughput claims in commit msg or diff | **Performance** |
|
||||
| Everything else | **General** |
|
||||
|
||||
If multiple rows apply (e.g., a diffusion model is also a new model), union the checklists.
|
||||
|
||||
### Step 4: Run Checklist
|
||||
|
||||
Ask: "Quick mode or full mode?" Then walk the checklist for the detected PR type from [references/checklists.md](references/checklists.md). Each item produces ✓, ✗, or ⚠.
|
||||
|
||||
### Step 5: Print Report
|
||||
|
||||
```
|
||||
Pre-check report for <branch>
|
||||
|
||||
Mode: quick | full
|
||||
Type: <new-model | diffusion-model | bug-fix | perf | general>
|
||||
|
||||
Dimension Result
|
||||
───────────────── ──────
|
||||
PR title format ✓
|
||||
PR desc integrity ✓
|
||||
Registry/config ✓
|
||||
Dead code ⚠ 2 warnings
|
||||
Accuracy ✓
|
||||
Benchmark ✗ missing software versions
|
||||
|
||||
Verdict: 1 blocking | 2 warnings | recommend fixing ✗ before PR
|
||||
```
|
||||
|
||||
**Severity:**
|
||||
|
||||
| Mark | Meaning |
|
||||
|------|---------|
|
||||
| ✗ | Blocking — fix before opening PR |
|
||||
| ⚠ | Warning — consider fixing |
|
||||
| ✓ | Pass |
|
||||
| — | Skipped (not applicable) |
|
||||
|
||||
### Stop Here
|
||||
|
||||
Do not post comments, open PRs, or modify files. The report is for the contributor's terminal only.
|
||||
@@ -0,0 +1,135 @@
|
||||
# Pre-Check Checklists
|
||||
|
||||
Each item is a concrete check to run against the diff. ✗ = fix before PR. ⚠ = consider fixing.
|
||||
|
||||
---
|
||||
|
||||
## All PRs: Title Format (check first)
|
||||
|
||||
PR title must follow the project convention documented in `docs/contributing/README.md`:
|
||||
|
||||
| Prefix | Applies to |
|
||||
|--------|-----------|
|
||||
| `[Bugfix]` | Bug fixes |
|
||||
| `[CI/Build]` | Build or CI improvements |
|
||||
| `[Doc]` | Documentation changes |
|
||||
| `[Model]` | New/improved models (include model name) |
|
||||
| `[Frontend]` | Frontend changes (API server, OmniLLM class, etc.) |
|
||||
| `[Kernel]` | CUDA/kernel changes |
|
||||
| `[Core]` | Core logic changes (OmniProcessor, OmniARScheduler, etc.) |
|
||||
| `[Hardware][Vendor]` | Hardware-specific (e.g., `[Hardware][Ascend]`) |
|
||||
| `[Misc]` | Other changes (use sparingly) |
|
||||
|
||||
**Examples:**
|
||||
```
|
||||
✓ [Model] Add MiniCPM-o 4.5 support
|
||||
✓ [Bugfix] Fix Ovis image text encoder dtype
|
||||
✓ [Core] Qwen3-Omni performance optimization
|
||||
✓ [Hardware][Ascend] Fix VoxCPM2 on Ascend
|
||||
✗ Fix Ovis image text encoder dtype ← missing prefix
|
||||
✗ [bugfix] Fix Ovis image text encoder ← wrong case
|
||||
✗ [WIP] Add model ← WIP in title
|
||||
```
|
||||
|
||||
- [ ] **Title has valid prefix** — ✗ if missing prefix, wrong case, or WIP/Draft in title. For `[Model]` PRs, the title MUST also include the model identifier (e.g., `[Model] Add <ModelName> ...`). ⚠ if missing the model name.
|
||||
|
||||
---
|
||||
|
||||
## Bug Fix PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body has 4 sections:** what broke (error message), repro steps, root cause, fix description
|
||||
- [ ] **Fix is minimal:** diff touches only files directly related to the bug. No unrelated refactors or style changes.
|
||||
- [ ] **Repro is runnable:** the PR body includes a copy-paste command that reproduces the bug
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Regression test exists:** `git diff --name-only` includes at least one `tests/` file
|
||||
- [ ] **Regression test reproduces the original error:** the test would fail on main, pass on this branch
|
||||
- [ ] **Fix matches root cause exactly:** no "fix the symptom + something else" — if the root cause is dtype, the fix is dtype, not dtype + reformatting
|
||||
- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added
|
||||
- [ ] **Upstream pattern match:** the fix follows the same pattern as existing code for similar cases (grep for analogous `from_pretrained` calls, etc.)
|
||||
- [ ] **Environment documented:** torch/transformers/vllm versions listed if the bug is version-dependent
|
||||
|
||||
---
|
||||
|
||||
## Performance PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **Before/after numbers in PR body:** at minimum, one metric (RTF, RPS, latency, VRAM) with main vs PR values
|
||||
- [ ] **Hardware specified:** GPU model, count, VRAM. "L4" or "H20" is sufficient. "tested on GPU" is not.
|
||||
- [ ] **Concurrency scaling:** if the optimization claims better throughput, numbers at 2+ concurrency levels
|
||||
- [ ] **No unexplained regressions:** if any metric goes down at any concurrency level, explain why
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Software versions:** torch, CUDA, vllm, vllm-omni versions in PR body
|
||||
- [ ] **Warmup stated:** "1 warmup + 3 measured" or equivalent
|
||||
- [ ] **Stddev or range:** not just a single best run. "average of 2 rounds" is minimum.
|
||||
- [ ] **VRAM measured:** peak GPU memory before and after
|
||||
- [ ] **Benchmark script checked in:** under `tests/` or `examples/`, runnable by reviewer
|
||||
- [ ] **Exact command line in PR:** not pseudocode, not "run the benchmark"
|
||||
- [ ] **Pytest summary included:** `36 passed, 3 skipped, 17 warnings in 2897.66s` — proves the whole suite passes
|
||||
- [ ] **Config changes documented:** if `max_num_seqs`, `gpu_memory_utilization`, or deploy YAML changed, old and new values stated with reasoning
|
||||
|
||||
---
|
||||
|
||||
## New Model PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body file list matches `git diff --name-only`:** every file path in the description exists in the diff
|
||||
- [ ] **Registry entries match claimed count:** count `_OMNI_MODELS` or `_DIFFUSION_MODELS` entries added vs claimed in body
|
||||
- [ ] **Pipeline config `custom_process_input_func` paths resolve:** each referenced module.function exists
|
||||
- [ ] **`__init__.py` exports match registry:** each registered class appears in `__init__.py` and `__all__`
|
||||
- [ ] **Shell script modes match README:** `run_curl.sh` supports every mode the README documents
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **Dead code scan (5 patterns):**
|
||||
1. Dead `forward()` — defined but never called through the pipeline path
|
||||
2. Dead factory/builder functions — instantiated nowhere
|
||||
3. Dead wrappers — thin wrappers with no callers
|
||||
4. Dead branch guards — `if version > X` where X is always satisfied
|
||||
5. Unused parameters — `__init__` args never accessed
|
||||
- [ ] **Copy-paste detection:** no string constant defined in 3+ files, no cross-module validation duplication, no near-identical shape coercion fns
|
||||
- [ ] **Import hygiene:** no re-export-only imports, no module-level side effects without docs, no `import os` in function body when already at top. Function-body imports are acceptable when the import is only reachable from a closure or factory return value (i.e., `import` would otherwise execute at module load even though the user may never trigger that path). Either lift to top-of-file or annotate with `# noqa: PLC0415`.
|
||||
- [ ] **Accuracy:** at least one test compares output against reference; audio = valid WAV at correct sample rate; determinism verified if seed claimed. ⚠ if missing and not explicitly deferred to a follow-up PR or RFC section; — if explicitly deferred (the deferral target SHOULD be linked in the PR body).
|
||||
- [ ] **Performance:** hardware, software versions, warmup, and metrics (RTF, VRAM, RPS, latency) stated. ⚠ if missing and not explicitly deferred; — if explicitly deferred.
|
||||
- [ ] **Benchmark settings:** model config (TP, PP, max_model_len, enforce_eager, quant), runtime config (batch, input/output spec), environment versions
|
||||
- [ ] **Benchmark script checked in + exact command line + pytest summary**
|
||||
|
||||
---
|
||||
|
||||
## Diffusion Model PRs
|
||||
|
||||
Diffusion models live under `vllm_omni/diffusion/`. In addition to the New Model checklist above, verify:
|
||||
|
||||
- [ ] **Transformer adapter follows contract:** inherits from `nn.Module`, implements `load_weights()`, uses `vllm_omni.diffusion.attention.layer.Attention`
|
||||
- [ ] **Pipeline has both offline and online paths:** `forward()` for batch generation, streaming path for online serving
|
||||
- [ ] **Cache-DiT integration:** at least one acceleration (TeaCache, FBCache, etc.) is wired up. ✗ for new **top-level** pipelines. For **subclass** pipelines that inherit `__call__` from a parent where Cache-DiT is already absent, downgrade to ⚠ and recommend filing a follow-up issue against the parent pipeline family.
|
||||
- [ ] **Parallelism config:** TP/SP/USP/CFG-Parallel options exposed in pipeline config
|
||||
- [ ] **Docs table updated:** model added to `docs/models/supported_models.md` with correct metadata
|
||||
- [ ] **E2E test exists:** at least one test exercises the full generate path
|
||||
|
||||
---
|
||||
|
||||
## General PRs
|
||||
|
||||
### Quick
|
||||
|
||||
- [ ] **PR body matches diff:** description doesn't claim changes not in the diff
|
||||
- [ ] **Branch is rebased:** `git merge-base HEAD origin/main` is recent, no merge conflicts
|
||||
- [ ] **CI gates passing:** DCO, pre-commit, build — check `gh pr view --json statusCheckRollup` if a PR already exists
|
||||
|
||||
### Full (adds)
|
||||
|
||||
- [ ] **No new dead code in changed files:** scan for unreachable paths, unused imports
|
||||
- [ ] **Test coverage:** if core code changed (`engine/`, `stages/`, `connectors/`), tests are added or existing tests cover the path
|
||||
- [ ] **Import hygiene:** no redundant imports added
|
||||
- [ ] **No unrelated changes:** diff scope matches PR description
|
||||
@@ -0,0 +1,103 @@
|
||||
---
|
||||
name: quantization
|
||||
description: Work on vLLM-Omni quantization for diffusion, autoregressive, omni, or multi-stage models. Use when choosing or adding methods such as fp8, int8, gguf, mxfp8, mxfp4, mxfp4_dualscale, ModelOpt, AutoRound, INC, msModelSlim, awq, or gptq; debugging quantized loading; or validating memory, speed, and output quality.
|
||||
---
|
||||
|
||||
# vLLM-Omni Quantization
|
||||
|
||||
Use this skill for quantization work in `vllm-omni`. Start from the local
|
||||
code and docs, then pick the smallest validated path for the model, method,
|
||||
hardware, and quality target.
|
||||
|
||||
## First Checks
|
||||
|
||||
Before changing code or recommending a command, identify:
|
||||
|
||||
- Model and task: AR text, omni/TTS, text-to-image, text-to-video, or multi-stage diffusion.
|
||||
- Hardware backend: CUDA generation, Ascend NPU, Intel XPU, or another platform.
|
||||
- Quantization mode: online from BF16/FP16 weights, pre-quantized checkpoint, runtime KV/FA quantization, or per-component routing.
|
||||
- Scope: transformer/DiT, thinker language model, text encoder, VAE, audio/vision encoder, talker, or a specific stage.
|
||||
- Validation target: memory reduction, latency/throughput, output quality, or all three.
|
||||
|
||||
Do not assume a method is supported just because the CLI accepts the string.
|
||||
Check `docs/user_guide/quantization/` and the closest implementation first.
|
||||
|
||||
## Quick Decision
|
||||
|
||||
| Task | Start With |
|
||||
|------|------------|
|
||||
| Choose a method or command | `references/methods.md` and `references/modality-compat.md` |
|
||||
| Use `build_quant_config()` or per-component routing | `references/methods.md` |
|
||||
| Work on diffusion quantization | `references/diffusion.md` |
|
||||
| Add quantization to a new model | `references/adding-models.md` |
|
||||
| Convert or load ModelOpt FP8 checkpoints | `references/modelopt-fp8.md` |
|
||||
| Debug ModelOpt, GGUF, AutoRound, MXFP, or serialized Int8 loading | `references/diffusion.md` and method docs under `docs/user_guide/quantization/` |
|
||||
|
||||
## Current Runtime Shape
|
||||
|
||||
The unified entrypoint is:
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
```
|
||||
|
||||
It supports method strings, flat method dictionaries, per-component
|
||||
dictionaries, existing `QuantizationConfig` objects, and `None`. The factory
|
||||
delegates generic methods to upstream `vllm` and keeps vLLM-Omni overrides for
|
||||
diffusion or omni-specific routing:
|
||||
|
||||
- `gguf`
|
||||
- `int8`
|
||||
- `mxfp8`
|
||||
- `mxfp4`
|
||||
- `mxfp4_dualscale`
|
||||
- `inc`, `auto-round`, `auto_round`
|
||||
- ModelOpt auto-detected configs: `modelopt`, `modelopt_fp4`, `modelopt_mixed`
|
||||
|
||||
Use `ComponentQuantizationConfig` when only one stage or component should be
|
||||
quantized. Pre-quantized ModelOpt-style checkpoints should not spill into
|
||||
vision/audio encoders that have no corresponding scale tensors.
|
||||
|
||||
## Ownership Boundary
|
||||
|
||||
- Upstream `vllm` owns generic quantization configs, kernels, loader semantics,
|
||||
hardware capability rules, and generic AR quantization methods.
|
||||
- `vllm-omni` owns unified config routing, diffusion-specific wrappers,
|
||||
component scoping, GGUF or ModelOpt checkpoint adapters, model-specific
|
||||
prefix mapping, docs, examples, and validation.
|
||||
|
||||
If a new method needs missing generic kernels or loader behavior, fix upstream
|
||||
`vllm` first. In `vllm-omni`, add thin integration and model wiring.
|
||||
|
||||
## Working Loop
|
||||
|
||||
1. Read the local method doc under `docs/user_guide/quantization/`.
|
||||
2. Inspect the active factory and config class in `vllm_omni/quantization/`.
|
||||
3. Inspect the target model pipeline and transformer for stable prefixes and
|
||||
whether every relevant vLLM linear layer receives `quant_config`.
|
||||
4. For pre-quantized checkpoints, inspect `transformer/config.json` and the
|
||||
checkpoint tensor names before running full generation.
|
||||
5. Validate with the same prompt, seed, size, scheduler, steps, dtype, eager or
|
||||
non-eager mode, and parallelism as the BF16 baseline.
|
||||
6. Report quality, latency, and memory separately. Do not call a method
|
||||
supported until quality is checked.
|
||||
|
||||
## Common Mistakes
|
||||
|
||||
| Symptom | Likely Cause | Fix |
|
||||
|---------|--------------|-----|
|
||||
| `--quantization` has no visible effect | Wrong scope or unsupported model path | Check component routing and method docs |
|
||||
| Some layers stay BF16 unexpectedly | `quant_config` was not threaded into all vLLM linear layers | Audit transformer constructors and prefixes |
|
||||
| Quality collapses but loading succeeds | Too many sensitive layers were quantized | Add model-specific `ignored_layers` and compare to BF16 |
|
||||
| ModelOpt checkpoint loads but output is corrupted | Prefixes, packed-module mapping, scale routing, or BF16 fallback is wrong | Use `references/modelopt-fp8.md` |
|
||||
| GGUF shape or tensor mismatch | Missing architecture-specific adapter | Add explicit adapter mapping; avoid generic fallback |
|
||||
| Online and offline MXFP4 disagree | Wrong `mxfp4` vs `mxfp4_dualscale` mode | Check `docs/user_guide/quantization/mxfp4.md` |
|
||||
| Quantized path is slower than BF16 | Kernel path, compile mode, dtype, or shape mismatch | Compare eager-to-eager and non-eager-to-non-eager |
|
||||
|
||||
## References
|
||||
|
||||
- General methods and config routing: [references/methods.md](references/methods.md)
|
||||
- Modality and model scope: [references/modality-compat.md](references/modality-compat.md)
|
||||
- Diffusion implementation patterns: [references/diffusion.md](references/diffusion.md)
|
||||
- Adding model support: [references/adding-models.md](references/adding-models.md)
|
||||
- ModelOpt FP8 conversion and loading: [references/modelopt-fp8.md](references/modelopt-fp8.md)
|
||||
@@ -0,0 +1,70 @@
|
||||
# Adding Quantization Support to a Model
|
||||
|
||||
Use this checklist when extending quantization to a new architecture or stage.
|
||||
|
||||
## Classify Ownership
|
||||
|
||||
| Case | Owner |
|
||||
|------|-------|
|
||||
| New generic AR method, kernel, or loader semantics | upstream `vllm` first |
|
||||
| Existing upstream method needs diffusion wiring | `vllm-omni` model and config integration |
|
||||
| GGUF or ModelOpt tensor mapping for a diffusion architecture | `vllm-omni` adapter |
|
||||
| NPU MXFP or msModelSlim diffusion path | `vllm-omni` platform/method wrapper and docs |
|
||||
| Multi-stage omni checkpoint quantizes only the thinker LM | `vllm-omni` component scoping |
|
||||
|
||||
Avoid building a private quantization stack in model code.
|
||||
|
||||
## Model Wiring Checklist
|
||||
|
||||
1. Find the closest supported model implementation.
|
||||
2. Confirm the model uses vLLM linear layers that accept `quant_config`.
|
||||
3. Add `quant_config` to the transformer constructor.
|
||||
4. Pass stable `prefix` values through every linear layer.
|
||||
5. Handle fused modules and packed mappings explicitly.
|
||||
6. Ensure `load_weights()` and any `WeightsMapper` remap `ignored_layers` or
|
||||
`modules_to_not_convert`.
|
||||
7. Keep tokenizer, scheduler, VAE, text encoder, audio/vision encoders, and
|
||||
talker BF16 unless method-specific docs validate them.
|
||||
8. Add or update a method doc under `docs/user_guide/quantization/`.
|
||||
|
||||
## Checkpoint Adapter Checklist
|
||||
|
||||
For pre-quantized checkpoints:
|
||||
|
||||
- Inspect `transformer/config.json` first.
|
||||
- Check `quant_method`, `method`, `producer.name`, `quant_algo`, serialized
|
||||
flags, `ignored_layers`, and any method-specific fields.
|
||||
- Verify tensor name mapping before generation.
|
||||
- Verify packed modules such as `to_qkv`, `add_kv_proj`, and `w13`.
|
||||
- Confirm scale tensor names and shapes.
|
||||
- Allow dequantization back to BF16 for intentionally skipped target layers.
|
||||
|
||||
## Sensitive Layer Discovery
|
||||
|
||||
Start conservative. Quantize the backbone first and keep output-side or routing
|
||||
layers BF16 until quality is proven.
|
||||
|
||||
Frequently sensitive areas:
|
||||
|
||||
- attention output projections such as `to_out`
|
||||
- FFN down projections such as `w2`
|
||||
- final projections such as `proj_out`, `final_layer`, and output heads
|
||||
- modulation, timestep, position, image, or text embedders
|
||||
- refiner blocks
|
||||
- MoE routers, experts selected by routing, and layernorm-adjacent modules
|
||||
- VAE, text/vision towers, audio encoders, and codec/talker stages
|
||||
|
||||
Do not reuse one model family's ignore list blindly.
|
||||
|
||||
## Tests and Evidence
|
||||
|
||||
Minimum evidence for a PR:
|
||||
|
||||
- loader/config test or smoke command
|
||||
- fixed-seed BF16 vs quantized output comparison
|
||||
- quality metrics and saved outputs for precision-changing work
|
||||
- memory and latency comparison under identical execution mode
|
||||
- failure behavior for unsupported models or missing checkpoint fields
|
||||
|
||||
When the PR adds a method to a support table, include the exact model,
|
||||
hardware, command, quality threshold, and artifact path used for validation.
|
||||
@@ -0,0 +1,113 @@
|
||||
# Diffusion Quantization Reference
|
||||
|
||||
Use this file for DiT quantization, config routing, checkpoint adapters, and
|
||||
quality or loader debugging.
|
||||
|
||||
## Supported Method Families
|
||||
|
||||
The current local diffusion quantization families are:
|
||||
|
||||
- CUDA/GPU online or checkpoint: `fp8`, `int8`, `gguf`, ModelOpt
|
||||
- Ascend NPU: `int8`, `mxfp8`, `mxfp4`, `mxfp4_dualscale`, msModelSlim paths
|
||||
- Intel XPU: AutoRound/INC and online MXFP8 paths where documented
|
||||
|
||||
Always cross-check the method doc before claiming model support. The CLI
|
||||
accepting a method string is not enough.
|
||||
|
||||
## Implementation Starting Points
|
||||
|
||||
| Area | Files |
|
||||
|------|-------|
|
||||
| Unified factory | `vllm_omni/quantization/factory.py` |
|
||||
| Per-component routing | `vllm_omni/quantization/component_config.py` |
|
||||
| GGUF diffusion method | `vllm_omni/quantization/gguf_config.py` |
|
||||
| Int8 diffusion method | `vllm_omni/quantization/int8_config.py` |
|
||||
| MXFP8 diffusion method | `vllm_omni/quantization/mxfp8_config.py` |
|
||||
| MXFP4 and dual-scale method | `vllm_omni/quantization/mxfp4_config.py` |
|
||||
| AutoRound/INC omni extension | `vllm_omni/quantization/inc_config.py` |
|
||||
| ModelOpt checkpoint adapters | `vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py` |
|
||||
| FP8 scaled-mm runtime patches | `vllm_omni/patch.py` |
|
||||
| Similarity tool | `vllm_omni/quantization/tools/compare_diffusion_trajectory_similarity.py` |
|
||||
|
||||
## Threading `quant_config`
|
||||
|
||||
For a diffusion transformer:
|
||||
|
||||
1. Accept `quant_config: QuantizationConfig | None` in the transformer.
|
||||
2. Pass `quant_config` and stable `prefix` values to every vLLM linear layer.
|
||||
3. Preserve packed-module mappings for fused QKV, FFN, or model-specific fused
|
||||
projections.
|
||||
4. Call `apply_vllm_mapper()` on configs that need HF-to-vLLM name remapping.
|
||||
5. For component configs, keep prefixes specific enough to avoid quantizing VAE
|
||||
or encoders by accident.
|
||||
|
||||
Common skipped areas are `proj_out`, output projections, modulation layers,
|
||||
embedders, refiner blocks, MoE routers, and layernorm-adjacent paths.
|
||||
|
||||
## Online vs Offline
|
||||
|
||||
Online quantization:
|
||||
|
||||
- starts from BF16/FP16 checkpoint weights
|
||||
- computes or stores quantized weights at load time
|
||||
- typically uses `--quantization <method>` or `quantization="<method>"`
|
||||
|
||||
Offline or pre-quantized checkpoints:
|
||||
|
||||
- include quantized weights and scales on disk
|
||||
- usually carry `quantization_config` in `transformer/config.json`
|
||||
- should auto-detect when the config is complete
|
||||
- may require model-specific adapter logic for tensor names and packed modules
|
||||
|
||||
Do not mix online flags with incompatible offline checkpoint configs. If the
|
||||
active config and disk config differ, `resolve_quant_config_from_disk()` should
|
||||
raise or rebuild rather than silently loading corrupted weights.
|
||||
|
||||
## GGUF Rules
|
||||
|
||||
- Require `quantization_config.gguf_model`.
|
||||
- Use the normal base model for tokenizer, scheduler, text encoder, and VAE.
|
||||
- Add explicit architecture adapters for tensor-name mapping.
|
||||
- Use `named_parameters()` and `named_buffers()`, not `state_dict()`, to
|
||||
discover loadable names.
|
||||
- Guard fused QKV/KV rewrites so tensors are not rewritten twice.
|
||||
- Flatten N-D activations around the GGUF matmul and restore shape.
|
||||
|
||||
Unsupported architectures should fail clearly instead of falling back to a fake
|
||||
generic mapper.
|
||||
|
||||
## MXFP Rules
|
||||
|
||||
- `mxfp8` is W8A8 microscaling, currently centered on NPU and documented XPU
|
||||
behavior.
|
||||
- `mxfp4` is online single-scale W4A4 on NPU.
|
||||
- `mxfp4_dualscale` is online/offline dual-scale W4A4 with BF16 fallback; the
|
||||
offline path is the production path when calibrated `mul_scale` exists.
|
||||
- Offline dual-scale checkpoints must inject per-transformer `ignored_layers`
|
||||
into config because cascade transformers can have different BF16 fallback
|
||||
sets.
|
||||
|
||||
Do not load an offline dual-scale checkpoint with the online single-scale path.
|
||||
|
||||
## Common Failures
|
||||
|
||||
| Symptom | Likely Cause | Fix |
|
||||
|---------|--------------|-----|
|
||||
| `Unknown quantization method` | Factory override or upstream registration is missing | Add the right override or upstream method |
|
||||
| Missing scale tensors | Wrong scope, wrong adapter, or BF16 layer not allowed to dequantize | Check config scope and adapter scale routing |
|
||||
| Output rank collapses from 3-D to 2-D | FP8 scaled-mm kernel ignored output shape | Check the patch in `vllm_omni/patch.py` |
|
||||
| Quality degrades silently | Sensitive layers quantized | Add `ignored_layers`; compare to BF16 |
|
||||
| Only one cascade transformer works | Disk config differs per transformer but active config is reused | Rebuild from each transformer's config |
|
||||
| Quantized path slower than BF16 | Kernel, dtype, compile, or benchmark mode mismatch | Compare eager-to-eager and non-eager-to-non-eager |
|
||||
|
||||
## Verification
|
||||
|
||||
For a new or changed diffusion quantization path, run:
|
||||
|
||||
1. Config and loader smoke test.
|
||||
2. One fixed-seed generation against BF16.
|
||||
3. Numeric comparison with the similarity tool or equivalent metrics.
|
||||
4. Serving test if the path is user-facing online.
|
||||
5. Memory and latency check under the same execution mode as the baseline.
|
||||
|
||||
Keep output artifacts when precision changes are part of the PR.
|
||||
@@ -0,0 +1,82 @@
|
||||
# General Quantization Methods
|
||||
|
||||
Use this file for method selection and unified `vllm-omni` config syntax. For
|
||||
diffusion-specific wiring, see `diffusion.md`. For ModelOpt checkpoint export
|
||||
and loading details, see `modelopt-fp8.md`.
|
||||
|
||||
## Unified Entry Point
|
||||
|
||||
`vllm_omni.quantization.build_quant_config()` accepts:
|
||||
|
||||
- `None` or `"none"` to disable quantization
|
||||
- a method string such as `"fp8"`, `"int8"`, `"mxfp8"`, `"awq"`, or `"auto-round"`
|
||||
- a flat dict with `"method"` or `"quant_method"`
|
||||
- a per-component dict such as `{"transformer": {"method": "fp8"}, "vae": None}`
|
||||
- an already-built `QuantizationConfig`
|
||||
|
||||
Per-component dicts are routed by `ComponentQuantizationConfig` using
|
||||
longest-prefix matching. Prefix stability matters: if a model remaps HF names
|
||||
to vLLM names, apply the mapper to `ignored_layers` or component prefixes.
|
||||
|
||||
## Local Overrides
|
||||
|
||||
The local factory delegates most methods to upstream `vllm`, but has vLLM-Omni
|
||||
overrides for methods that need diffusion, NPU, XPU, or component behavior:
|
||||
|
||||
| Method | Local Config | Main Use |
|
||||
|--------|--------------|----------|
|
||||
| `gguf` | `DiffusionGGUFConfig` | Pre-quantized diffusion transformer weights |
|
||||
| `int8` | `DiffusionInt8Config` | Online or serialized W8A8 diffusion transformers |
|
||||
| `mxfp8` | `DiffusionMXFP8Config` | W8A8 MXFP8 diffusion on NPU and online XPU path |
|
||||
| `mxfp4` | `DiffusionMXFP4Config` | Online W4A4 MXFP4 diffusion on NPU |
|
||||
| `mxfp4_dualscale` | `DiffusionMXFP4DualScaleMixedConfig` | Online/offline W4A4 dual-scale plus BF16 fallback on NPU |
|
||||
| `inc`, `auto-round`, `auto_round` | `OmniINCConfig` | AutoRound/INC checkpoint loading, including MXFP8 checkpoint metadata |
|
||||
| ModelOpt config dicts | upstream ModelOpt configs through auto-detect | FP8, NVFP4, or mixed ModelOpt checkpoints |
|
||||
|
||||
Check `SUPPORTED_QUANTIZATION_METHODS` in `vllm_omni/quantization/factory.py`
|
||||
before adding a new method name.
|
||||
|
||||
## AR and Generic Quantization
|
||||
|
||||
For AR-backed models, upstream `vllm` methods usually apply:
|
||||
|
||||
- `awq` and `gptq`: 4-bit or 8-bit weight quantization with calibration,
|
||||
commonly used for text and AR stages.
|
||||
- `fp8`: FP8 weight path or KV-cache FP8 depending on flags and hardware.
|
||||
- `inc` or `auto-round`: pre-quantized checkpoints with Intel Neural
|
||||
Compressor / AutoRound metadata.
|
||||
|
||||
Use these for the AR language model stage only when the checkpoint and model
|
||||
guide support it. Omni and TTS models often keep audio, vision, talker, and
|
||||
codec stages in BF16.
|
||||
|
||||
## Diffusion and Offline Methods
|
||||
|
||||
For diffusion transformers, prefer the method-specific docs:
|
||||
|
||||
- `docs/user_guide/quantization/fp8.md`
|
||||
- `docs/user_guide/quantization/int8.md`
|
||||
- `docs/user_guide/quantization/gguf.md`
|
||||
- `docs/user_guide/quantization/modelopt.md`
|
||||
- `docs/user_guide/quantization/mxfp8.md`
|
||||
- `docs/user_guide/quantization/mxfp4.md`
|
||||
- `docs/user_guide/quantization/autoround.md`
|
||||
- `docs/user_guide/quantization/msmodelslim.md`
|
||||
|
||||
Pre-quantized checkpoints should normally auto-detect from config. Runtime
|
||||
flags such as `--force-cutlass-fp8`, `--linear-backend cutlass`, and
|
||||
`--moe-backend cutlass` select kernels; they do not quantize a BF16 checkpoint.
|
||||
|
||||
## Validation
|
||||
|
||||
Use fixed-seed A/B comparisons and record:
|
||||
|
||||
- model path and quantized checkpoint path
|
||||
- method config and `ignored_layers`
|
||||
- prompt, size, frame count, steps, scheduler, dtype, and seed
|
||||
- eager/non-eager mode, compile flags, and parallelism
|
||||
- output metrics such as PSNR, MAE, cosine similarity, SSIM, LPIPS, plus visual artifacts
|
||||
- latency and peak memory, reported separately from quality
|
||||
|
||||
The helper `vllm_omni.quantization.tools.compare_diffusion_trajectory_similarity`
|
||||
is the default lightweight image/video comparison tool for diffusion paths.
|
||||
@@ -0,0 +1,47 @@
|
||||
# Quantization Compatibility by Modality
|
||||
|
||||
Use this as a routing guide. The source of truth is the local method docs under
|
||||
`docs/user_guide/quantization/` and the current model implementation.
|
||||
|
||||
## Summary
|
||||
|
||||
- AR and omni language-model quantization mostly follows upstream `vllm`, with
|
||||
vLLM-Omni scoping for multi-stage models.
|
||||
- Diffusion quantization is method-specific and model-specific. It usually
|
||||
targets the diffusion transformer only.
|
||||
- Pre-quantized checkpoints should quantize only the component that has
|
||||
quantized weights and scale tensors.
|
||||
|
||||
## Model Type Matrix
|
||||
|
||||
| Model Type | Typical Quantization Scope | Notes |
|
||||
|------------|----------------------------|-------|
|
||||
| AR text models | Language model | Upstream `vllm` methods usually apply |
|
||||
| Omni/TTS models | Thinker or AR language-model stage | Audio, vision, talker, and codec stages usually stay BF16 |
|
||||
| DiT image models | Diffusion transformer | Check FP8, Int8, GGUF, ModelOpt, AutoRound docs |
|
||||
| DiT video models | Diffusion transformer | HunyuanVideo FP8 and Wan2.2 MXFP paths are documented; validate other video paths |
|
||||
| Multi-stage diffusion | Selected stage only | Use per-component routing and per-stage validation |
|
||||
| Audio diffusion | Usually unsupported | Prefer dtype, offload, cache, or parallelism unless docs say otherwise |
|
||||
|
||||
## Method Hints
|
||||
|
||||
| Goal | Candidate |
|
||||
|------|-----------|
|
||||
| AR memory reduction | AWQ, GPTQ, AutoRound, ModelOpt where checkpoint supports it |
|
||||
| CUDA diffusion W8A8 | FP8 or Int8 |
|
||||
| Pre-quantized CUDA diffusion checkpoint | ModelOpt or GGUF |
|
||||
| Ascend NPU diffusion | Int8, MXFP8, MXFP4, MXFP4 DualScale, msModelSlim |
|
||||
| Intel XPU checkpoint | AutoRound/INC paths |
|
||||
| Conservative image quality | selective FP8/Int8 with `ignored_layers` |
|
||||
| Unsupported diffusion model | cache acceleration, CPU/layerwise offload, dtype, TP/SP/CFG parallelism |
|
||||
|
||||
## Validation Rule
|
||||
|
||||
Compatibility means more than successful loading. A method is compatible only
|
||||
after it has:
|
||||
|
||||
- correct component scope
|
||||
- complete checkpoint/config metadata if offline
|
||||
- fixed-seed quality comparison against BF16
|
||||
- acceptable latency and memory behavior
|
||||
- docs or examples that state model, method, hardware, and caveats
|
||||
@@ -0,0 +1,122 @@
|
||||
# ModelOpt FP8 Conversion and Loading
|
||||
|
||||
Use this when exporting or loading ModelOpt FP8 checkpoints for diffusion
|
||||
transformers. This reference captures the practical conversion lessons from the
|
||||
2026-04-21 ModelOpt FP8 paste and aligns them with the current
|
||||
`vllm-omni` ModelOpt adapter.
|
||||
|
||||
## Core Rule
|
||||
|
||||
Treat ModelOpt FP8 as selective checkpoint quantization, not one-click
|
||||
whole-model FP8. Quantize the main transformer or DiT backbone, keep sensitive
|
||||
or peripheral modules BF16, serialize a complete ModelOpt checkpoint, and let
|
||||
vLLM-Omni auto-detect it from config.
|
||||
|
||||
Do not pass `--quantization fp8` to force a ModelOpt checkpoint path. The
|
||||
checkpoint's `quantization_config` should select `modelopt`.
|
||||
|
||||
## Conversion Strategy
|
||||
|
||||
Prefer:
|
||||
|
||||
- backbone attention Q/K/V and FFN expansion projections
|
||||
- conservative per-model ignore lists
|
||||
- MHA quantizers disabled until basic linear quantization is stable
|
||||
- full module-name prefixes in ignore/routing logic
|
||||
- BF16 fallback for output-side and routing-sensitive layers
|
||||
|
||||
Avoid quantizing by default:
|
||||
|
||||
- VAE
|
||||
- text, vision, or audio encoders
|
||||
- embedders and modulation layers
|
||||
- final heads and output projections
|
||||
- refiner modules
|
||||
- MoE routers and layernorm-sensitive blocks
|
||||
|
||||
## Model-Specific Starting Points
|
||||
|
||||
| Model Family | Quantize First | Keep BF16 First |
|
||||
|--------------|----------------|-----------------|
|
||||
| Z-Image | attention q/k/v, FFN w1/w3 | `to_out`, `w2`, `noise_refiner`, `context_refiner`, embedders, modulation, final/output |
|
||||
| Qwen-Image | transformer backbone linear layers | `proj_out`, `img_in`, `txt_in`, norms, modulation, position/time/embed layers |
|
||||
| FLUX2-dev / FLUX2-klein | backbone linear layers | `proj_out`, context/x embedders, `norm_out`, stream modulation, time embeddings |
|
||||
| HunyuanImage-3.0 | `model.model` backbone | VAE, vision tower, token/time/image embedders, final image head, LM head, routers, layernorm-sensitive blocks |
|
||||
|
||||
These are starting points, not support claims. Validate with BF16 outputs.
|
||||
|
||||
## Required Checkpoint Metadata
|
||||
|
||||
`transformer/config.json` must include a complete `quantization_config`.
|
||||
At minimum it should identify ModelOpt and the algorithm:
|
||||
|
||||
```json
|
||||
{
|
||||
"quantization_config": {
|
||||
"producer": {"name": "modelopt"},
|
||||
"quant_method": "modelopt",
|
||||
"quant_algo": "FP8"
|
||||
}
|
||||
}
|
||||
```
|
||||
|
||||
Current vLLM-Omni auto-detects:
|
||||
|
||||
- `FP8` and `FP8_PER_CHANNEL_PER_TOKEN` -> `modelopt`
|
||||
- `NVFP4` -> `modelopt_fp4`
|
||||
- `MIXED_PRECISION` -> `modelopt_mixed`
|
||||
|
||||
Packed modules and scales must map correctly:
|
||||
|
||||
- `to_qkv` from `to_q`, `to_k`, `to_v`
|
||||
- `add_kv_proj` from `add_q_proj`, `add_k_proj`, `add_v_proj`
|
||||
- `w13` from `w1`, `w3`
|
||||
- `.input_scale`, `.weight_scale`, `.weight_scale_2`, `.weight_scale_inv`
|
||||
|
||||
For selective FP8, target layers that remain BF16 must be allowed to
|
||||
dequantize from FP8 weights using the matching `weight_scale`, or the adapter
|
||||
must skip the quantized scale tensors for that BF16 target.
|
||||
|
||||
## Runtime Signals
|
||||
|
||||
Healthy loading should show ModelOpt auto-detection and a ModelOpt linear
|
||||
kernel path. Useful signals include:
|
||||
|
||||
- auto-detected quantization from config
|
||||
- selected ModelOpt FP8 linear method
|
||||
- CUTLASS or FlashInfer FP8 scaled-mm backend when enabled and supported
|
||||
- non-eager path compiling the transformer when expected
|
||||
|
||||
If online serving fails while offline inference works, inspect warmup and
|
||||
serving initialization before changing checkpoint conversion.
|
||||
|
||||
## Validation Order
|
||||
|
||||
1. Check config completeness.
|
||||
2. Run an export checker such as `examples/quantization/check_modelopt_fp8_export.py`.
|
||||
3. Confirm actual FP8 tensors and expected disk-size reduction.
|
||||
4. Run one offline generation with fixed prompt, seed, size, and steps.
|
||||
5. Run serving with `vllm serve --omni` and verify `/v1/images/generations`.
|
||||
6. Compare against BF16 using visual review plus metrics such as LPIPS, PSNR,
|
||||
MAE, cosine similarity, and SSIM.
|
||||
7. Compare performance only under matched modes:
|
||||
- BF16 eager vs FP8 eager
|
||||
- BF16 non-eager vs FP8 non-eager
|
||||
|
||||
FP8 is not automatically faster. Check whether execution reaches FP8 kernels,
|
||||
whether compile mode is comparable, and whether the workload shape fits the
|
||||
kernel path.
|
||||
|
||||
## Triage
|
||||
|
||||
| Symptom | First Check |
|
||||
|---------|-------------|
|
||||
| Startup failure | `config.json`, `model_index.json`, `quantization_config`, missing scales |
|
||||
| Loads but output is corrupted | Over-quantized sensitive layers or wrong full-name prefixes |
|
||||
| Offline works but online fails | Dummy warmup, serving init, compile path |
|
||||
| FP8 is not faster | Kernel selection, dtype, compile mode, workload shape |
|
||||
| BF16 fallback fails | Missing `weight_scale` for dequantization or wrong packed-module mapping |
|
||||
|
||||
The success criterion is stable, correctly scoped checkpoint loading with
|
||||
quality preserved enough for the target use case, not maximizing the number of
|
||||
FP8 layers.
|
||||
@@ -21,11 +21,18 @@ 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
|
||||
- `generate-release-note`: helps prepare release notes for repository changes
|
||||
- `precheck-pr`: self-check a branch before creating a PR — validates PR title
|
||||
format, catches dead code, verifies accuracy/perf claims, and confirms merge
|
||||
readiness
|
||||
- `review-pr`: provides a structured workflow for reviewing pull requests
|
||||
- `vllm-omni-npu-model-runner-upgrade`: upgrades NPU model runners to align with the
|
||||
latest vllm-ascend NPUModelRunner
|
||||
|
||||
## Maintenance Guidelines
|
||||
|
||||
|
||||
@@ -1,21 +1,15 @@
|
||||
<!-- markdownlint-disable -->
|
||||
PLEASE FILL IN THE PR DESCRIPTION HERE ENSURING ALL CHECKLIST ITEMS (AT THE BOTTOM) HAVE BEEN CONSIDERED.
|
||||
PLEASE FILL IN THE PR DESCRIPTION HERE.
|
||||
|
||||
## Purpose
|
||||
|
||||
## Test Plan
|
||||
|
||||
**vLLM Version:**
|
||||
|
||||
**vLLM-Omni Commit:**
|
||||
|
||||
## Test Result
|
||||
|
||||
---
|
||||
<details>
|
||||
<summary> Essential Elements of an Effective PR Description Checklist </summary>
|
||||
|
||||
- [ ] The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
|
||||
- [ ] The test plan. Please provide the test scripts & test commands. Please state the reasons if your codes don't require additional test scripts. For test file guidelines, please check the [test style doc](https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/ci/tests_style/)
|
||||
- [ ] The test results. Please paste the results comparison before and after, or the e2e results.
|
||||
- [ ] (Optional) The necessary documentation update, such as updating `supported_models.md` and `examples` for a new model. **Please run `mkdocs serve` to sync the documentation editions to `./docs`.**
|
||||
- [ ] (Optional) Release notes update. If your change is user-facing, please update the release notes draft.
|
||||
</details>
|
||||
|
||||
**BEFORE SUBMITTING, PLEASE READ <https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md>** (anything written below this line will be removed by GitHub Actions)
|
||||
**BEFORE SUBMITTING:** read [CONTRIBUTING.md](https://github.com/vllm-project/vllm-omni/blob/main/CONTRIBUTING.md) and run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
|
||||
(anything written below this line will be removed by GitHub Actions)
|
||||
@@ -1,3 +1,5 @@
|
||||
# Contributing to vLLM-Omni
|
||||
|
||||
You may find information about contributing to vLLM-Omni on [Contributing](https://vllm-omni.readthedocs.io/en/latest/contributing/)
|
||||
|
||||
Before submitting a PR, run the [precheck-pr skill](https://github.com/vllm-project/vllm-omni/blob/main/.claude/skills/precheck-pr/SKILL.md) with the code agent for a self-check against project conventions.
|
||||
@@ -16,6 +16,7 @@ Easy, fast, and cheap omni-modality model serving for everyone
|
||||
---
|
||||
|
||||
*Latest News* 🔥
|
||||
- [2026/06] We released [0.22.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.22.0) - an **omnimodal world-model** release aligned with vLLM 0.22, featuring [Nvidia Cosmos3](recipes/cosmos3/Cosmos3-Nano.md)/DreamZero world model support, expanded quantization coverage across Blackwell/NPU/XPU, TTS production improvements, new models including MiniCPM-o 4.5, MOSS-TTS, and Lance, plus RL integration with [VeRL-Omni](https://github.com/verl-project/verl-omni).
|
||||
- [2026/05] We released [0.20.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.20.0) - refreshes the serving/runtime stack for large-scale omni workloads, and improves diffusion model performance, quantization, and hardware readiness across CUDA, ROCm, MUSA, NPU, and XPU backends.
|
||||
- [2026/03] We released [0.18.0](https://github.com/vllm-project/vllm-omni/releases/tag/v0.18.0) - strengthens the core runtime through a large entrypoint refactor and scheduler/runtime cleanups, expands unified quantization and diffusion execution, broadens multimodal model coverage, and improves production readiness across audio, omni, image, video, RL, and multi-platform deployments.
|
||||
- [2026/03] Check out our first public [project deepdive](https://youtu.be/sgwNfsNnR9I) at the vLLM Hong Kong Meetup!
|
||||
@@ -56,8 +57,9 @@ vLLM-Omni is flexible and easy to use with:
|
||||
|
||||
vLLM-Omni seamlessly supports most popular open-source models on HuggingFace, including:
|
||||
|
||||
- Omni-modality models (e.g. Qwen-Omni)
|
||||
- Multi-modality generation models (e.g. Qwen-Image)
|
||||
- **Omni-modality models** (e.g. Qwen3-Omni, Cosmos, HunyuanImage, BAGEL)
|
||||
- **TTS models** (e.g. Qwen3-TTS, VoxCPM2, Ming-Omni-TTS, CosyVoice3)
|
||||
- **Diffusion models** — image, video, and audio generation (e.g. Qwen-Image, Wan2.2, FLUX)
|
||||
|
||||
## Getting Started
|
||||
|
||||
|
||||
@@ -33,6 +33,15 @@ models:
|
||||
voice_clone:
|
||||
task_type: Base
|
||||
|
||||
bosonai/higgs-audio-v3-tts-4b:
|
||||
supported_tasks: [default_voice, voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
trust_remote_code: true
|
||||
task_extra_body:
|
||||
default_voice: {}
|
||||
voice_clone: {}
|
||||
|
||||
openbmb/VoxCPM2:
|
||||
supported_tasks: [voice_clone, default_voice]
|
||||
trust_remote_code: true
|
||||
|
||||
@@ -87,6 +87,8 @@ nav:
|
||||
- AutoRound: user_guide/quantization/autoround.md
|
||||
- msModelSlim: user_guide/quantization/msmodelslim.md
|
||||
- ComfyUI: features/comfyui.md
|
||||
- Fault Tolerance:
|
||||
- Failure Modes: user_guide/fault_injection_reliability_matrix.md
|
||||
- Developer Guide:
|
||||
- General:
|
||||
- contributing/README.md
|
||||
|
||||
@@ -92,7 +92,6 @@ Core scheduling and caching components.
|
||||
Configuration classes.
|
||||
|
||||
- [vllm_omni.config.model.OmniModelConfig][]
|
||||
- [vllm_omni.config.stage_config.ModelPipeline][]
|
||||
- [vllm_omni.config.stage_config.StageConfig][]
|
||||
- [vllm_omni.config.stage_config.StageConfigFactory][]
|
||||
- [vllm_omni.config.stage_config.StageType][]
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 408 KiB After Width: | Height: | Size: 403 KiB |
@@ -116,6 +116,16 @@ Only specific types of PRs will be reviewed. The PR title is prefixed appropriat
|
||||
!!! note
|
||||
If the PR spans more than one category, please include all relevant prefixes.
|
||||
|
||||
### 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:
|
||||
|
||||
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.
|
||||
|
||||
### Local Test
|
||||
Please run the L1 and L2 test cases locally first and attach the results before contacting us to add the "ready" label. Please refer to the [test instructions](./ci/test_guide.md) for running the test cases.
|
||||
|
||||
|
||||
@@ -268,6 +268,101 @@ 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
|
||||
@@ -302,6 +397,7 @@ A clear directory structure is key to managing test cases efficiently.
|
||||
### 1.4 Execution Method and Example
|
||||
|
||||
- ***Trigger Timing***: **`PR with ready label`**. That is, when a developer adds a "ready for review" or similar label to a PR on platforms like GitHub, L1 and L2 tests are automatically triggered.
|
||||
- ***Diff-aware step skipping***: On L2, 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***:
|
||||
|
||||
@@ -404,7 +500,8 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
### 2.4 Execution Method and Example
|
||||
|
||||
- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch.
|
||||
- ***Trigger Timing***: **`PR Merged`**. Automatically triggered after code review is approved and merged into the main branch (typically via `merge-test` label on the PR before merge).
|
||||
- ***Diff-aware step skipping***: Same [`source_file_dependencies`](#diff-aware-buildkite-uploads-source_file_dependencies) mechanism as L2—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).
|
||||
- ***Execution Environment***: ***GPU*** servers.
|
||||
- ***Script Example***:
|
||||
|
||||
@@ -520,6 +617,8 @@ L3 level testing executes after code is merged into the main branch. Its core pu
|
||||
|
||||
**Implicit Validation**: The `send_omni_request` and `send_diffusion_request` methods internally includes validation logic dynamically selected based on the --run-level parameter: core_model performs basic validation, while advanced_model and full_model perform deep validation.
|
||||
|
||||
**Audio output debugging**: Deep validation may transcribe returned audio via `convert_audio_bytes_to_text` (Whisper). If an audio keyword or text–audio similarity assertion fails, set `VLLM_OMNI_KEEP_REQUEST_MEDIA=1` before running pytest to keep the intermediate WAV files for inspection (see [Test helper environment variables](#test-helper-environment-variables)).
|
||||
|
||||
- ***Run Command (L3 merge)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}.py -m advanced_model --run-level=advanced_model`
|
||||
|
||||
- ***Run Command (L4 nightly expansion)***: `pytest -s -v /tests/e2e/online_serving/test_{model_name}_expansion.py -m full_model --run-level=full_model`
|
||||
|
||||
@@ -69,6 +69,8 @@ If the failure is already tracked, leave a comment on your PR noting that the fa
|
||||
|
||||
If the failure appears to be new, investigate whether your changes caused it.
|
||||
|
||||
For resilience behavior under injected process signals/OOM pressure (validation scenarios, not root-cause equivalence), see [Failure Mode Injection Scenarios and Expected Behavior Matrix](../../user_guide/fault_injection_reliability_matrix.md).
|
||||
|
||||
### Reading Buildkite Logs
|
||||
|
||||
1. Click **Details** next to the Buildkite check on your PR.
|
||||
|
||||
@@ -8,6 +8,7 @@ Backend-specific setup lives in separate docs:
|
||||
- [SharedMemoryConnector](omni_connectors/shared_memory_connector.md)
|
||||
- [MooncakeStoreConnector](omni_connectors/mooncake_store_connector.md)
|
||||
- [MooncakeTransferEngineConnector](omni_connectors/mooncake_transfer_engine_connector.md)
|
||||
- [MoriTransferEngineConnector](omni_connectors/mori_transfer_engine_connector.md)
|
||||
- [YuanrongConnector](omni_connectors/yuanrong_connector.md)
|
||||
- [YuanrongTransferEngineConnector](omni_connectors/yuanrong_transfer_engine_connector.md)
|
||||
|
||||
@@ -23,6 +24,7 @@ Current connectors operate in D2H2D (device to host to device) mode.
|
||||
| Single node | SharedMemoryConnector | Auto-configured if no connector is specified. |
|
||||
| Multi node (Mooncake Store) | MooncakeStoreConnector | TCP-based, requires Mooncake Master + metadata server. |
|
||||
| Multi node (Mooncake RDMA) | MooncakeTransferEngineConnector | RDMA/TCP direct transfer with managed memory pool. Fastest. |
|
||||
| Multi node (Mori RDMA) | MoriTransferEngineConnector | RDMA direct transfer via Mori IOEngine. |
|
||||
| Multi node (Yuanrong) | YuanrongConnector | Requires Yuanrong Datasystem + etcd. |
|
||||
| Ascend NPU P2P (Yuanrong TE) | YuanrongTransferEngineConnector | Uses Yuanrong TransferEngine directly. Configure NPU device IPv4 and `memory_pool_device: "npu"`. |
|
||||
|
||||
|
||||
@@ -0,0 +1,76 @@
|
||||
# MoriTransferEngineConnector
|
||||
|
||||
## When to Use
|
||||
|
||||
Currently supports intra-node deployment with Mori.
|
||||
|
||||
As noted in #1742, inter-node support will be added back in a future
|
||||
refactor.
|
||||
|
||||
## Mechanism
|
||||
|
||||
Uses Mori's `IOEngine` / `MemoryDesc` API for zero-copy RDMA transfers.
|
||||
|
||||
- Data Plane: RDMA (InfiniBand/RoCE) with managed memory pool.
|
||||
- Control Plane: ZMQ for pull-request handshake and async completion.
|
||||
|
||||
## Installation
|
||||
|
||||
See the [Mori repository](https://github.com/ROCm/mori) for installation instructions.
|
||||
|
||||
## Configuration
|
||||
|
||||
Mori is configured through the new deploy-config schema (see
|
||||
[`docs/configuration/stage_configs.md`](../../../configuration/stage_configs.md)).
|
||||
Define the connector at the top level of the deploy YAML and reference it
|
||||
by name from each stage's `input_connectors` / `output_connectors`:
|
||||
|
||||
```yaml
|
||||
connectors:
|
||||
mori_connector:
|
||||
name: MoriTransferEngineConnector
|
||||
extra:
|
||||
host: "auto"
|
||||
zmq_port: 50051
|
||||
device_name: ""
|
||||
memory_pool_size: 536870912
|
||||
memory_pool_device: "cuda"
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
output_connectors:
|
||||
to_stage_1: mori_connector
|
||||
|
||||
- stage_id: 1
|
||||
input_connectors:
|
||||
from_stage_0: mori_connector
|
||||
```
|
||||
|
||||
A ready-to-run intra-node example for Qwen3-Omni-MoE on AMD MI300X lives
|
||||
at `vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml` and can be loaded
|
||||
with:
|
||||
|
||||
```bash
|
||||
vllm-omni serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --log-stats \
|
||||
--deploy-config vllm_omni/deploy/qwen3_omni_moe_mori_intranode.yaml
|
||||
```
|
||||
|
||||
The yaml wires `MoriTransferEngineConnector` (with `backend_type: xgmi`) to
|
||||
the chunk_transfer_adapter path (`async_chunk: true`) so stage-to-stage
|
||||
hidden-state and codec-frame streams ship GPU-to-GPU over AMD Infinity
|
||||
Fabric instead of SHM. Qwen2.5-Omni + Mori is not yet functional on
|
||||
the chunk path: it needs `thinker2talker_async_chunk` /
|
||||
`talker2code2wav_async_chunk` input processors that do not exist yet
|
||||
(the orchestrator-level path the upstream PR originally targeted was
|
||||
lost during the entrypoints → engine refactor and removed in #1742).
|
||||
|
||||
Parameters:
|
||||
|
||||
- host: local RDMA IP (`"auto"` for auto-detect).
|
||||
- zmq_port: ZMQ base port for control-plane communication.
|
||||
- device_name: RDMA device (e.g., `"mlx5_0"`), empty for auto-detect.
|
||||
- memory_pool_size: RDMA memory pool size in bytes.
|
||||
- memory_pool_device: `"cpu"` (pinned) or `"cuda"` (GPUDirect / XGMI RDMA).
|
||||
|
||||
For more details, refer to the
|
||||
[Mori repository](https://github.com/ROCm/mori).
|
||||
@@ -305,8 +305,8 @@ class GPUWorker:
|
||||
)
|
||||
|
||||
# Load model
|
||||
model_loader = DiffusersPipelineLoader(load_config)
|
||||
self.pipeline = model_loader.load_model(od_config, load_device=f"cuda:{rank}")
|
||||
model_loader = DiffusersPipelineLoader(load_config, self.od_config)
|
||||
self.pipeline = model_loader.load_model(load_device=f"cuda:{rank}")
|
||||
|
||||
# Setup cache backend
|
||||
from vllm_omni.diffusion.cache.selector import get_cache_backend
|
||||
|
||||
@@ -45,7 +45,7 @@ If you do not need to modify source code of vLLM, you can directly install the s
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
The 0.22.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The 0.22.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -70,8 +70,8 @@ git checkout v0.22.0
|
||||
```
|
||||
Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file.
|
||||
```bash
|
||||
#For CUDA 13.0 (the default for v0.22.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_35_x86_64.whl
|
||||
#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
|
||||
@@ -22,7 +22,7 @@ source .venv/bin/activate
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.22.0+rocm721 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm721
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
|
||||
@@ -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, 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` | ✅︎ | ✅︎ | | |
|
||||
@@ -50,6 +50,8 @@ th {
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
|
||||
| `VoxCPM2TalkerForConditionalGeneration` | VoxCPM2 | `openbmb/VoxCPM2` | ✅︎ | | | |
|
||||
| `MammothModa2ForConditionalGeneration` | MammothModa2-Preview | `bytedance-research/MammothModa2-Preview` | ✅︎ | ✅︎ | | |
|
||||
| `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | |
|
||||
@@ -61,7 +63,12 @@ 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` | ✅︎ | | | |
|
||||
| `MossTTSRealtime` | MOSS-TTS-Realtime | `OpenMOSS-Team/MOSS-TTS-Realtime` | ✅︎ | | | |
|
||||
| `HiggsAudioV2ForConditionalGeneration` | Higgs-Audio v2 | `bosonai/higgs-audio-v2-generation-3B-base` | ✅︎ | | | |
|
||||
| `NextStep11Pipeline` | NextStep-1.1 | `stepfun-ai/NextStep-1.1` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `MiMoAudioModel` | MiMo-Audio-7B-Instruct | `XiaomiMiMo/MiMo-Audio-7B-Instruct` | ✅︎ | ✅︎ | | |
|
||||
| `MiMoV2ASRForCausalLM` | MiMo-V2.5-ASR | `XiaomiMiMo/MiMo-V2.5-ASR` | ✅︎ | ✅︎ | | |
|
||||
@@ -69,11 +76,13 @@ th {
|
||||
| `FishSpeechSlowARForConditionalGeneration` | Fish Speech S2 Pro | `fishaudio/s2-pro` | ✅︎ | ✅︎ | | |
|
||||
| `DreamIDOmniPipeline` | DreamID-Omni | `XuGuo699/DreamID-Omni` | ✅︎ | ✅︎ | | |
|
||||
| `SenseNovaU1Pipeline` | SenseNova-U1 (DiT-only) | `SenseNova/SenseNova-U1-8B-MoT` | ✅︎ | | | |
|
||||
| `LancePipeline` | Lance | `bytedance-research/Lance` | ✅︎ | | | |
|
||||
| `HunyuanVideo15Pipeline` | HunyuanVideo-1.5-T2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | ✅︎ | ✅︎ | | |
|
||||
| `HunyuanVideo15ImageToVideoPipeline` | HunyuanVideo-1.5-I2V | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v`, `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_i2v` | ✅︎ | ✅︎ | | |
|
||||
| `VoxtralTTSForConditionalGeneration` | Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | ✅︎ | ✅︎ | | |
|
||||
| `CovoAudioForConditionalGeneration` | Covo-Audio-Chat | `tencent/Covo-Audio-Chat` | ✅︎ | | | |
|
||||
|`DyninOmniForConditionalGeneration` | Dynin-Omni | `snu-aidas/Dynin-Omni` | ✅︎ | | | |
|
||||
| `MiniCPMO45OmniForConditionalGeneration` | MiniCPM-o 4.5 | `openbmb/MiniCPM-o-4_5` | ✅︎ | | ✅︎ | |
|
||||
| `ErnieImagePipeline` | ERNIE-Image | `baidu/ERNIE-Image`, `baidu/ERNIE-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
|`HiDreamImagePipeline` | HiDream-I1-Full | `HiDream-ai/HiDream-I1-Full` | ✅︎ | ✅︎ | | |
|
||||
|
||||
|
||||
@@ -114,9 +114,10 @@ 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:** Auto, Chinese, English, Japanese, Korean, German, French, Russian, Portuguese, Spanish, Italian
|
||||
**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.
|
||||
|
||||
#### Voice Clone Parameters (Base task)
|
||||
|
||||
@@ -332,7 +333,8 @@ curl -X POST http://localhost:8091/v1/audio/speech \
|
||||
"input": "Hello, this is a cloned voice",
|
||||
"task_type": "Base",
|
||||
"ref_audio": "https://example.com/reference.wav",
|
||||
"ref_text": "Original transcript of the reference audio"
|
||||
"ref_text": "Original transcript of the reference audio",
|
||||
"non_streaming_mode": true
|
||||
}' --output cloned.wav
|
||||
```
|
||||
|
||||
@@ -454,6 +456,7 @@ Content-Type: application/json
|
||||
| `ref_audio` | string | null | Default reference audio (Base task) |
|
||||
| `ref_text` | string | null | Default reference transcript (Base task) |
|
||||
| `max_new_tokens` | integer | null | Default max tokens |
|
||||
| `non_streaming_mode` | bool | null | Default Qwen3-TTS prompt construction mode override. Does not affect HTTP response streaming or async-chunk pipelining. When null, Qwen3-TTS uses model defaults: Base=false, CustomVoice/VoiceDesign=true. |
|
||||
|
||||
Each item in the `items` array requires only `input` (the text). All other fields are optional and override the batch-level defaults when set:
|
||||
|
||||
@@ -469,6 +472,7 @@ Each item in the `items` array requires only `input` (the text). All other field
|
||||
| `ref_audio` | string | Override reference audio |
|
||||
| `ref_text` | string | Override reference transcript |
|
||||
| `max_new_tokens` | integer | Override max tokens |
|
||||
| `non_streaming_mode` | bool | Override Qwen3-TTS prompt construction mode. Does not affect HTTP response streaming or async-chunk pipelining. When null, inherits the batch-level value (then the model default). |
|
||||
|
||||
### Response Format
|
||||
|
||||
|
||||
@@ -67,8 +67,10 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `input_reference` | file | null | Uploaded reference image for image-to-video requests |
|
||||
| `image_reference` | string | null | JSON-encoded reference image payload; do not combine with `input_reference` |
|
||||
| `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 |
|
||||
@@ -80,6 +82,8 @@ curl -L "http://localhost:8091/v1/videos/${video_id}/content" -o output.mp4
|
||||
| `flow_shift` | number | null | Scheduler flow-shift value |
|
||||
| `true_cfg_scale` | number | null | True CFG scale when supported by the model |
|
||||
| `seed` | integer | null | Random seed for reproducibility |
|
||||
| `generate_sound` | boolean | false | Request model-generated audio for video models that support sound generation |
|
||||
| `sound_duration` | number | null | Duration in seconds for generated audio; defaults to generated video duration |
|
||||
| `negative_prompt` | string | null | Text describing what to avoid in the generated video |
|
||||
| `enable_frame_interpolation` | boolean | null | Enable post-generation frame interpolation |
|
||||
| `frame_interpolation_exp` | integer | null | Interpolation exponent; `1=2x`, `2=4x`, and so on |
|
||||
@@ -123,6 +127,60 @@ curl -s http://localhost:8091/v1/videos \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
### Video-to-Video
|
||||
|
||||
For models that support video conditioning, upload the reference video with
|
||||
`input_reference`:
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F "input_reference=@input.mp4;type=video/mp4" \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
You can also pass a JSON-safe video URL or `data:video/...;base64,...` payload
|
||||
through `video_reference`. Do not send `video_reference` together with
|
||||
`input_reference` or `image_reference`.
|
||||
|
||||
```bash
|
||||
curl -s http://localhost:8091/v1/videos \
|
||||
-F "prompt=continue this motion with consistent subjects and lighting" \
|
||||
-F 'video_reference={"video_url":"https://example.com/input.mp4"}' \
|
||||
-F "width=1280" \
|
||||
-F "height=720" \
|
||||
-F "num_frames=80" \
|
||||
-F "fps=16"
|
||||
```
|
||||
|
||||
JSON references currently support `image_url`/`video_url`; `file_id` references
|
||||
are not implemented yet. Models may expose additional V2V controls through
|
||||
`extra_params`. For example, Cosmos3 supports
|
||||
`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
|
||||
@@ -146,7 +204,12 @@ export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete text-to-video and image-to-video walkthroughs, see:
|
||||
For complete text-to-video, image-to-video, and model-specific video-to-video
|
||||
walkthroughs, see:
|
||||
|
||||
- [Text-to-Video](../user_guide/examples/online_serving/text_to_video.md)
|
||||
- [Image-to-Video](../user_guide/examples/online_serving/image_to_video.md)
|
||||
- [Speech-to-Video](../user_guide/examples/online_serving/speech_to_video.md)
|
||||
for Wan2.2-S2V audio-driven lip-sync generation
|
||||
- [Cosmos3 recipes](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-Nano.md)
|
||||
for model-specific video-to-video examples and conditioning controls
|
||||
@@ -148,7 +148,7 @@ The following tables show which models support each feature:
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Helios** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (encode/decode) | ✅ | ❌ |
|
||||
| **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
|
||||
@@ -0,0 +1,128 @@
|
||||
# Speech-To-Video
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/online_serving/speech_to_video>.
|
||||
|
||||
This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Model ID |
|
||||
|-------|----------|
|
||||
| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` |
|
||||
|
||||
## Start Server
|
||||
|
||||
### Basic Start
|
||||
|
||||
```bash
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size 2 \
|
||||
--flow-shift 3.0 \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--cache-backend cache_dit \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
### Start with Script
|
||||
|
||||
```bash
|
||||
bash run_server.sh
|
||||
```
|
||||
|
||||
The script allows overriding:
|
||||
- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`)
|
||||
- `PORT` (default: `8091`)
|
||||
- `FLOW_SHIFT` (default: `3.0`)
|
||||
- `TP` (default: `2`)
|
||||
- `CACHE_BACKEND` (default: `cache_dit`)
|
||||
|
||||
### Start with 4 GPUs
|
||||
|
||||
```bash
|
||||
TP=4 bash run_server.sh
|
||||
```
|
||||
|
||||
## API Endpoints
|
||||
|
||||
The server provides OpenAI-compatible video generation endpoints:
|
||||
|
||||
- `POST /v1/videos`: create a video generation job (async)
|
||||
- `POST /v1/videos/sync`: generate a video and return raw bytes (sync, for benchmarks)
|
||||
- `GET /v1/videos/{video_id}`: retrieve the current job status and metadata
|
||||
- `GET /v1/videos`: list stored video jobs
|
||||
- `GET /v1/videos/{video_id}/content`: download the generated video file
|
||||
- `DELETE /v1/videos/{video_id}`: delete the job and any stored output
|
||||
|
||||
## Sync API (Recommended for Testing)
|
||||
|
||||
`POST /v1/videos/sync` blocks until generation completes and returns the raw
|
||||
video bytes (`video/mp4`) directly in the response body.
|
||||
|
||||
```bash
|
||||
bash run_curl_speech_to_video.sh
|
||||
```
|
||||
|
||||
## Async Job API
|
||||
|
||||
`POST /v1/videos` creates an asynchronous job and returns immediately.
|
||||
|
||||
```bash
|
||||
no_proxy=127.0.0.1 \
|
||||
create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A person singing" \
|
||||
-F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
|
||||
-F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
|
||||
-F "width=832" -F "height=480" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.5" \
|
||||
-F "fps=16")
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
|
||||
# Poll until complete
|
||||
while true; do
|
||||
status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
# Download
|
||||
curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4
|
||||
```
|
||||
|
||||
## Request Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `prompt` | str | - | Text description of the desired video |
|
||||
| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `width` | int | None | Video width in pixels |
|
||||
| `height` | int | None | Video height in pixels |
|
||||
| `fps` | int | None | Output video frame rate |
|
||||
| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) |
|
||||
| `guidance_scale` | float | None | CFG guidance scale |
|
||||
| `seed` | int | None | Random seed for reproducibility |
|
||||
|
||||
## Notes
|
||||
|
||||
- S2V requires both a reference image (`image_reference`) and an audio reference
|
||||
(`audio_reference`). The generated video will show a person matching the reference
|
||||
image with lip movements synchronized to the audio.
|
||||
- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be:
|
||||
- An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`)
|
||||
- A base64 data URL (e.g., `data:audio/mp3;base64,...`)
|
||||
- `--model-class-name WanS2VPipeline` is required on the server to select the
|
||||
S2V pipeline (distinct from the T2V/I2V pipelines).
|
||||
- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps.
|
||||
- Audio is muxed into the output MP4 automatically.
|
||||
- Always pass `fps=16` to ensure correct video/audio alignment.
|
||||
@@ -0,0 +1,117 @@
|
||||
# vLLM-Omni Failure Mode Injection Scenarios and Expected Behavior Matrix
|
||||
|
||||
## Background and Goals
|
||||
|
||||
The online serving pipeline of vLLM-Omni can be abstracted as: `Serve` (service entry process) -> `Engine` (model orchestration unit) -> `Worker` (execution process) -> `GPU` (compute resource).
|
||||
|
||||
- `Serve` provides external APIs, accepts requests, and manages lifecycle.
|
||||
- `Engine` handles request scheduling and inference orchestration; the number and organization of engines can vary by model type.
|
||||
- `Worker` performs concrete model computation, usually bound to underlying GPU resources.
|
||||
- `GPU` is the resource layer that ultimately carries inference compute and memory usage.
|
||||
|
||||
<div markdown style="text-align: center">
|
||||
|
||||
```mermaid
|
||||
%%{init: {'flowchart': {'useMaxWidth': false}}}%%
|
||||
flowchart TD
|
||||
R(request)--> A(serve)
|
||||
A --> B(engine)
|
||||
A --> C(engine)
|
||||
A --> D(engine)
|
||||
B --> E(worker)
|
||||
C --> F(worker)
|
||||
D --> I(worker)
|
||||
E --> G(GPU)
|
||||
F --> H(GPU)
|
||||
I --> L(GPU)
|
||||
```
|
||||
|
||||
</div>
|
||||
|
||||
The diagram above shows an abstract relationship. Omni models and Diffusion models share the same hierarchy; differences mainly come from engine decomposition, worker count, and scheduling strategy.
|
||||
|
||||
In real customer environments, the system will inevitably encounter accidental operations or environmental disturbances, such as unintended process termination, container/node jitter, or short-term resource contention. From the system perspective, many of these events map to two fault categories: "process receives abnormal signals" or "GPU memory is squeezed." For example, pressing `Ctrl+C` essentially sends a `SIGINT` signal to the `serve` process; running `kill -15 <pid>` maps to graceful termination via `SIGTERM`; running `kill -9 <pid>` maps to forced termination via `SIGKILL`; when `docker kill` is executed or a Pod is deleted, the container runtime usually sends `SIGKILL`.
|
||||
|
||||
This document currently covers two core failure modes, and the system is expected to show predictable and observable behavior under each mode:
|
||||
|
||||
- **Failure Mode 1: Process receives abnormal signals (`SIGINT` / `SIGTERM` / `SIGKILL`)**
|
||||
Focus on three aspects: whether processes exit as expected and complete cleanup, whether requests fail fast or connections are interrupted, and whether GPU resources are eventually fully released with no residue.
|
||||
- **Failure Mode 2: OOM (occupy all free GPU memory via an extra process)**
|
||||
Focus on two aspects: whether service health degrades to the expected state (for example, `503`), and whether requests fail within an acceptable time without hanging.
|
||||
|
||||
More failure modes (for example, network jitter and network interruption) will be added in future iterations to improve end-to-end reliability validation coverage.
|
||||
|
||||
> Note: independent fault injection on the `engine` component is not covered in the current version and will be added gradually in future versions.
|
||||
|
||||
## Fault Injection Scenario Matrix
|
||||
|
||||
| Scenario | Fault Type | System Behavior | Current Status |
|
||||
|------|----------------------------|----------|----------|
|
||||
| No load | Send `SIGKILL` to Worker | Worker process is killed immediately; main service detects child-process loss and turns unavailable; API enters stable 5xx | |
|
||||
| No load | Send `SIGTERM` to Worker | Worker exits after receiving termination signal; main service is marked unavailable; API enters stable 5xx | |
|
||||
| No load | Send `SIGKILL` to serve main process | serve main process exits instantly; request connections are interrupted; related child processes are cleaned up with no residue; GPU memory is released quickly | [#3725](https://github.com/vllm-project/vllm-omni/issues/3725) <br>[#43060](https://github.com/vllm-project/vllm/issues/43060) |
|
||||
| No load | Send `SIGTERM` to serve main process | serve enters graceful shutdown and stops serving; then exits and completes cleanup; GPU memory is released | |
|
||||
| No load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | Triggers serve shutdown path; service stops responding and becomes unavailable; related child processes exit and resources are released | |
|
||||
| No load | Send `SIGKILL` to all related processes | All related processes terminate immediately; service becomes unavailable at once; no residual processes remain; GPU memory is released quickly | |
|
||||
| No load | Send `SIGTERM` to all related processes | All processes enter exit flow and complete shutdown; service becomes unavailable; resource release is completed | |
|
||||
| Under load | Send `SIGKILL` to Worker | In-flight requests are hard interrupted (5xx/connection drop); main service becomes unavailable | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGTERM` to Worker | In-flight requests are canceled or fail fast; main service becomes unavailable | |
|
||||
| Under load | Send `SIGKILL` to serve main process | serve is hard-killed and current connections are interrupted; in-flight requests fail; after cleanup there are no residual processes and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGTERM` to serve main process | serve stops accepting new requests and executes shutdown flow; in-flight requests fail; no residue remains and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGINT` to serve main process (equivalent to `Ctrl+C`) | `Ctrl+C`-style serve shutdown; in-flight requests fail (5xx/connection interruption); service unavailable; after exit there is no residue and GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| Under load | Send `SIGKILL` to all related processes | All processes terminate instantly; all in-flight requests fail; service becomes unavailable immediately; GPU memory is released quickly | |
|
||||
| Under load | Send `SIGTERM` to all related processes | All processes exit gracefully; in-flight requests fail; service unavailable; GPU memory is released | [#3683](https://github.com/vllm-project/vllm-omni/issues/3683) |
|
||||
| OOM | Occupy all free GPU memory via an extra process | After OOM injection process starts, GPU memory is continuously saturated; service enters unavailable/degraded state and health check drops to 503; different request types (chat/speech, etc.) fail fast within a fixed time and return 500 (no hanging) | [#4285](https://github.com/vllm-project/vllm-omni/issues/4285) |
|
||||
|
||||
## Source of Conclusions
|
||||
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`.
|
||||
|
||||
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
|
||||
|
||||
The following full log is from a real run where a `SIGTERM` signal was sent to the serve root process.
|
||||
|
||||
```text
|
||||
Launching OmniServer with: /workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
Server ready on 127.0.0.1:60675 (OmniServer startup took 206.023s)
|
||||
OmniServer started successfully
|
||||
[reliability][process-kill] current_server_proc pid=556855 name=python3 cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[reliability][process-kill] current_server_proc pid=557118 name=python3 cmdline=/workspace/.venv/bin/python3 -c from multiprocessing.resource_tracker import main;main(57)
|
||||
[reliability][process-kill] current_server_proc pid=557119 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
[reliability][process-kill] current_server_proc pid=557122 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
[reliability][process-kill] current_server_proc pid=558201 name=VLLM::StageEngineCoreProc_noid_replica0_DP0 cmdline=VLLM::StageEngineCoreProc_noid_replica0_DP0
|
||||
|
||||
[reliability][process-kill] root-kill pid=556855 name=python3 signal=SIGTERM cmdline=/workspace/.venv/bin/python3 -m vllm_omni.entrypoints.cli.main serve Qwen/Qwen3-Omni-30B-A3B-Instruct --host 127.0.0.1 --port 60675 --omni --async-chunk --stage-init-timeout 600 --init-timeout 900 --log-stats --stage-configs-path vllm-omni/vllm_omni/deploy/qwen3_omni_moe.yaml
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [async_omni_engine.py:2397] [AsyncOmniEngine] Shutting down Orchestrator
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:351] [Orchestrator] Received shutdown signal
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [orchestrator.py:1359] [Orchestrator] Shutting down all 3 client(s)
|
||||
|
||||
[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
|
||||
[0;36m(StageEngineCoreProc pid=557119)[0;0m [32mINFO[0m [90m05-19 12:24:51[0m [90m[core.py:1265][0m Shutdown complete
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Shutting down
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Waiting for application shutdown.
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Application shutdown complete.
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO: Finished server process [556855]
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:51 [omni_base.py:463] [AsyncOmni] Shutting down
|
||||
[rank0]:[W519 12:24:53.427894833 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:24:56 [stage_pool.py:731] [StagePool] Stage 0 replica 0 shut down
|
||||
[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
|
||||
[0;36m(StageEngineCoreProc pid=557122)[0;0m [32mINFO[0m [90m05-19 12:24:56[0m [90m[core.py:1265][0m Shutdown complete
|
||||
[rank0]:[W519 12:24:57.569425065 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
|
||||
[0;36m(APIServer pid=556855)[0;0m INFO 05-19 12:25:00 [stage_pool.py:731] [StagePool] Stage 1 replica 0 shut down
|
||||
[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1242][0m Shutdown initiated (timeout=0)
|
||||
[0;36m(StageEngineCoreProc pid=558201)[0;0m [32mINFO[0m [90m05-19 12:25:00[0m [90m[core.py:1265][0m Shutdown complete
|
||||
[0;36m(APIServer pid=556855)[0;0m WARNING 05-19 12:25:01 [async_omni_engine.py:2406] [AsyncOmniEngine] Orchestrator thread did not exit in time
|
||||
[rank0]:[W519 12:25:01.867242016 ProcessGroupNCCL.cpp:1575] Warning: WARNING: destroy_process_group() was not called before program exit, which can leak resources. For more info, please see https://pytorch.org/docs/stable/distributed.html#shutdown (function operator())
|
||||
```
|
||||
|
||||
### How this maps to the matrix
|
||||
|
||||
- This log corresponds to the `Send SIGTERM to serve main process` scenario in the matrix.
|
||||
- The `root-kill ... signal=SIGTERM` line confirms the fault injection action is applied as expected.
|
||||
- The sequence from `Shutting down Orchestrator` to stage-level `Shutdown complete` and `Stage ... shut down` matches the expected graceful shutdown path and child-process cleanup behavior described in the matrix.
|
||||
|
||||
## Related Documentation
|
||||
|
||||
- For practical CI debugging procedures and triage guidance, see [CI Failures](../contributing/ci/failures.md).
|
||||
@@ -28,6 +28,50 @@ in deep DiT blocks.
|
||||
Legend: `✅` supported, `❌` unsupported, `⭕` not verified in this
|
||||
guide. FP8 on Ampere may use a weight-only path where available.
|
||||
|
||||
### Faster FP8 GEMM on Blackwell (quack)
|
||||
|
||||
On Blackwell (SM 100+), vLLM runs FP8 linears through the FlashInfer kernel, which
|
||||
applies the bias as a separate kernel after the GEMM. On the small GEMMs in video
|
||||
DiTs this bias add is a significant overhead. Installing the optional `quack` kernel
|
||||
lets vLLM-Omni fuse `alpha * (A @ B) + bias` into a single CuteDSL GEMM, recovering
|
||||
that overhead (e.g. HunyuanVideo-1.5 FP8 goes from slower-than-BF16 to faster).
|
||||
|
||||
```bash
|
||||
# CUDA 12.9
|
||||
pip install vllm-omni[quack]
|
||||
|
||||
# CUDA 13.x
|
||||
pip install 'quack-kernels[cu13]' --extra-index-url https://download.pytorch.org/whl/cu130
|
||||
```
|
||||
|
||||
It is enabled automatically once installed (no flag needed) and is **Blackwell-only**:
|
||||
on Hopper/Ada the CUTLASS FP8 kernel already fuses bias, so quack is not used there.
|
||||
Set `VLLM_OMNI_USE_QUACK_FP8=0` to force the FlashInfer path. If `quack-kernels` is
|
||||
not installed, FP8 still works — it just keeps the unfused FlashInfer path.
|
||||
|
||||
#### Compile cache and warmup
|
||||
|
||||
quack JIT-compiles its kernel once per distinct GEMM shape (tens of seconds, longer
|
||||
the first time across all autotuned configs). The compiled `.o` files are cached on
|
||||
disk and reused on later runs, so this is a one-time cost — **not** per request.
|
||||
|
||||
vLLM-Omni points that cache at `~/.cache/vllm_omni/quack` (override with
|
||||
`QUACK_CACHE_DIR`) instead of quack's default under `/tmp`, so it survives restarts.
|
||||
In containers, set `QUACK_CACHE_DIR` to a mounted/persistent path — or bake it into
|
||||
the image — so the first cold start does not recompile. The engine's startup dummy
|
||||
run already exercises the kernels, so with a warm cache the first real request is fast.
|
||||
|
||||
To pre-warm specific shapes (e.g. at image build time):
|
||||
|
||||
```python
|
||||
from vllm_omni.quantization.quack_fp8 import warmup_quack_fp8
|
||||
# (M, K, N) per linear; M = number of tokens for your resolution/frame count
|
||||
warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
```
|
||||
|
||||
> The PyPI package is `quack-kernels` (imported as `quack`); plain `pip install
|
||||
> quack` is an unrelated statistics library. Requires CUDA 12.9+ and Python 3.12.
|
||||
|
||||
## Model Type Support
|
||||
|
||||
### Diffusion Model (Qwen-Image, Wan2.2)
|
||||
@@ -52,17 +96,6 @@ guide. FP8 on Ampere may use a weight-only path where available.
|
||||
Audio encoder, vision encoder, talker, and code2wav stay in BF16 unless a
|
||||
model-specific guide says otherwise.
|
||||
|
||||
### MoT Diffusion Model (SenseNova-U1)
|
||||
|
||||
| Model | HF models | Scope | Status | Notes |
|
||||
|-------|-----------|-------|--------|-------|
|
||||
| SenseNova-U1 | `SenseNova/SenseNova-U1-8B-MoT` | Gen-path only (`*_mot_gen` layers) | Validated | Understanding-path layers stay in BF16; no `ignored_layers` needed |
|
||||
|
||||
SenseNova-U1 uses Mixture-of-Tokenizers (MoT) with duplicate weights per
|
||||
layer: understanding path (`qkv_proj`, `mlp`) and generation path
|
||||
(`qkv_proj_mot_gen`, `mlp_mot_gen`). When `quantization="fp8"` is set,
|
||||
the pipeline automatically applies FP8 only to gen-path GEMMs.
|
||||
|
||||
### Multi-Stage Diffusion Model (BAGEL, GLM-Image)
|
||||
|
||||
| Model | Scope | Status | Notes |
|
||||
|
||||
@@ -0,0 +1,14 @@
|
||||
{
|
||||
"prompt": "A robotic arm, primarily white with black joints and cables, is shown in a clean, modern indoor setting with a white tabletop. The arm, equipped with a gripper holding a small, light green pitcher, is positioned above a clear glass containing a reddish-brown liquid and a spoon. The robotic arm is in the process of pouring a transparent liquid into the glass. To the left of the pitcher, there is an opened jar with a similar reddish-brown substance visible through its transparent body. In the background, a vase with white flowers and a brown couch are partially visible, adding to the contemporary ambiance. The lighting is bright, casting soft shadows on the table. The robotic arm's movements are smooth and controlled, demonstrating precision in its task. As the video progresses, the robotic arm completes the pour, leaving the glass half-filled with the reddish-brown liquid. The jar remains untouched throughout the sequence, and the spoon inside the glass remains stationary. The other robotic arm on the right side also stays stationary throughout the video. The final frame captures the robotic arm with the pitcher finishing the pour, with the glass now filled to a higher level, while the pitcher is slightly tilted but still held securely by the gripper.",
|
||||
"negative_prompt": "",
|
||||
"vision_path": "https://github.com/nvidia-cosmos/cosmos-dependencies/raw/refs/heads/assets/cosmos3/inputs/vision/robot_pouring.mp4",
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"flow_shift": 10.0,
|
||||
"fps": 24,
|
||||
"condition_frame_indexes_vision": [0, 1],
|
||||
"condition_video_keep": "first"
|
||||
}
|
||||
@@ -143,6 +143,20 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
required=False,
|
||||
)
|
||||
parser.add_argument(
|
||||
"--width",
|
||||
type=int,
|
||||
default=None,
|
||||
metavar="W",
|
||||
help="Output image width in pixels. Default: None (pipeline's default).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--height",
|
||||
type=int,
|
||||
default=None,
|
||||
metavar="H",
|
||||
help="Output image height in pixels. Default: None (pipeline's default).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
@@ -236,8 +250,8 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--resolution",
|
||||
type=int,
|
||||
default=640,
|
||||
help="Bucket in (640, 1024) to determine the condition and output resolution",
|
||||
default=None,
|
||||
help="Bucket in (640, 1024) to determine the condition and output resolution. If width and height are not provided, this will be set to default 640.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
@@ -384,6 +398,15 @@ def parse_args() -> argparse.Namespace:
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
if args.resolution and (args.width or args.height):
|
||||
raise ValueError("--resolution and --width/--height cannot be specified together")
|
||||
if args.width is not None and args.width <= 0:
|
||||
raise ValueError("--width must be a positive integer")
|
||||
if args.height is not None and args.height <= 0:
|
||||
raise ValueError("--height must be a positive integer")
|
||||
if not args.width and not args.height and not args.resolution:
|
||||
args.resolution = 640
|
||||
|
||||
# Validate input images exist and load them
|
||||
input_images = []
|
||||
for image_path in args.image:
|
||||
@@ -464,6 +487,7 @@ def main():
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
|
||||
)
|
||||
print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
@@ -473,22 +497,26 @@ def main():
|
||||
omni.start_profile()
|
||||
|
||||
# Generate edited image
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
width=args.width,
|
||||
height=args.height,
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_2,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_outputs_per_prompt,
|
||||
layers=args.layers,
|
||||
)
|
||||
if args.resolution:
|
||||
sampling_params.resolution = args.resolution
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"multi_modal_data": {"image": input_image},
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_2,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_outputs_per_prompt,
|
||||
layers=args.layers,
|
||||
resolution=args.resolution,
|
||||
),
|
||||
sampling_params,
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -17,6 +17,7 @@ list of supported architectures across all modalities, see
|
||||
| CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | 2 (talker + code2wav) | ✓ | ✓ | — | 24 kHz |
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | dual-AR | ✓ | ✓ | — | 44.1 kHz |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | 2 (AR + DiT) | ✓ (required) | ✓ | — | 24 kHz |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | 2 (AR + audio VAE) | ✓ | ✓ | style / IP / dialect / TTA / podcast | 44.1 kHz |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | single (talker only) | — (caption-controlled) | — | style / IP / basic captions | 44.1 kHz |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | single (AR + codec) | ✓ (required) | ✓ | voice_clone, continuation | 48 kHz |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
@@ -159,6 +160,46 @@ Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
---
|
||||
|
||||
## Ming-omni-tts
|
||||
|
||||
Dense 0.5B two-stage TTS pipeline (`AR + flow` + audio VAE) at 44.1 kHz. The example covers style, IP voice, music-only generation, text-to-audio events, emotion, dialect, zero-shot cloning, podcast, speech+BGM, and speech+environment-sound cases.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case style \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Voice cloning
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case zero_shot \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Streaming
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case basic \
|
||||
--ref-audio /path/to/reference.wav \
|
||||
--streaming \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `style`, `ip`, `bgm`, and `tta` do not require reference audio.
|
||||
- Reference-audio cases use `--ref-audio`; `zero_shot` also requires `--ref-text`.
|
||||
- `podcast` uses multiple references via `--ref-audio-paths`.
|
||||
- Full case details live in [`ming_tts/README.md`](ming_tts/README.md).
|
||||
|
||||
---
|
||||
|
||||
## Ming-flash-omni-TTS
|
||||
|
||||
Standalone talker-only deployment of Ming-flash-omni-2.0 at 44.1 kHz. Voice is controlled through caption fields (`风格` / `IP` / `语速`/`基频`/`音量`) rather than reference audio.
|
||||
|
||||
@@ -0,0 +1,204 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline higgs-audio v3 inference example.
|
||||
|
||||
Runs Stage 0 (Qwen3 talker) + Stage 1 (HiggsAudio codec) end-to-end
|
||||
through the vLLM-Omni engine without going through the HTTP server, and
|
||||
saves a 24 kHz mono WAV per prompt.
|
||||
|
||||
Example (plain TTS):
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
--output-dir results/higgs_v3_wavs
|
||||
|
||||
Example (voice clone):
|
||||
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \\
|
||||
--texts "Hello world." \\
|
||||
--ref-audio path/to/reference.wav \\
|
||||
--ref-text "Transcript of the reference clip." \\
|
||||
--output-dir results/higgs_v3_wavs
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
|
||||
os.environ.setdefault("VLLM_USE_DEEP_GEMM", "0")
|
||||
os.environ.setdefault("VLLM_MOE_USE_DEEP_GEMM", "0")
|
||||
|
||||
import re
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
|
||||
from vllm_omni import Omni
|
||||
|
||||
SAMPLE_RATE = 24_000
|
||||
DEFAULT_TEXTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"Today is a beautiful day for a walk in the park.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
import argparse
|
||||
|
||||
parser = argparse.ArgumentParser(description="Offline higgs-audio v3 inference")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
default="bosonai/higgs-audio-v3-tts-4b",
|
||||
help="Stage-0 talker model id or path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--texts",
|
||||
type=str,
|
||||
nargs="+",
|
||||
default=list(DEFAULT_TEXTS),
|
||||
help="One or more plain-text prompts to synthesize.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
type=str,
|
||||
default="results/higgs_v3_wavs",
|
||||
help="Directory to write per-prompt WAV files.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override the deploy config path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-new-tokens",
|
||||
type=int,
|
||||
default=2048,
|
||||
help="Cap on Stage-0 codec frames per request.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Reference audio for voice clone (WAV/FLAC/MP3 path).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference audio. Optional but improves fidelity.",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _slugify(text: str) -> str:
|
||||
slug = re.sub(r"\s+", "_", text.strip().lower())
|
||||
slug = re.sub(r"[^a-z0-9_]+", "", slug)
|
||||
return slug[:48] or "prompt"
|
||||
|
||||
|
||||
def _extract_pcm(multimodal_output: dict) -> torch.Tensor:
|
||||
audio = multimodal_output.get("model_outputs")
|
||||
if audio is None:
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
raise ValueError(f"no audio key in multimodal_output: {list(multimodal_output.keys())}")
|
||||
if isinstance(audio, list):
|
||||
valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None]
|
||||
if not valid:
|
||||
raise ValueError("audio list is empty")
|
||||
return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0]
|
||||
return torch.as_tensor(audio).float().cpu().reshape(-1)
|
||||
|
||||
|
||||
def _pcm_to_int16(pcm: torch.Tensor) -> np.ndarray:
|
||||
arr = pcm.numpy()
|
||||
if arr.dtype.kind == "f":
|
||||
arr = np.clip(arr, -1.0, 1.0)
|
||||
arr = (arr * 32767.0).astype(np.int16)
|
||||
else:
|
||||
arr = arr.astype(np.int16)
|
||||
return arr
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
engine = Omni(model=args.model, deploy_config=args.deploy_config, trust_remote_code=True)
|
||||
|
||||
from transformers import AutoTokenizer
|
||||
|
||||
from vllm_omni.model_executor.models.higgs_audio_v3.higgs_audio_v3_tokenizer import (
|
||||
HiggsAudioV3TokenizerAdapter,
|
||||
apply_delay_pattern,
|
||||
encode_reference_audio,
|
||||
)
|
||||
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=True)
|
||||
adapter = HiggsAudioV3TokenizerAdapter(tokenizer)
|
||||
|
||||
# Load and encode reference audio once (if voice cloning)
|
||||
ref_codes_delayed: torch.Tensor | None = None
|
||||
if args.ref_audio is not None:
|
||||
ref_wav, ref_sr = sf.read(args.ref_audio, always_2d=False)
|
||||
if ref_wav.ndim == 2:
|
||||
ref_wav = ref_wav.mean(axis=1)
|
||||
ref_codes_raw = encode_reference_audio(ref_wav, int(ref_sr))
|
||||
ref_codes_delayed = apply_delay_pattern(ref_codes_raw)
|
||||
print(f"Reference : {args.ref_audio}")
|
||||
print(f"Ref codes : {ref_codes_raw.shape[0]} frames -> {ref_codes_delayed.shape[0]} delayed")
|
||||
if args.ref_text:
|
||||
print(f"Ref text : {args.ref_text}")
|
||||
|
||||
print(f"Model : {args.model}")
|
||||
print(f"Prompts : {len(args.texts)}")
|
||||
print(f"Output dir : {output_dir}")
|
||||
print(f"Voice clone : {'yes' if ref_codes_delayed is not None else 'no'}")
|
||||
|
||||
total_elapsed = 0.0
|
||||
total_dur = 0.0
|
||||
for text in args.texts:
|
||||
if ref_codes_delayed is not None:
|
||||
prompt_ids = adapter.build_prompt(
|
||||
text,
|
||||
num_ref_tokens=int(ref_codes_delayed.shape[0]),
|
||||
reference_text=args.ref_text,
|
||||
)
|
||||
prompt = {
|
||||
"prompt_token_ids": prompt_ids,
|
||||
"additional_information": {
|
||||
"audio_input_ids": ref_codes_delayed.to(torch.long),
|
||||
"audio_input_ids_mask": torch.ones(ref_codes_delayed.shape[0], dtype=torch.bool),
|
||||
},
|
||||
}
|
||||
else:
|
||||
prompt_ids = adapter.build_prompt(text)
|
||||
prompt = {"prompt_token_ids": prompt_ids}
|
||||
|
||||
t_start = time.perf_counter()
|
||||
outputs = engine.generate([prompt])
|
||||
elapsed = time.perf_counter() - t_start
|
||||
total_elapsed += elapsed
|
||||
|
||||
mm = outputs[0].outputs[0].multimodal_output
|
||||
pcm = _extract_pcm(mm)
|
||||
slug = _slugify(text)
|
||||
out_path = output_dir / f"{slug}.wav"
|
||||
sf.write(str(out_path), _pcm_to_int16(pcm), SAMPLE_RATE, format="WAV", subtype="PCM_16")
|
||||
dur = pcm.numel() / SAMPLE_RATE
|
||||
total_dur += dur
|
||||
print(f" {slug:<50} dur={dur:5.2f}s -> {out_path}")
|
||||
|
||||
rtf = total_elapsed / total_dur if total_dur > 0 else float("inf")
|
||||
print(f"Total infer : {total_elapsed:.2f}s total audio: {total_dur:.2f}s RTF: {rtf:.3f}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,248 @@
|
||||
# Ming-omni-tts Offline Inference
|
||||
|
||||
`end2end.py` runs Ming dense 0.5B end to end with vLLM-Omni. It uses the in-repo Ming prompt assembly helper directly, so the example request shape matches the real integration instead of a simplified wrapper.
|
||||
|
||||
## Files
|
||||
|
||||
| File | Purpose |
|
||||
|---|---|
|
||||
| `end2end.py` | Driver: CLI, case loading, prompt construction, orchestration (~150 lines) |
|
||||
| `cases.yaml` | All 11 built-in case definitions (prompt, text, instruction, ref-audio flags, flow controls) |
|
||||
| `runner.py` | Engine management and audio output (streaming + blocking paths) |
|
||||
|
||||
## Model Overview
|
||||
|
||||
Ming dense 0.5B is exposed here as a two-stage offline pipeline:
|
||||
|
||||
- **Stage 0**: Qwen2-based AR generation with Ming prompt formatting and inline flow controls
|
||||
- **Stage 1**: audio VAE decode to mono 44.1 kHz waveform
|
||||
|
||||
`config_ming_tts.py` adapts the checkpoint's HuggingFace config fields
|
||||
(LLM, DiT, aggregator, AudioVAE, token ids). `vllm_omni/deploy/ming_tts.yaml`
|
||||
selects the vLLM-Omni pipeline and stage runtime topology, including
|
||||
connectors, async chunking, memory limits, and sampling defaults.
|
||||
|
||||
The example supports both:
|
||||
|
||||
- **Blocking eager** via `vllm_omni/deploy/ming_tts.yaml`
|
||||
- **Async chunk eager** via `vllm_omni/deploy/ming_tts.yaml` (default `async_chunk: true`)
|
||||
|
||||
## Setup
|
||||
|
||||
Install vLLM-Omni with the platform requirements for your accelerator:
|
||||
|
||||
```bash
|
||||
uv pip install -e .
|
||||
```
|
||||
|
||||
The Ming offline example does not require a separate upstream Ming package.
|
||||
Reference-audio cases use the repo dependencies for audio loading,
|
||||
resampling, and CampPlus speaker extraction, including `soundfile`,
|
||||
`torchaudio`, and `onnxruntime-rocm` in the documented ROCm environment.
|
||||
|
||||
The tested ROCm environment is summarized in the
|
||||
[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
|
||||
|
||||
## Supported Cases
|
||||
|
||||
These cases cover the upstream dense 0.5B cookbook surface that maps cleanly onto the current vLLM-Omni example:
|
||||
|
||||
- `style`: zero-speaker style-conditioned speech
|
||||
- `ip`: zero-speaker IP voice generation
|
||||
- `bgm`: music-only generation
|
||||
- `tta`: text-to-audio event generation with FlowLoss controls
|
||||
- `emotion`: reference-audio speech with emotion control
|
||||
- `basic`: reference-audio speech with speed / pitch / volume control
|
||||
- `dialect`: reference-audio speech with dialect control
|
||||
- `zero_shot`: reference-audio cloning with explicit transcript
|
||||
- `podcast`: multi-reference dialogue generation with automatic speaker embedding extraction
|
||||
- `speech_bgm`: speech with background music conditioning
|
||||
- `speech_sound`: speech with environmental sound conditioning
|
||||
|
||||
## Quick Start
|
||||
|
||||
Run the zero-speaker style example:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case style \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Run zero-shot cloning with a transcript:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case zero_shot \
|
||||
--ref-audio /path/to/10002287-00000094.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Run emotion-controlled speech:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case emotion \
|
||||
--ref-audio /path/to/emotion_prompt.wav \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Run podcast generation with two reference clips:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case podcast \
|
||||
--ref-audio-paths /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
The script automatically extracts one 192-d speaker embedding per reference WAV using the Ming model's `campplus.onnx`.
|
||||
|
||||
If you already have precomputed multi-speaker embeddings, you can override extraction with:
|
||||
|
||||
```bash
|
||||
--speaker-embedding /path/to/podcast_speaker_embeddings.json
|
||||
```
|
||||
|
||||
where the JSON is a list of speaker embeddings, one 192-d vector per speaker.
|
||||
|
||||
Run text-to-audio event generation:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case tta \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Use async_chunk streaming:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case basic \
|
||||
--ref-audio /path/to/10002287-00000095.wav \
|
||||
--streaming \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
`--streaming` uses `AsyncOmni` and the async_chunk stage config. It currently
|
||||
supports one prompt per process invocation; use blocking mode for
|
||||
`--num-prompts > 1`.
|
||||
|
||||
Collect runtime stats and a manifest:
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/ming_tts/end2end.py \
|
||||
--case style \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--enforce-eager \
|
||||
--enable-stats \
|
||||
--stats-log-file output_audio/ming_style_pipeline.log \
|
||||
--metadata-json output_audio/ming_style_manifest.json
|
||||
```
|
||||
|
||||
## Reference Fixtures
|
||||
|
||||
The upstream Ming cookbook uses these public audio fixtures from `inclusionAI/Ming-omni-tts/data/wavs`:
|
||||
|
||||
- `10002287-00000094.wav` for zero-shot cloning
|
||||
- `10002287-00000095.wav` for `basic`
|
||||
- `emotion_prompt.wav` for `emotion`
|
||||
- `yue_prompt.wav` for `dialect`
|
||||
- `00000309-00000300.wav` for `speech_bgm` and `speech_sound`
|
||||
- `CTS-CN-F2F-2019-11-11-423-012-A.wav` and `CTS-CN-F2F-2019-11-11-423-012-B.wav` for `podcast`
|
||||
|
||||
## Validation Matrix
|
||||
|
||||
The repo-facing example is intended to cover the same dense TTS workflows used
|
||||
by the local Ming validation script:
|
||||
|
||||
| Case | Blocking `deploy/ming_tts.yaml` | Async chunk `deploy/ming_tts.yaml` | Extra inputs |
|
||||
|---|---:|---:|---|
|
||||
| `style` | Yes | Optional smoke test | none |
|
||||
| `ip` | Yes | Optional smoke test | none |
|
||||
| `bgm` | Yes | Optional smoke test | none |
|
||||
| `tta` | Yes | Optional smoke test | none |
|
||||
| `emotion` | Yes | Yes | `--ref-audio emotion_prompt.wav` |
|
||||
| `basic` | Yes | Yes | `--ref-audio 10002287-00000095.wav` |
|
||||
| `dialect` | Yes | Yes | `--ref-audio yue_prompt.wav` |
|
||||
| `zero_shot` | Yes | Yes | `--ref-audio 10002287-00000094.wav --ref-text ...` |
|
||||
| `podcast` | Yes | Yes | two `--ref-audio-paths` |
|
||||
| `speech_bgm` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
|
||||
| `speech_sound` | Yes | Yes | `--ref-audio 00000309-00000300.wav` |
|
||||
|
||||
## Validated Outputs
|
||||
|
||||
The following measurements are retained from an earlier L4 CUDA validation;
|
||||
they are not ROCm benchmark results. Default async_chunk matched blocking
|
||||
output frame counts and Stage-1 patch counts for every case:
|
||||
|
||||
| Case | Blocking frames / patches / sec | Async chunk frames / patches / sec |
|
||||
|---|---:|---:|
|
||||
| `style` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
|
||||
| `ip` | 183456 / 13 / 4.16 | 183456 / 13 / 4.16 |
|
||||
| `bgm` | 1326528 / 94 / 30.08 | 1326528 / 94 / 30.08 |
|
||||
| `tta` | 465696 / 33 / 10.56 | 465696 / 33 / 10.56 |
|
||||
| `emotion` | 324576 / 23 / 7.36 | 324576 / 23 / 7.36 |
|
||||
| `basic` | 211680 / 15 / 4.80 | 211680 / 15 / 4.80 |
|
||||
| `dialect` | 239904 / 17 / 5.44 | 239904 / 17 / 5.44 |
|
||||
| `zero_shot` | 409248 / 29 / 9.28 | 409248 / 29 / 9.28 |
|
||||
| `podcast` | 437472 / 31 / 9.92 | 437472 / 31 / 9.92 |
|
||||
| `speech_bgm` | 296352 / 21 / 6.72 | 296352 / 21 / 6.72 |
|
||||
| `speech_sound` | 352800 / 25 / 8.00 | 352800 / 25 / 8.00 |
|
||||
|
||||
## Key Arguments
|
||||
|
||||
| Argument | Description |
|
||||
|---|---|
|
||||
| `--model` | Hugging Face repo or local Ming checkpoint path |
|
||||
| `--deploy-config` | Deploy config YAML. Use `vllm_omni/deploy/ming_tts.yaml` |
|
||||
| `--case` | Built-in demo case |
|
||||
| `--ref-audio` | Single reference wav path for cloning-style cases |
|
||||
| `--ref-audio-paths` | Multiple reference wav paths, used by `podcast` |
|
||||
| `--ref-text` | Reference transcript. Required for `zero_shot` |
|
||||
| `--instructions` | Free-form Ming instruction string |
|
||||
| `--instruction-json` | Structured Ming instruction JSON |
|
||||
| `--speaker-embedding` | JSON file containing a 192-d speaker embedding |
|
||||
| `--extract-speaker-embeddings` | Force CampPlus speaker extraction from the provided reference audio paths |
|
||||
| `--max-decode-steps` | Override `ming_max_decode_steps` |
|
||||
| `--num-prompts` | Repeat the same case N times. Outputs are indexed when `N > 1` |
|
||||
| `--streaming` | Use `AsyncOmni` and async_chunk transport |
|
||||
| `--enforce-eager` | Recommended for Ming dense; non-eager is out of scope |
|
||||
| `--enable-stats` / `--log-stats` | Enable vLLM-Omni per-request stats logging |
|
||||
| `--stats-log-file` | Optional path for the stats log |
|
||||
| `--metadata-json` | Optional path for the run manifest JSON |
|
||||
| `--stage-init-timeout` | Per-stage initialization timeout in seconds |
|
||||
| `--init-timeout` | Total initialization timeout in seconds |
|
||||
| `--batch-timeout` | Batch timeout in seconds |
|
||||
| `--worker-backend` | `multi_process` or `ray` |
|
||||
| `--ray-address` | Ray cluster address when using `--worker-backend ray` |
|
||||
|
||||
## Output
|
||||
|
||||
- The script writes one mono 44.1 kHz WAV file per run
|
||||
- Default output directory: `output_audio/`
|
||||
- Default filename: `ming_<case>.wav`
|
||||
- When `--num-prompts > 1`, outputs are indexed as `ming_<case>_00000.wav`, `..._00001.wav`, etc.
|
||||
- When stats are enabled, the script can also write:
|
||||
- a stats log file such as `ming_style_pipeline.log`
|
||||
- a manifest JSON with per-output metadata, stage durations, peak memory info,
|
||||
and streaming client latency metrics when `--streaming` is used
|
||||
|
||||
## Notes
|
||||
|
||||
- `style` and `ip` are zero-speaker paths and do not require a reference clip
|
||||
- `emotion`, `basic`, `dialect`, `speech_bgm`, and `speech_sound` require one reference clip
|
||||
- `zero_shot` requires both `--ref-audio` and `--ref-text`
|
||||
- `podcast` requires at least two reference clips via `--ref-audio-paths`
|
||||
- `podcast` automatically extracts one speaker embedding per reference clip
|
||||
- `--speaker-embedding` may contain either one 192-d vector or a list of 192-d vectors
|
||||
- `--enforce-eager` was used for the validated runs
|
||||
- The earlier L4 validation used SDPA for the Ming audio VAE instead of
|
||||
FlashAttention2, which is the preferred default when available.
|
||||
@@ -0,0 +1,110 @@
|
||||
style:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我会一直在这里陪着你,直到你慢慢、慢慢地沉入那个最温柔的梦里……好吗?"
|
||||
instruction:
|
||||
风格: >-
|
||||
这是一种ASMR耳语,属于一种旨在引发特殊感官体验的创意风格。这个女性使用轻柔的普通话进行耳语,声音气音成分重。音量极低,紧贴麦克风,语速极慢,旨在制造触发听者颅内快感的声学刺激。
|
||||
use_zero_spk_emb: true
|
||||
max_decode_steps: 200
|
||||
|
||||
ip:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "这款产品的名字,叫变态坑爹牛肉丸。"
|
||||
instruction:
|
||||
IP: "灵小甄"
|
||||
use_zero_spk_emb: true
|
||||
max_decode_steps: 200
|
||||
|
||||
bgm:
|
||||
prompt: "Please generate music based on the following description.\n"
|
||||
text: " Genre: 电子舞曲. Mood: 自信 / 坚定. Instrument: 架子鼓. Theme: 节日. Duration: 30s."
|
||||
instruction: null
|
||||
use_zero_spk_emb: false
|
||||
max_decode_steps: 400
|
||||
|
||||
tta:
|
||||
prompt: "Please generate audio events based on given text.\n"
|
||||
text: "Thunder and a gentle rain"
|
||||
instruction: null
|
||||
use_zero_spk_emb: false
|
||||
max_decode_steps: 200
|
||||
cfg: 4.5
|
||||
sigma: 0.3
|
||||
temperature: 2.5
|
||||
|
||||
emotion:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我竟然抢到了陈奕迅的演唱会门票!太棒了!终于可以现场听一听他的歌声了!"
|
||||
instruction:
|
||||
情感: "高兴"
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
basic:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "简单地说,这相当于惠普把消费领域市场拱手相让了。"
|
||||
instruction:
|
||||
语速: "快速"
|
||||
基频: "中"
|
||||
音量: "中"
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
dialect:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我觉得社会企业同个人都有责任"
|
||||
instruction:
|
||||
方言: "广粤话"
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
zero_shot:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。"
|
||||
instruction: null
|
||||
requires_ref_audio: true
|
||||
requires_ref_text: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
podcast:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: " speaker_1:你可以说一下,就大概说一下,可能虽然我也不知道,我看过那部电影没有。\n speaker_2:就是那个叫什么,变相一节课的嘛。\n speaker_1:嗯。\n speaker_2:一部搞笑的电影。\n speaker_1:一部搞笑的。\n"
|
||||
instruction: null
|
||||
prompt_text: " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试,做个,当服务员还要去笔试了\n speaker_2:对啊,这真的很奇怪,就是 单纯的因,单纯自己工资不高,只是因为可能人家那个店比较出名一点,就对你苛刻要求\n"
|
||||
requires_ref_audio_count: 2
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
speech_bgm:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
|
||||
instruction:
|
||||
BGM:
|
||||
Genre: "当代古典音乐."
|
||||
Mood: "温暖 / 友善."
|
||||
Instrument: "电吉他"
|
||||
Theme: "节日."
|
||||
SNR: 10.0
|
||||
ENV: null
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
|
||||
speech_sound:
|
||||
prompt: "Please generate speech based on the following description.\n"
|
||||
text: "此次业绩下滑原因,可归结为企业停止服务某些品牌,而带来的负面影响。"
|
||||
instruction:
|
||||
BGM:
|
||||
ENV: "Birds chirping"
|
||||
SNR: 10.0
|
||||
Genre: null
|
||||
Mood: null
|
||||
Instrument: null
|
||||
Theme: null
|
||||
requires_ref_audio: true
|
||||
auto_extract_speaker_embeddings: true
|
||||
max_decode_steps: 200
|
||||
@@ -0,0 +1,200 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference demo for Ming-omni-tts via vLLM Omni."""
|
||||
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
import torchaudio
|
||||
import yaml
|
||||
from transformers import AutoTokenizer
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
KEY_CFG,
|
||||
KEY_MAX_DECODE_STEPS,
|
||||
KEY_SIGMA,
|
||||
KEY_TEMPERATURE,
|
||||
SAMPLE_RATE,
|
||||
)
|
||||
from vllm_omni.model_executor.models.ming_tts.prompt_assembly import build_ming_dense_prompt
|
||||
from vllm_omni.model_executor.models.ming_tts.speaker_extractor import MingSpeakerEmbeddingExtractor
|
||||
|
||||
try:
|
||||
from .runner import (
|
||||
build_manifest,
|
||||
build_sampling_params,
|
||||
resolve_metadata_json,
|
||||
resolve_stats_log_file,
|
||||
run_generation,
|
||||
)
|
||||
except ImportError:
|
||||
from runner import (
|
||||
build_manifest,
|
||||
build_sampling_params,
|
||||
resolve_metadata_json,
|
||||
resolve_stats_log_file,
|
||||
run_generation,
|
||||
)
|
||||
|
||||
_DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
|
||||
_REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
_DEFAULT_DEPLOY_CONFIG = str(_REPO_ROOT / "vllm_omni/deploy/ming_tts.yaml")
|
||||
_CASES_FILE = Path(__file__).with_name("cases.yaml")
|
||||
|
||||
CASE_DEFAULTS = yaml.safe_load(_CASES_FILE.read_text(encoding="utf-8")) or {}
|
||||
if not CASE_DEFAULTS:
|
||||
raise RuntimeError(f"Empty or missing case definitions in {_CASES_FILE}")
|
||||
|
||||
|
||||
def _build_parser() -> FlexibleArgumentParser:
|
||||
p = FlexibleArgumentParser(description="Offline Ming-omni-tts example")
|
||||
p.add_argument("--model", default=_DEFAULT_MODEL, help="Model name or local path")
|
||||
p.add_argument("--deploy-config", default=None, help="Deploy config YAML; auto-selected when omitted")
|
||||
p.add_argument("--case", choices=sorted(CASE_DEFAULTS), default="style", help="Built-in demo case")
|
||||
p.add_argument("--text", default=None, help="Override case text")
|
||||
p.add_argument("--prompt", default=None, help="Override the system prompt prefix")
|
||||
p.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
|
||||
p.add_argument(
|
||||
"--instruction-json", default=None, help='Structured Ming instruction JSON, e.g. \'{"方言":"广粤话"}\''
|
||||
)
|
||||
p.add_argument("--ref-audio", default=None, help="Single reference audio path for cloning-style cases")
|
||||
p.add_argument("--ref-audio-paths", nargs="+", default=None, help="Multiple reference audio paths (podcast)")
|
||||
p.add_argument("--ref-text", default=None, help="Reference transcript; required for zero_shot")
|
||||
p.add_argument("--speaker-embedding", default=None, help="Path to a JSON speaker embedding file")
|
||||
p.add_argument(
|
||||
"--extract-speaker-embeddings",
|
||||
action="store_true",
|
||||
help="Extract speaker embeddings from ref audio via campplus.onnx",
|
||||
)
|
||||
p.add_argument("--max-decode-steps", type=int, default=None, help="Override ming_max_decode_steps")
|
||||
p.add_argument("--output-dir", default="output_audio", help="Directory for output wav files")
|
||||
p.add_argument("--output-name", default=None, help="Output wav filename")
|
||||
p.add_argument("--num-prompts", type=int, default=1, help="Repeat the same prompt N times")
|
||||
p.add_argument("--streaming", action="store_true", help="Use AsyncOmni with async_chunk streaming")
|
||||
p.add_argument("--trust-remote-code", action="store_true")
|
||||
p.add_argument("--enforce-eager", action="store_true")
|
||||
p.add_argument(
|
||||
"--log-stats", "--enable-stats", dest="log_stats", action="store_true", help="Enable Omni stats logging"
|
||||
)
|
||||
p.add_argument("--stats-log-file", default=None, help="Path for the Omni stats log file")
|
||||
p.add_argument("--metadata-json", default=None, help="Path for the run manifest JSON")
|
||||
p.add_argument("--stage-init-timeout", type=int, default=300, help="Per-stage init timeout (s)")
|
||||
p.add_argument("--init-timeout", type=int, default=600, help="Total init timeout (s)")
|
||||
p.add_argument("--batch-timeout", type=int, default=5, help="Batch timeout (s)")
|
||||
p.add_argument("--shm-threshold-bytes", type=int, default=65536)
|
||||
p.add_argument("--worker-backend", default="multi_process", choices=["multi_process", "ray"])
|
||||
p.add_argument("--ray-address", default=None, help="Ray cluster address (--worker-backend ray)")
|
||||
return p
|
||||
|
||||
|
||||
def _finalize_args(args) -> None:
|
||||
if args.instructions is not None and args.instruction_json is not None:
|
||||
raise RuntimeError("Use either --instructions or --instruction-json, not both")
|
||||
if args.num_prompts < 1:
|
||||
raise RuntimeError("--num-prompts must be at least 1")
|
||||
if args.streaming and args.num_prompts != 1:
|
||||
raise RuntimeError("--streaming currently supports exactly one prompt")
|
||||
if args.deploy_config is None:
|
||||
args.deploy_config = _DEFAULT_DEPLOY_CONFIG
|
||||
|
||||
|
||||
def _load_waveform(path: str) -> torch.Tensor:
|
||||
samples, sr = sf.read(path, dtype="float32")
|
||||
wav = torch.as_tensor(samples, dtype=torch.float32)
|
||||
if wav.ndim == 2:
|
||||
wav = wav.mean(dim=1)
|
||||
wav = wav.reshape(1, -1)
|
||||
if int(sr) != SAMPLE_RATE:
|
||||
wav = torchaudio.functional.resample(wav, int(sr), SAMPLE_RATE)
|
||||
return wav
|
||||
|
||||
|
||||
def _load_speaker_embedding(path: str) -> torch.Tensor:
|
||||
return torch.as_tensor(json.loads(Path(path).read_text(encoding="utf-8")), dtype=torch.float32)
|
||||
|
||||
|
||||
def _ref_audio_paths(args) -> list[str]:
|
||||
if args.ref_audio is not None and args.ref_audio_paths is not None:
|
||||
raise RuntimeError("Use either --ref-audio or --ref-audio-paths, not both")
|
||||
if args.ref_audio_paths is not None:
|
||||
return list(args.ref_audio_paths)
|
||||
return [args.ref_audio] if args.ref_audio else []
|
||||
|
||||
|
||||
def _resolve_reference_inputs(args, case: dict, paths: list[str]):
|
||||
required = int(case.get("requires_ref_audio_count", 0))
|
||||
if required > 0 and len(paths) < required:
|
||||
raise RuntimeError(f"Case '{args.case}' needs {required} ref audio paths via --ref-audio-paths")
|
||||
if required <= 0 and case.get("requires_ref_audio") and not paths:
|
||||
raise RuntimeError(f"--ref-audio required for case '{args.case}'")
|
||||
if not paths:
|
||||
return None
|
||||
wavs = [_load_waveform(p) for p in paths]
|
||||
return wavs[0] if len(wavs) == 1 else wavs
|
||||
|
||||
|
||||
def _resolve_speaker_embedding(args, case: dict, paths: list[str]):
|
||||
if args.speaker_embedding:
|
||||
return _load_speaker_embedding(args.speaker_embedding)
|
||||
if not (case.get("auto_extract_speaker_embeddings") or args.extract_speaker_embeddings) or not paths:
|
||||
return None
|
||||
embs = MingSpeakerEmbeddingExtractor(args.model).extract_many(paths)
|
||||
if not embs:
|
||||
raise RuntimeError("Speaker extraction produced no embeddings")
|
||||
return embs[0] if len(embs) == 1 else torch.stack(embs, dim=0)
|
||||
|
||||
|
||||
def _build_prompt_payload(tokenizer, args):
|
||||
case = CASE_DEFAULTS[args.case]
|
||||
paths = _ref_audio_paths(args)
|
||||
prompt_text = args.ref_text if args.ref_text is not None else case.get("prompt_text")
|
||||
if case.get("requires_ref_text") and not prompt_text:
|
||||
raise RuntimeError(f"--ref-text required for case '{args.case}'")
|
||||
prompt_waveform = _resolve_reference_inputs(args, case, paths) if prompt_text is not None else None
|
||||
speaker_embedding = _resolve_speaker_embedding(args, case, paths)
|
||||
use_zero_spk_emb = bool(case.get("use_zero_spk_emb")) and prompt_waveform is None and speaker_embedding is None
|
||||
runtime_controls = {KEY_MAX_DECODE_STEPS: args.max_decode_steps or case["max_decode_steps"]}
|
||||
for key, field in [(KEY_CFG, "cfg"), (KEY_SIGMA, "sigma"), (KEY_TEMPERATURE, "temperature")]:
|
||||
if field in case:
|
||||
runtime_controls[key] = case[field]
|
||||
instruction = (
|
||||
json.loads(args.instruction_json) if args.instruction_json else (args.instructions or case.get("instruction"))
|
||||
)
|
||||
return build_ming_dense_prompt(
|
||||
tokenizer,
|
||||
prompt=args.prompt or case["prompt"],
|
||||
text=args.text or case["text"],
|
||||
runtime_controls=runtime_controls,
|
||||
instruction=instruction,
|
||||
prompt_text=prompt_text,
|
||||
prompt_waveform=prompt_waveform,
|
||||
speaker_embedding=speaker_embedding,
|
||||
use_zero_spk_emb=use_zero_spk_emb,
|
||||
)
|
||||
|
||||
|
||||
def main():
|
||||
args = _build_parser().parse_args()
|
||||
_finalize_args(args)
|
||||
tokenizer = AutoTokenizer.from_pretrained(args.model, trust_remote_code=False)
|
||||
prompt_payload = _build_prompt_payload(tokenizer, args)
|
||||
case = CASE_DEFAULTS[args.case]
|
||||
sampling_params_list = build_sampling_params(args.max_decode_steps or case["max_decode_steps"])
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
stats_log_file = resolve_stats_log_file(args)
|
||||
summaries = run_generation(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
|
||||
metadata_json = resolve_metadata_json(args)
|
||||
manifest = build_manifest(args, prompt_payload, stats_log_file, summaries)
|
||||
if metadata_json is not None:
|
||||
Path(metadata_json).write_text(json.dumps(manifest, ensure_ascii=False, indent=2), encoding="utf-8")
|
||||
print(f"Saved run manifest to {metadata_json}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
main()
|
||||
@@ -0,0 +1,277 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import time
|
||||
import uuid
|
||||
import wave
|
||||
from pathlib import Path
|
||||
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
|
||||
from vllm_omni import AsyncOmni, Omni
|
||||
from vllm_omni.model_executor.models.ming_tts.config_ming_tts import (
|
||||
KEY_SPEAKER_EMBEDDING,
|
||||
SAMPLE_RATE,
|
||||
TEXT_EOS_TOKEN_ID,
|
||||
)
|
||||
|
||||
|
||||
def coerce_audio_tensor(audio, *, async_chunk: bool) -> torch.Tensor:
|
||||
if isinstance(audio, list):
|
||||
parts = []
|
||||
for item in audio:
|
||||
tensor = torch.as_tensor(item, dtype=torch.float32).reshape(-1)
|
||||
if tensor.numel() > 0:
|
||||
parts.append(tensor)
|
||||
if not parts:
|
||||
return torch.zeros((0,), dtype=torch.float32)
|
||||
return torch.cat(parts, dim=0)
|
||||
|
||||
return torch.as_tensor(audio, dtype=torch.float32).reshape(-1)
|
||||
|
||||
|
||||
def resolve_sr(sr) -> int:
|
||||
if isinstance(sr, list):
|
||||
sr = sr[-1]
|
||||
if hasattr(sr, "item"):
|
||||
return int(sr.item())
|
||||
return int(sr)
|
||||
|
||||
|
||||
def extract_sample_rate(multimodal_output: dict) -> int:
|
||||
sr = multimodal_output.get("sr")
|
||||
if sr is None:
|
||||
raise RuntimeError("Expected multimodal_output['sr']")
|
||||
return resolve_sr(sr)
|
||||
|
||||
|
||||
def write_wav(path: str, audio: torch.Tensor, sample_rate: int) -> None:
|
||||
audio = audio.clamp(-1.0, 1.0)
|
||||
pcm16 = (audio * 32767.0).round().to(torch.int16).cpu().numpy()
|
||||
with wave.open(path, "wb") as wav_file:
|
||||
wav_file.setnchannels(1)
|
||||
wav_file.setsampwidth(2)
|
||||
wav_file.setframerate(int(sample_rate))
|
||||
wav_file.writeframes(pcm16.tobytes())
|
||||
|
||||
|
||||
def request_index(request_id: str | None, fallback: int) -> int:
|
||||
try:
|
||||
return int(request_id)
|
||||
except (TypeError, ValueError):
|
||||
if isinstance(request_id, str):
|
||||
head = request_id.split("_", 1)[0]
|
||||
if head.isdigit():
|
||||
return int(head)
|
||||
return fallback
|
||||
|
||||
|
||||
def audio_summary(audio: torch.Tensor, sample_rate: int) -> dict:
|
||||
waveform = audio.detach().cpu().reshape(-1).to(torch.float32)
|
||||
return {
|
||||
"sample_rate": int(sample_rate),
|
||||
"num_samples": int(waveform.numel()),
|
||||
"duration_seconds": float(waveform.numel()) / float(sample_rate),
|
||||
"max_abs_amplitude": float(waveform.abs().max().item()) if waveform.numel() > 0 else 0.0,
|
||||
}
|
||||
|
||||
|
||||
def resolve_output_name(output_name: str | None, case: str, index: int, total: int) -> str:
|
||||
if total == 1:
|
||||
return output_name or f"ming_{case}.wav"
|
||||
base = Path(output_name or f"ming_{case}.wav")
|
||||
return f"{base.stem}_{index:05d}{base.suffix or '.wav'}"
|
||||
|
||||
|
||||
def resolve_stats_log_file(args) -> str | None:
|
||||
if not args.log_stats:
|
||||
return None
|
||||
if args.stats_log_file:
|
||||
return args.stats_log_file
|
||||
base = Path(args.output_name or f"ming_{args.case}.wav").stem
|
||||
return str(Path(args.output_dir) / f"{base}_pipeline.log")
|
||||
|
||||
|
||||
def resolve_metadata_json(args) -> str | None:
|
||||
if args.metadata_json:
|
||||
return args.metadata_json
|
||||
if args.log_stats:
|
||||
base = Path(args.output_name or f"ming_{args.case}.wav").stem
|
||||
return str(Path(args.output_dir) / f"{base}_manifest.json")
|
||||
return None
|
||||
|
||||
|
||||
def build_manifest(args, prompt_payload, stats_log_file: str | None, outputs: list[dict]) -> dict:
|
||||
additional_information = {}
|
||||
if isinstance(prompt_payload, dict):
|
||||
additional_information = dict(prompt_payload.get("additional_information", {}))
|
||||
return {
|
||||
"model": args.model,
|
||||
"case": args.case,
|
||||
"streaming": bool(args.streaming),
|
||||
"deploy_config": args.deploy_config,
|
||||
"enforce_eager": bool(args.enforce_eager),
|
||||
"num_prompts": int(args.num_prompts),
|
||||
"log_stats": bool(args.log_stats),
|
||||
"stats_log_file": stats_log_file,
|
||||
"prompt_text": additional_information.get("prompt_text"),
|
||||
"instruction": additional_information.get("instruction"),
|
||||
"speaker_embedding_shape": (
|
||||
list(additional_information[KEY_SPEAKER_EMBEDDING].shape)
|
||||
if KEY_SPEAKER_EMBEDDING in additional_information
|
||||
and hasattr(additional_information[KEY_SPEAKER_EMBEDDING], "shape")
|
||||
else None
|
||||
),
|
||||
"outputs": outputs,
|
||||
"generated_at_unix": time.time(),
|
||||
}
|
||||
|
||||
|
||||
def build_engine_kwargs(args, stats_log_file: str | None) -> dict:
|
||||
kwargs = {
|
||||
"model": args.model,
|
||||
"deploy_config": args.deploy_config,
|
||||
"enforce_eager": args.enforce_eager,
|
||||
"trust_remote_code": args.trust_remote_code,
|
||||
"log_stats": args.log_stats,
|
||||
"stage_init_timeout": args.stage_init_timeout,
|
||||
"init_timeout": args.init_timeout,
|
||||
"batch_timeout": args.batch_timeout,
|
||||
"shm_threshold_bytes": args.shm_threshold_bytes,
|
||||
"worker_backend": args.worker_backend,
|
||||
}
|
||||
if stats_log_file is not None:
|
||||
kwargs["log_file"] = stats_log_file
|
||||
if args.ray_address is not None:
|
||||
kwargs["ray_address"] = args.ray_address
|
||||
return kwargs
|
||||
|
||||
|
||||
def build_sampling_params(max_decode_steps: int) -> list[SamplingParams]:
|
||||
return [
|
||||
SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=max_decode_steps + 1,
|
||||
stop_token_ids=[int(TEXT_EOS_TOKEN_ID)],
|
||||
),
|
||||
SamplingParams(temperature=0.0, max_tokens=1),
|
||||
]
|
||||
|
||||
|
||||
async def run_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
|
||||
engine = AsyncOmni(**build_engine_kwargs(args, stats_log_file))
|
||||
try:
|
||||
all_audio_chunks = []
|
||||
accumulated_samples = 0
|
||||
chunk_idx = 0
|
||||
start_time = time.time()
|
||||
chunk_times = []
|
||||
ttfp_seconds = None
|
||||
final_stage_output = None
|
||||
async for stage_output in engine.generate(
|
||||
prompt=prompt_payload,
|
||||
request_id=str(uuid.uuid4()),
|
||||
sampling_params_list=sampling_params_list,
|
||||
):
|
||||
final_stage_output = stage_output
|
||||
multimodal_output = stage_output.multimodal_output or {}
|
||||
audio = multimodal_output.get("audio")
|
||||
if audio is None:
|
||||
continue
|
||||
|
||||
finished = stage_output.finished
|
||||
if isinstance(audio, torch.Tensor):
|
||||
if finished:
|
||||
audio_chunk = audio[accumulated_samples:].float().detach().cpu()
|
||||
else:
|
||||
audio_chunk = audio.float().detach().cpu()
|
||||
elif isinstance(audio, list):
|
||||
audio_chunk = torch.as_tensor(audio[chunk_idx], dtype=torch.float32).reshape(-1).cpu()
|
||||
else:
|
||||
audio_chunk = torch.as_tensor(audio, dtype=torch.float32).reshape(-1).cpu()
|
||||
|
||||
accumulated_samples += int(audio_chunk.numel())
|
||||
chunk_idx += 1
|
||||
if audio_chunk.numel() > 0:
|
||||
now = time.time()
|
||||
if ttfp_seconds is None:
|
||||
ttfp_seconds = now - start_time
|
||||
chunk_times.append(now)
|
||||
all_audio_chunks.append(audio_chunk)
|
||||
|
||||
if not all_audio_chunks:
|
||||
raise RuntimeError("Streaming Ming example produced no audio chunks")
|
||||
|
||||
waveform = torch.cat(all_audio_chunks, dim=0)
|
||||
output_name = resolve_output_name(args.output_name, args.case, 0, 1)
|
||||
output_path = str(output_dir / output_name)
|
||||
write_wav(output_path, waveform, SAMPLE_RATE)
|
||||
summary = {
|
||||
"request_id": getattr(final_stage_output, "request_id", None),
|
||||
"stage_id": getattr(final_stage_output, "stage_id", None),
|
||||
"output_path": output_path,
|
||||
"stage_durations": getattr(final_stage_output, "stage_durations", {}),
|
||||
"peak_memory_mb": getattr(final_stage_output, "peak_memory_mb", 0.0),
|
||||
"ttfp_seconds": ttfp_seconds,
|
||||
"mean_inter_chunk_seconds": (
|
||||
sum(t1 - t0 for t0, t1 in zip(chunk_times, chunk_times[1:])) / (len(chunk_times) - 1)
|
||||
if len(chunk_times) > 1
|
||||
else None
|
||||
),
|
||||
}
|
||||
summary.update(audio_summary(waveform, SAMPLE_RATE))
|
||||
print(f"Saved streaming output to {output_path}")
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
return [summary]
|
||||
finally:
|
||||
engine.shutdown()
|
||||
|
||||
|
||||
def run_non_streaming(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
|
||||
engine = Omni(**build_engine_kwargs(args, stats_log_file))
|
||||
try:
|
||||
outputs = engine.generate(
|
||||
prompts=[prompt_payload for _ in range(args.num_prompts)],
|
||||
sampling_params_list=sampling_params_list,
|
||||
py_generator=False,
|
||||
)
|
||||
summaries = []
|
||||
for fallback_index, output in enumerate(outputs):
|
||||
if output.final_output_type != "audio":
|
||||
continue
|
||||
multimodal_output = output.multimodal_output or {}
|
||||
waveform = coerce_audio_tensor(multimodal_output.get("audio"), async_chunk=False)
|
||||
sample_rate = extract_sample_rate(multimodal_output)
|
||||
output_name = resolve_output_name(
|
||||
args.output_name,
|
||||
args.case,
|
||||
request_index(output.request_id, fallback_index),
|
||||
args.num_prompts,
|
||||
)
|
||||
output_path = str(output_dir / output_name)
|
||||
write_wav(output_path, waveform, sample_rate)
|
||||
summary = {
|
||||
"request_id": output.request_id,
|
||||
"stage_id": output.stage_id,
|
||||
"output_path": output_path,
|
||||
"stage_durations": output.stage_durations,
|
||||
"peak_memory_mb": output.peak_memory_mb,
|
||||
}
|
||||
summary.update(audio_summary(waveform, sample_rate))
|
||||
summaries.append(summary)
|
||||
print(f"Saved output to {output_path}")
|
||||
print(json.dumps(summary, ensure_ascii=False, indent=2))
|
||||
if not summaries:
|
||||
raise RuntimeError("Non-streaming Ming example produced no audio outputs")
|
||||
return summaries
|
||||
finally:
|
||||
engine.close()
|
||||
|
||||
|
||||
def run_generation(args, prompt_payload, sampling_params_list, output_dir: Path, stats_log_file: str | None):
|
||||
if args.streaming:
|
||||
return asyncio.run(run_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file))
|
||||
return run_non_streaming(args, prompt_payload, sampling_params_list, output_dir, stats_log_file)
|
||||
@@ -212,7 +212,7 @@ def run_non_streaming(inputs, sampling_params_list, model_name, args, output_dir
|
||||
output_audio_dur = 0.0
|
||||
|
||||
for batch_idx, o in enumerate(outputs):
|
||||
audio_tensor = torch.cat(o.multimodal_output["audio"])
|
||||
audio_tensor = torch.cat((o.multimodal_output["audio"],))
|
||||
audio_array = audio_tensor.tolist()
|
||||
output_audio_dur += float(len(audio_array)) / 24000
|
||||
if args.write_audio:
|
||||
|
||||
@@ -0,0 +1,115 @@
|
||||
# Speech-To-Video
|
||||
|
||||
This example demonstrates how to deploy the Wan2.2 speech-to-video (S2V) model for online video generation using vLLM-Omni.
|
||||
|
||||
## Supported Models
|
||||
|
||||
| Model | Model ID |
|
||||
|-------|----------|
|
||||
| Wan2.2 S2V (14B) | `Wan-AI/Wan2.2-S2V-14B` |
|
||||
|
||||
## Start Server
|
||||
|
||||
### Basic Start
|
||||
|
||||
```bash
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size 2 \
|
||||
--flow-shift 3.0 \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--cache-backend cache_dit \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
### Start with Script
|
||||
|
||||
```bash
|
||||
bash run_server.sh
|
||||
```
|
||||
|
||||
The script allows overriding:
|
||||
- `MODEL` (default: `Wan-AI/Wan2.2-S2V-14B`)
|
||||
- `PORT` (default: `8091`)
|
||||
- `FLOW_SHIFT` (default: `3.0`)
|
||||
- `TP` (default: `2`)
|
||||
- `CACHE_BACKEND` (default: `cache_dit`)
|
||||
|
||||
### Start with 4 GPUs
|
||||
|
||||
```bash
|
||||
TP=4 bash run_server.sh
|
||||
```
|
||||
|
||||
## Sync API (Recommended for Testing)
|
||||
|
||||
`POST /v1/videos/sync` blocks until generation completes and returns the raw
|
||||
video bytes (`video/mp4`) directly in the response body.
|
||||
|
||||
```bash
|
||||
bash run_curl_speech_to_video.sh
|
||||
```
|
||||
|
||||
## Async Job API
|
||||
|
||||
`POST /v1/videos` creates an asynchronous job and returns immediately.
|
||||
|
||||
```bash
|
||||
no_proxy=127.0.0.1 \
|
||||
create_response=$(curl -sS -X POST http://127.0.0.1:8091/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
-F "prompt=A person singing" \
|
||||
-F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
|
||||
-F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
|
||||
-F "width=832" -F "height=480" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.5" \
|
||||
-F "fps=16")
|
||||
|
||||
video_id=$(echo "$create_response" | jq -r '.id')
|
||||
|
||||
# Poll until complete
|
||||
while true; do
|
||||
status=$(curl -s "http://127.0.0.1:8091/v1/videos/${video_id}" | jq -r '.status')
|
||||
if [ "$status" = "completed" ]; then
|
||||
break
|
||||
fi
|
||||
if [ "$status" = "failed" ]; then
|
||||
echo "Video generation failed"
|
||||
exit 1
|
||||
fi
|
||||
sleep 2
|
||||
done
|
||||
|
||||
# Download
|
||||
curl -L "http://127.0.0.1:8091/v1/videos/${video_id}/content" -o s2v_output.mp4
|
||||
```
|
||||
|
||||
## Request Parameters
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `prompt` | str | - | Text description of the desired video |
|
||||
| `image_reference` | JSON str | - | Image reference: `{"image_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `audio_reference` | JSON str | - | Audio reference: `{"audio_url": "..."}` — supports HTTP(s) URLs or base64 data URLs |
|
||||
| `width` | int | None | Video width in pixels |
|
||||
| `height` | int | None | Video height in pixels |
|
||||
| `fps` | int | None | Output video frame rate |
|
||||
| `num_inference_steps` | int | None | Number of denoising steps (40 recommended) |
|
||||
| `guidance_scale` | float | None | CFG guidance scale |
|
||||
| `seed` | int | None | Random seed for reproducibility |
|
||||
|
||||
## Notes
|
||||
|
||||
- S2V requires both a reference image (`image_reference`) and an audio reference
|
||||
(`audio_reference`). The generated video will show a person matching the reference
|
||||
image with lip movements synchronized to the audio.
|
||||
- `audio_reference` accepts a JSON string: `{"audio_url": "..."}` where the URL can be:
|
||||
- An HTTP/HTTPS URL (e.g., `https://example.com/audio.mp3`)
|
||||
- A base64 data URL (e.g., `data:audio/mp3;base64,...`)
|
||||
- `--model-class-name WanS2VPipeline` is required on the server to select the
|
||||
S2V pipeline (distinct from the T2V/I2V pipelines).
|
||||
- `--cache-backend cache_dit` enables DiT caching for ~2x speedup on cached steps.
|
||||
- Audio is muxed into the output MP4 automatically.
|
||||
- Always pass `fps=16` to ensure correct video/audio alignment.
|
||||
@@ -0,0 +1,44 @@
|
||||
#!/bin/bash
|
||||
# Wan2.2 S2V (speech-to-video) curl example using the sync video API.
|
||||
|
||||
set -euo pipefail
|
||||
|
||||
BASE_URL="${BASE_URL:-http://localhost:8091}"
|
||||
OUTPUT_PATH="${OUTPUT_PATH:-s2v_480p_serve.mp4}"
|
||||
IMAGE_URL="${IMAGE_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png}"
|
||||
AUDIO_URL="${AUDIO_URL:-https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3}"
|
||||
PROMPT="${PROMPT:-A person singing}"
|
||||
WIDTH="${WIDTH:-832}"
|
||||
HEIGHT="${HEIGHT:-480}"
|
||||
NUM_INFERENCE_STEPS="${NUM_INFERENCE_STEPS:-40}"
|
||||
GUIDANCE_SCALE="${GUIDANCE_SCALE:-4.5}"
|
||||
FPS="${FPS:-16}"
|
||||
|
||||
echo "Sending S2V request..."
|
||||
echo " Image URL: $IMAGE_URL"
|
||||
echo " Audio URL: $AUDIO_URL"
|
||||
echo " Prompt: $PROMPT"
|
||||
echo " Resolution: ${WIDTH}x${HEIGHT}"
|
||||
echo " Steps: $NUM_INFERENCE_STEPS"
|
||||
echo " FPS: $FPS"
|
||||
|
||||
IMAGE_REF_JSON="{\"image_url\": \"${IMAGE_URL}\"}"
|
||||
AUDIO_REF_JSON="{\"audio_url\": \"${AUDIO_URL}\"}"
|
||||
|
||||
no_proxy=127.0.0.1 \
|
||||
curl -X POST "${BASE_URL}/v1/videos/sync" \
|
||||
-F "prompt=${PROMPT}" \
|
||||
-F "image_reference=${IMAGE_REF_JSON}" \
|
||||
-F "audio_reference=${AUDIO_REF_JSON}" \
|
||||
-F "width=${WIDTH}" -F "height=${HEIGHT}" \
|
||||
-F "num_inference_steps=${NUM_INFERENCE_STEPS}" \
|
||||
-F "guidance_scale=${GUIDANCE_SCALE}" \
|
||||
-F "fps=${FPS}" \
|
||||
--output "${OUTPUT_PATH}"
|
||||
|
||||
if [ -f "$OUTPUT_PATH" ] && [ -s "$OUTPUT_PATH" ]; then
|
||||
echo "Saved video to ${OUTPUT_PATH} ($(du -h "$OUTPUT_PATH" | cut -f1))"
|
||||
else
|
||||
echo "ERROR: Output file is empty or missing"
|
||||
exit 1
|
||||
fi
|
||||
+29
@@ -0,0 +1,29 @@
|
||||
#!/bin/bash
|
||||
# Wan2.2 S2V (speech-to-video) online serving startup script
|
||||
|
||||
MODEL="${MODEL:-Wan-AI/Wan2.2-S2V-14B}"
|
||||
PORT="${PORT:-8091}"
|
||||
FLOW_SHIFT="${FLOW_SHIFT:-3.0}"
|
||||
TP="${TP:-2}"
|
||||
CACHE_BACKEND="${CACHE_BACKEND:-cache_dit}"
|
||||
|
||||
echo "Starting Wan2.2 S2V server..."
|
||||
echo "Model: $MODEL"
|
||||
echo "Port: $PORT"
|
||||
echo "Flow shift: $FLOW_SHIFT"
|
||||
echo "Tensor parallel size: $TP"
|
||||
echo "Cache backend: $CACHE_BACKEND"
|
||||
|
||||
CACHE_BACKEND_FLAG=""
|
||||
if [ "$CACHE_BACKEND" != "none" ]; then
|
||||
CACHE_BACKEND_FLAG="--cache-backend $CACHE_BACKEND"
|
||||
fi
|
||||
|
||||
VLLM_WORKER_MULTIPROC_METHOD=spawn \
|
||||
vllm serve "$MODEL" --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size "$TP" \
|
||||
--flow-shift "$FLOW_SHIFT" \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--port "$PORT" \
|
||||
$CACHE_BACKEND_FLAG
|
||||
@@ -16,6 +16,7 @@ For the full list of supported architectures across all modalities, see
|
||||
|---|---|---|---|---|---|
|
||||
| Fish Speech S2 Pro | `fishaudio/s2-pro` | ✓ (`ref_audio`+`ref_text`) | ✓ (PCM stream) | — | ✓ |
|
||||
| GLM-TTS | `zai-org/GLM-TTS` | ✓ (`ref_audio`+`ref_text`, required) | ✓ (PCM stream) | — | ✓ |
|
||||
| Ming-omni-tts | `inclusionAI/Ming-omni-tts-0.5B` | ✓ (`ref_audio` / `speaker_embedding`) | ✓ (PCM stream) | IP labels + structured `instructions` | — |
|
||||
| Ming-flash-omni-TTS | `Jonathan1909/Ming-flash-omni-2.0` | — (caption-controlled) | — | caption fields (`instructions`) | — |
|
||||
| MOSS-TTS-Nano | `OpenMOSS-Team/MOSS-TTS-Nano` | ✓ (`ref_audio` required) | ✓ (PCM stream) | — | ✓ |
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | ✓ | — | — | — |
|
||||
@@ -202,6 +203,51 @@ python fish_speech/gradio_demo.py --api-base http://localhost:8091 # if server
|
||||
|
||||
---
|
||||
|
||||
## Ming-omni-tts
|
||||
|
||||
Dense 0.5B two-stage TTS served through `/v1/audio/speech`. Ming uses the standard speech endpoint plus structured controls in `instructions`, `voice`, `language`, `ref_audio`, `ref_text`, and `speaker_embedding`.
|
||||
|
||||
### Launch
|
||||
```bash
|
||||
bash examples/online_serving/text_to_speech/ming_tts/run_server.sh
|
||||
```
|
||||
Equivalent manual command:
|
||||
```bash
|
||||
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--host 0.0.0.0 --port 8091 \
|
||||
--enforce-eager --omni
|
||||
```
|
||||
|
||||
### Sending requests
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "你好,这是 Ming 在线语音合成测试。"
|
||||
```
|
||||
|
||||
Structured dialect control:
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' \
|
||||
--ref-audio /path/to/yue_prompt.wav
|
||||
```
|
||||
|
||||
Zero-shot cloning:
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
|
||||
--ref-audio /path/to/10002287-00000094.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。"
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `run_curl.sh` keeps a small sanity subset; use the Ming README for the broader request cookbook.
|
||||
- Online serving is speech-shaped today; music-only `bgm` and text-to-audio `tta` remain offline examples.
|
||||
- Full request details live in [`ming_tts/README.md`](ming_tts/README.md).
|
||||
|
||||
---
|
||||
|
||||
## Ming-flash-omni-TTS
|
||||
|
||||
Standalone talker-only deployment of Ming-flash-omni-2.0. Voice is controlled through caption text passed via `instructions`.
|
||||
|
||||
@@ -51,13 +51,18 @@ def build_payload(
|
||||
response_format: str = "wav",
|
||||
stream: bool = False,
|
||||
) -> dict:
|
||||
"""Build the /v1/audio/speech request payload."""
|
||||
"""Build the /v1/audio/speech request payload.
|
||||
|
||||
Fish Speech has no built-in speakers; omit ``voice`` so the model
|
||||
generates with a random timbre (controllable via text tags like
|
||||
``[cool]``, ``[warm]``, etc.). When ``ref_audio`` is supplied we
|
||||
enter voice-cloning mode instead.
|
||||
"""
|
||||
if not text or not text.strip():
|
||||
raise gr.Error("Please enter text to synthesize.")
|
||||
|
||||
payload: dict = {
|
||||
"input": text.strip(),
|
||||
"voice": "default",
|
||||
"response_format": "pcm" if stream else response_format,
|
||||
"stream": stream,
|
||||
}
|
||||
@@ -193,7 +198,8 @@ def build_interface(api_base: str, stream_chunk_seconds: float = 0.5):
|
||||
|
||||
with gr.Accordion("Voice Cloning (optional)", open=False):
|
||||
gr.Markdown(
|
||||
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice."
|
||||
"Upload or link a short reference audio (10-30s) and provide its transcript to clone the voice. "
|
||||
"Leave empty for random timbre; use text tags like `[cool]`, `[warm]`, `[happy]` to steer the output."
|
||||
)
|
||||
ref_audio = gr.Audio(
|
||||
label="Reference Audio",
|
||||
|
||||
@@ -39,7 +39,6 @@ def run_tts(args) -> None:
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"voice": "default",
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,40 @@
|
||||
# Higgs-Audio V3 Online Serving
|
||||
|
||||
## Start the server
|
||||
|
||||
```bash
|
||||
# Default: GPU 0, port 8095
|
||||
./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
|
||||
|
||||
# Custom GPU / port
|
||||
PORT=8096 GPUS=0,1 ./examples/online_serving/text_to_speech/higgs_audio_v3/run_server.sh
|
||||
```
|
||||
|
||||
## Plain text TTS
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_batch \
|
||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
||||
```
|
||||
|
||||
## Voice clone
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_clone \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "transcript of the reference clip" \
|
||||
--prompts "Text to synthesize in the cloned voice."
|
||||
```
|
||||
|
||||
## curl example
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8095/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{"model": "higgs_audio_v3", "input": "Hello world."}' \
|
||||
--output hello.wav
|
||||
```
|
||||
@@ -0,0 +1,115 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Batch client for the higgs-audio v3 online server.
|
||||
|
||||
Sends prompts to ``/v1/audio/speech`` and saves the returned WAV files.
|
||||
|
||||
Usage (plain text -> speech):
|
||||
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_batch \
|
||||
--prompts "Hello world." "The quick brown fox jumps over the lazy dog."
|
||||
|
||||
Usage (voice clone):
|
||||
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--output-dir /tmp/higgs_v3_clone \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "the transcript of the reference clip" \
|
||||
--prompts "Hello world."
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
DEFAULT_PROMPTS = (
|
||||
"Hello world.",
|
||||
"The quick brown fox jumps over the lazy dog.",
|
||||
"Today is a beautiful day for a walk in the park.",
|
||||
"Innovation distinguishes between a leader and a follower.",
|
||||
)
|
||||
|
||||
|
||||
def _slug(text: str) -> str:
|
||||
import re
|
||||
|
||||
s = re.sub(r"\s+", "_", text.strip().lower())
|
||||
return re.sub(r"[^a-z0-9_]+", "", s)[:32] or "prompt"
|
||||
|
||||
|
||||
def main() -> int:
|
||||
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
parser.add_argument("--base-url", default="http://localhost:8095")
|
||||
parser.add_argument("--model", default="higgs_audio_v3")
|
||||
parser.add_argument("--prompts", nargs="+", default=list(DEFAULT_PROMPTS))
|
||||
parser.add_argument("--output-dir", type=Path, default=Path("/tmp/higgs_v3_batch"))
|
||||
parser.add_argument("--format", choices=("wav", "pcm"), default="wav")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=2048)
|
||||
parser.add_argument("--seed", type=int, default=42)
|
||||
parser.add_argument("--timeout-s", type=float, default=120.0)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
type=Path,
|
||||
default=None,
|
||||
help="Reference clip for voice clone (WAV/FLAC/MP3 path). Pair with --ref-text.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-text",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Transcript of the reference clip. Optional but improves fidelity.",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
ref_audio_data_url: str | None = None
|
||||
if args.ref_audio is not None:
|
||||
if not args.ref_audio.exists():
|
||||
print(f"ref-audio file not found: {args.ref_audio}", file=sys.stderr)
|
||||
return 2
|
||||
mime = "audio/wav" if args.ref_audio.suffix.lower() == ".wav" else "audio/mpeg"
|
||||
ref_b64 = base64.b64encode(args.ref_audio.read_bytes()).decode("ascii")
|
||||
ref_audio_data_url = f"data:{mime};base64,{ref_b64}"
|
||||
|
||||
try:
|
||||
import httpx
|
||||
except ImportError:
|
||||
print("this client needs `httpx`. Install with `pip install httpx`.", file=sys.stderr)
|
||||
return 2
|
||||
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
url = args.base_url.rstrip("/") + "/v1/audio/speech"
|
||||
failures = 0
|
||||
with httpx.Client(timeout=args.timeout_s) as client:
|
||||
for prompt in args.prompts:
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": prompt,
|
||||
"response_format": args.format,
|
||||
"max_new_tokens": args.max_new_tokens,
|
||||
"seed": args.seed,
|
||||
}
|
||||
if ref_audio_data_url is not None:
|
||||
payload["ref_audio"] = ref_audio_data_url
|
||||
if args.ref_text:
|
||||
payload["ref_text"] = args.ref_text
|
||||
resp = client.post(url, json=payload)
|
||||
if resp.status_code != 200:
|
||||
print(f"[FAIL] {prompt!r} -> {resp.status_code}: {resp.text[:200]}", file=sys.stderr)
|
||||
failures += 1
|
||||
continue
|
||||
suffix = ".wav" if args.format == "wav" else ".pcm"
|
||||
out = args.output_dir / f"{_slug(prompt)}{suffix}"
|
||||
out.write_bytes(resp.content)
|
||||
print(f"[ ok ] {prompt!r} -> {out} ({len(resp.content)} bytes)")
|
||||
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -0,0 +1,32 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for higgs-audio v3.
|
||||
#
|
||||
# Supports plain text TTS and voice cloning via /v1/audio/speech.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh # default port 8095, GPU 0
|
||||
# PORT=8096 GPUS=0,1 ./run_server.sh
|
||||
# MODEL=/path/to/local/checkpoint ./run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
MODEL="${MODEL:-bosonai/higgs-audio-v3-tts-4b}"
|
||||
PORT="${PORT:-8095}"
|
||||
GPUS="${GPUS:-0}"
|
||||
GPU_UTIL="${GPU_UTIL:-0.6}"
|
||||
|
||||
echo "Starting higgs-audio v3 server"
|
||||
echo " MODEL=$MODEL"
|
||||
echo " PORT=$PORT"
|
||||
echo " CUDA_VISIBLE_DEVICES=$GPUS"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="$GPUS" \
|
||||
VLLM_USE_DEEP_GEMM=0 \
|
||||
VLLM_MOE_USE_DEEP_GEMM=0 \
|
||||
vllm-omni serve "$MODEL" \
|
||||
--deploy-config vllm_omni/deploy/higgs_multimodal_qwen3.yaml \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--gpu-memory-utilization "$GPU_UTIL" \
|
||||
--trust-remote-code \
|
||||
--omni
|
||||
@@ -0,0 +1,113 @@
|
||||
# Ming-omni-tts Online Serving
|
||||
|
||||
Serve the dense `inclusionAI/Ming-omni-tts-0.5B` two-stage TTS model through
|
||||
the OpenAI-compatible `/v1/audio/speech` endpoint.
|
||||
|
||||
## Start Server
|
||||
|
||||
```bash
|
||||
vllm-omni serve inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
Or:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/text_to_speech/ming_tts
|
||||
./run_server.sh
|
||||
```
|
||||
|
||||
The tested ROCm environment is summarized in the
|
||||
[Ming recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/inclusionAI/Ming-omni-tts-0.5B.md).
|
||||
|
||||
## Send Requests
|
||||
|
||||
The Python client targets `http://localhost:8091/v1` with `api_key=EMPTY`; it
|
||||
does not call OpenAI's hosted API.
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "你好,这是 Ming 在线语音合成测试。" \
|
||||
--max-new-tokens 200
|
||||
```
|
||||
|
||||
Style or dialect controls can be plain text or Ming JSON. The upstream
|
||||
dialect example also uses `yue_prompt.wav` for speaker conditioning:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' \
|
||||
--ref-audio /path/to/yue_prompt.wav \
|
||||
--max-new-tokens 200
|
||||
```
|
||||
|
||||
When `--ref-audio` is supplied without `--ref-text`, the server extracts the
|
||||
Ming speaker embedding, matching upstream `use_spk_emb=True`, without using the
|
||||
audio as a zero-shot prompt.
|
||||
|
||||
Reference-audio cloning:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
|
||||
--ref-audio /path/to/10002287-00000094.wav \
|
||||
--ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--max-new-tokens 200
|
||||
```
|
||||
|
||||
Podcast-style multi-speaker prompt:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text " speaker_1:你可以说一下,就大概说一下,可能虽然我也不知道,我看过那部电影没有。
|
||||
speaker_2:就是那个叫什么,变相一节课的嘛。
|
||||
speaker_1:嗯。
|
||||
speaker_2:一部搞笑的电影。
|
||||
speaker_1:一部搞笑的。" \
|
||||
--ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-A.wav \
|
||||
--ref-audio /path/to/CTS-CN-F2F-2019-11-11-423-012-B.wav \
|
||||
--ref-text " speaker_1:并且我们还要进行每个月还要考核 笔试的话还要进行笔试,做个,当服务员还要去笔试了
|
||||
speaker_2:对啊,这真的很奇怪,就是 单纯的因,单纯自己工资不高,只是因为可能人家那个店比较出名一点,就对你苛刻要求"
|
||||
```
|
||||
|
||||
Streaming PCM:
|
||||
|
||||
```bash
|
||||
python openai_speech_client.py \
|
||||
--text "你好,这是流式输出测试。" \
|
||||
--stream \
|
||||
--output ming_output.pcm
|
||||
```
|
||||
|
||||
`run_curl.sh` keeps small smoke checks:
|
||||
|
||||
```bash
|
||||
./run_curl.sh basic
|
||||
REF_AUDIO=/path/to/reference.wav REF_TEXT="在此奉劝大家别乱打美白针。" ./run_curl.sh zero_shot
|
||||
./run_curl.sh stream
|
||||
```
|
||||
|
||||
## Request Fields
|
||||
|
||||
| Field | Ming meaning |
|
||||
|-------|--------------|
|
||||
| `input` | target text |
|
||||
| `instructions` | plain style text, or JSON object for structured Ming controls |
|
||||
| `voice` | Ming IP voice label unless it resolves to an uploaded speaker |
|
||||
| `language` | Ming `方言` control |
|
||||
| `ref_audio` | speaker reference; with `ref_text`, also supplies the prompt waveform |
|
||||
| `ref_text` | transcript enabling zero-shot or podcast prompt-latent conditioning |
|
||||
| `speaker_embedding` | 192-d Ming speaker embedding |
|
||||
| `max_new_tokens` | Ming `max_decode_steps` |
|
||||
|
||||
## Notes
|
||||
|
||||
- `ref_audio` accepts local paths through the client, remote URLs, `file://`,
|
||||
or `data:` URLs.
|
||||
- Non-streaming responses return WAV bytes; streaming responses return PCM.
|
||||
- Music-only `bgm` generation is offline-only until the API exposes Ming
|
||||
prompt-mode selection.
|
||||
@@ -0,0 +1,225 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech.
|
||||
|
||||
Examples:
|
||||
python openai_speech_client.py --text "你好,世界"
|
||||
python openai_speech_client.py --text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' --ref-audio yue_prompt.wav --max-new-tokens 200
|
||||
python openai_speech_client.py \
|
||||
--text "我们的愿景是构建未来服务业的数字化基础设施,为世界带来更多微小而美好的改变。" \
|
||||
--ref-audio 10002287-00000094.wav --ref-text "在此奉劝大家别乱打美白针。" \
|
||||
--max-new-tokens 200
|
||||
python openai_speech_client.py --text "speaker_1:你好。 speaker_2:你好。" \
|
||||
--ref-audio speaker_1.wav --ref-audio speaker_2.wav --ref-text "speaker_1:你好。 speaker_2:你好。"
|
||||
python openai_speech_client.py --text "你好,这是流式输出测试。" \
|
||||
--stream --output ming_output.pcm
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import json
|
||||
import os
|
||||
|
||||
import httpx
|
||||
|
||||
DEFAULT_API_BASE = "http://localhost:8091"
|
||||
DEFAULT_API_KEY = "EMPTY"
|
||||
DEFAULT_MODEL = "inclusionAI/Ming-omni-tts-0.5B"
|
||||
EXPECTED_SPEAKER_EMBEDDING_DIM = 192
|
||||
|
||||
|
||||
def encode_audio_to_base64(audio_path: str) -> str:
|
||||
"""Encode a local audio file to a base64 data URL."""
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
ext = audio_path.lower().rsplit(".", 1)[-1]
|
||||
mime_map = {
|
||||
"wav": "audio/wav",
|
||||
"mp3": "audio/mpeg",
|
||||
"flac": "audio/flac",
|
||||
"ogg": "audio/ogg",
|
||||
"aac": "audio/aac",
|
||||
}
|
||||
mime_type = mime_map.get(ext, "audio/wav")
|
||||
with open(audio_path, "rb") as f:
|
||||
audio_b64 = base64.b64encode(f.read()).decode("utf-8")
|
||||
return f"data:{mime_type};base64,{audio_b64}"
|
||||
|
||||
|
||||
def load_speaker_embedding(path: str) -> list[float]:
|
||||
"""Load and validate a 192-d Ming speaker embedding JSON file."""
|
||||
with open(path, encoding="utf-8") as f:
|
||||
data = json.load(f)
|
||||
|
||||
if not isinstance(data, list):
|
||||
raise ValueError("speaker_embedding file must contain a JSON list")
|
||||
if len(data) != EXPECTED_SPEAKER_EMBEDDING_DIM:
|
||||
raise ValueError(
|
||||
f"Ming dense speaker_embedding must have {EXPECTED_SPEAKER_EMBEDDING_DIM} values, got {len(data)}"
|
||||
)
|
||||
|
||||
values = []
|
||||
for index, value in enumerate(data):
|
||||
try:
|
||||
values.append(float(value))
|
||||
except (TypeError, ValueError) as exc:
|
||||
raise ValueError(f"speaker_embedding[{index}] must be a number, got {value!r}") from exc
|
||||
return values
|
||||
|
||||
|
||||
def build_instruction_payload(args) -> str | None:
|
||||
"""Return a string payload for the API `instructions` field."""
|
||||
if args.instructions and args.instruction_json:
|
||||
raise ValueError("Use either --instructions or --instruction-json, not both")
|
||||
if args.instruction_json:
|
||||
parsed = json.loads(args.instruction_json)
|
||||
return json.dumps(parsed, ensure_ascii=False)
|
||||
return args.instructions
|
||||
|
||||
|
||||
def validate_args(args) -> None:
|
||||
"""Fail fast on invalid combinations before hitting the server."""
|
||||
if args.ref_text and not args.ref_audio:
|
||||
raise ValueError("--ref-audio is required when --ref-text is provided")
|
||||
if args.speaker_embedding and args.ref_audio and len(args.ref_audio) > 1:
|
||||
raise ValueError("--speaker-embedding cannot be combined with multiple --ref-audio values")
|
||||
|
||||
|
||||
def run_tts(args) -> None:
|
||||
"""Generate speech via the OpenAI-compatible /v1/audio/speech API."""
|
||||
validate_args(args)
|
||||
|
||||
payload = {
|
||||
"model": args.model,
|
||||
"input": args.text,
|
||||
"response_format": args.response_format,
|
||||
}
|
||||
|
||||
if args.voice:
|
||||
payload["voice"] = args.voice
|
||||
if args.task_type:
|
||||
payload["task_type"] = args.task_type
|
||||
if args.dialect:
|
||||
payload["language"] = args.dialect
|
||||
|
||||
instructions = build_instruction_payload(args)
|
||||
if instructions:
|
||||
payload["instructions"] = instructions
|
||||
|
||||
if args.ref_audio:
|
||||
ref_audio = []
|
||||
for audio in args.ref_audio:
|
||||
if audio.startswith(("http://", "https://", "data:")):
|
||||
ref_audio.append(audio)
|
||||
else:
|
||||
ref_audio.append(encode_audio_to_base64(audio))
|
||||
payload["ref_audio"] = ref_audio[0] if len(ref_audio) == 1 else ref_audio
|
||||
if args.ref_text:
|
||||
payload["ref_text"] = args.ref_text
|
||||
if args.speaker_embedding:
|
||||
payload["speaker_embedding"] = load_speaker_embedding(args.speaker_embedding)
|
||||
if args.max_new_tokens:
|
||||
payload["max_new_tokens"] = args.max_new_tokens
|
||||
if args.stream:
|
||||
payload["stream"] = True
|
||||
payload["response_format"] = "pcm"
|
||||
|
||||
api_url = f"{args.api_base}/v1/audio/speech"
|
||||
headers = {
|
||||
"Content-Type": "application/json",
|
||||
"Authorization": f"Bearer {args.api_key}",
|
||||
}
|
||||
|
||||
print(f"Model: {args.model}")
|
||||
print(f"Text: {args.text}")
|
||||
print(f"Payload keys: {sorted(payload)}")
|
||||
|
||||
if args.stream:
|
||||
output_path = args.output or "ming_output.pcm"
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
with client.stream("POST", api_url, json=payload, headers=headers) as response:
|
||||
if response.status_code != 200:
|
||||
print(f"Error: {response.status_code}")
|
||||
print(response.read().decode())
|
||||
return
|
||||
with open(output_path, "wb") as f:
|
||||
for chunk in response.iter_bytes():
|
||||
f.write(chunk)
|
||||
print(f"Streamed PCM audio to: {output_path}")
|
||||
return
|
||||
|
||||
with httpx.Client(timeout=300.0) as client:
|
||||
response = client.post(api_url, json=payload, headers=headers)
|
||||
|
||||
if response.status_code != 200:
|
||||
print(f"Error: {response.status_code}")
|
||||
print(response.text)
|
||||
return
|
||||
|
||||
try:
|
||||
text = response.content.decode("utf-8")
|
||||
if text.startswith('{"error"'):
|
||||
print(f"Error: {text}")
|
||||
return
|
||||
except UnicodeDecodeError:
|
||||
pass
|
||||
|
||||
output_path = args.output or "ming_output.wav"
|
||||
with open(output_path, "wb") as f:
|
||||
f.write(response.content)
|
||||
print(f"Audio saved to: {output_path}")
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="OpenAI-compatible client for Ming-omni-tts via /v1/audio/speech")
|
||||
parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL")
|
||||
parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key")
|
||||
parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path")
|
||||
parser.add_argument("--text", required=True, help="Text to synthesize")
|
||||
parser.add_argument(
|
||||
"--task-type",
|
||||
default=None,
|
||||
choices=["CustomVoice", "VoiceDesign", "Base"],
|
||||
help="Optional compatibility task type. Ming accepts the same field but primarily uses prompt metadata.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--voice",
|
||||
default=None,
|
||||
help="Maps to Ming `IP` when using built-in character voices, or to an uploaded voice sample name",
|
||||
)
|
||||
parser.add_argument("--dialect", default=None, help="Maps to Ming `方言`")
|
||||
parser.add_argument("--instructions", default=None, help="Free-form Ming instruction string")
|
||||
parser.add_argument(
|
||||
"--instruction-json",
|
||||
default=None,
|
||||
help='Structured Ming instruction JSON, for example \'{"情感":"高兴"}\'',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--ref-audio",
|
||||
action="append",
|
||||
default=None,
|
||||
help="Reference audio path, URL, or data URL. Repeat for podcast-style multi-speaker prompts.",
|
||||
)
|
||||
parser.add_argument("--ref-text", default=None, help="Reference transcript for cloning")
|
||||
parser.add_argument(
|
||||
"--speaker-embedding", default=None, help="Path to a JSON file containing a 192-d speaker embedding"
|
||||
)
|
||||
parser.add_argument("--max-new-tokens", type=int, default=None, help="Override ming_max_decode_steps")
|
||||
parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output")
|
||||
parser.add_argument(
|
||||
"--response-format",
|
||||
default="wav",
|
||||
choices=["wav", "mp3", "flac", "pcm", "aac", "opus"],
|
||||
help="Audio format when not streaming",
|
||||
)
|
||||
parser.add_argument("--output", "-o", default=None, help="Output file path")
|
||||
args = parser.parse_args()
|
||||
try:
|
||||
run_tts(args)
|
||||
except (FileNotFoundError, ValueError, json.JSONDecodeError) as exc:
|
||||
raise SystemExit(f"Error: {exc}") from exc
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,76 @@
|
||||
#!/bin/bash
|
||||
set -euo pipefail
|
||||
|
||||
MODE="${1:-basic}"
|
||||
HOST="${HOST:-localhost}"
|
||||
PORT="${PORT:-8091}"
|
||||
MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
|
||||
API_URL="http://${HOST}:${PORT}/v1/audio/speech"
|
||||
TEXT="${TEXT:-你好,这是 Ming 在线语音合成测试。}"
|
||||
OUTPUT="${OUTPUT:-ming_output.wav}"
|
||||
STREAM_OUTPUT="${STREAM_OUTPUT:-ming_output.pcm}"
|
||||
REF_AUDIO="${REF_AUDIO:-}"
|
||||
REF_TEXT="${REF_TEXT:-}"
|
||||
|
||||
post_json() {
|
||||
local payload="$1"
|
||||
local output_path="$2"
|
||||
curl -X POST "$API_URL" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer EMPTY" \
|
||||
-d "$payload" \
|
||||
--output "$output_path"
|
||||
}
|
||||
|
||||
case "$MODE" in
|
||||
basic)
|
||||
post_json "{
|
||||
\"model\": \"${MODEL}\",
|
||||
\"input\": \"${TEXT}\",
|
||||
\"response_format\": \"wav\"
|
||||
}" "$OUTPUT"
|
||||
;;
|
||||
zero_shot)
|
||||
if [ -z "$REF_AUDIO" ] || [ -z "$REF_TEXT" ]; then
|
||||
echo "zero_shot requires REF_AUDIO and REF_TEXT" >&2
|
||||
exit 1
|
||||
fi
|
||||
python - <<'PY' > /tmp/ming_zero_shot_payload.json
|
||||
import base64
|
||||
import json
|
||||
import mimetypes
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
path = Path(os.environ["REF_AUDIO"])
|
||||
mime_type = mimetypes.guess_type(path.name)[0] or "audio/wav"
|
||||
payload = {
|
||||
"model": os.environ["MODEL"],
|
||||
"input": os.environ["TEXT"],
|
||||
"ref_audio": f"data:{mime_type};base64,{base64.b64encode(path.read_bytes()).decode('utf-8')}",
|
||||
"ref_text": os.environ["REF_TEXT"],
|
||||
"response_format": "wav",
|
||||
}
|
||||
print(json.dumps(payload, ensure_ascii=False))
|
||||
PY
|
||||
curl -X POST "$API_URL" \
|
||||
-H "Content-Type: application/json" \
|
||||
-H "Authorization: Bearer EMPTY" \
|
||||
--data-binary @/tmp/ming_zero_shot_payload.json \
|
||||
--output "$OUTPUT"
|
||||
rm -f /tmp/ming_zero_shot_payload.json
|
||||
;;
|
||||
stream)
|
||||
post_json "{
|
||||
\"model\": \"${MODEL}\",
|
||||
\"input\": \"${TEXT}\",
|
||||
\"stream\": true,
|
||||
\"response_format\": \"pcm\"
|
||||
}" "$STREAM_OUTPUT"
|
||||
;;
|
||||
*)
|
||||
echo "Unknown mode: $MODE" >&2
|
||||
echo "Supported sanity checks: basic, zero_shot, stream" >&2
|
||||
exit 1
|
||||
;;
|
||||
esac
|
||||
@@ -0,0 +1,25 @@
|
||||
#!/bin/bash
|
||||
# Launch vLLM-Omni server for Ming-omni-tts.
|
||||
#
|
||||
# Usage:
|
||||
# ./run_server.sh
|
||||
# PORT=8000 ./run_server.sh
|
||||
|
||||
set -e
|
||||
|
||||
DIR="$(cd "$(dirname "$0")" && pwd)"
|
||||
ROOT="$(cd "$DIR/../../../.." && pwd)"
|
||||
|
||||
MODEL="${MODEL:-inclusionAI/Ming-omni-tts-0.5B}"
|
||||
PORT="${PORT:-8091}"
|
||||
DEPLOY_CONFIG="${DEPLOY_CONFIG:-$ROOT/vllm_omni/deploy/ming_tts.yaml}"
|
||||
|
||||
echo "Starting Ming-omni-tts server with model: $MODEL"
|
||||
echo "Deploy config: $DEPLOY_CONFIG"
|
||||
|
||||
vllm-omni serve "$MODEL" \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--host 0.0.0.0 \
|
||||
--port "$PORT" \
|
||||
--enforce-eager \
|
||||
--omni
|
||||
@@ -408,7 +408,10 @@ def create_app(api_base: str):
|
||||
await resp.aclose()
|
||||
await client.aclose()
|
||||
|
||||
return StreamingResponse(relay(), media_type="application/octet-stream")
|
||||
# Use audio/pcm to bypass browsers (e.g. Chromium) MIME type sniffing, which
|
||||
# buffers application/octet-stream responses before delivering chunks to
|
||||
# fetch().body.getReader(), inflating TTFP to total_ms.
|
||||
return StreamingResponse(relay(), media_type="audio/pcm")
|
||||
|
||||
# ── Gradio UI ────────────────────────────────────────────────
|
||||
voices = fetch_voices(api_base)
|
||||
|
||||
@@ -0,0 +1,247 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Verify a ModelOpt FP8 diffusers checkpoint exported by
|
||||
quantize_wan2_2_modelopt_fp8.py / quantize_hunyuanvideo_15_modelopt_fp8.py.
|
||||
|
||||
Three checks:
|
||||
A. transformer/config.json has a sane quantization_config block.
|
||||
B. transformer/*.safetensors contains FP8 (float8_e4m3fn) quantized tensors.
|
||||
C. transformer disk size is materially smaller than a BF16 baseline.
|
||||
|
||||
Example:
|
||||
python examples/quantization/check_modelopt_fp8_export.py \\
|
||||
--output ./hv15-480p-modelopt-fp8
|
||||
|
||||
# Optional: compare disk size against a local or HF BF16 baseline.
|
||||
python examples/quantization/check_modelopt_fp8_export.py \\
|
||||
--output ./hv15-480p-modelopt-fp8 \\
|
||||
--baseline hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import sys
|
||||
from collections import Counter
|
||||
from pathlib import Path
|
||||
|
||||
SUPPORTED_ALGOS = ("FP8",)
|
||||
|
||||
|
||||
def _check_config(transformer_dir: Path) -> tuple[int, str | None]:
|
||||
"""Returns (status, quant_algo). status: 0 pass, 1 fail, 2 warn."""
|
||||
cfg_path = transformer_dir / "config.json"
|
||||
if not cfg_path.exists():
|
||||
print(f"[FAIL] {cfg_path} missing.")
|
||||
return 1, None
|
||||
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
qc = cfg.get("quantization_config")
|
||||
if not isinstance(qc, dict):
|
||||
print(f"[FAIL] No `quantization_config` block in {cfg_path}.")
|
||||
return 1, None
|
||||
|
||||
print(f"[A] quantization_config from {cfg_path}:")
|
||||
print(json.dumps(qc, indent=2))
|
||||
|
||||
quant_algo = qc.get("quant_algo")
|
||||
issues = []
|
||||
if qc.get("quant_method") != "modelopt":
|
||||
issues.append(f"quant_method={qc.get('quant_method')!r} (expected 'modelopt')")
|
||||
if quant_algo not in SUPPORTED_ALGOS:
|
||||
issues.append(
|
||||
f"quant_algo={quant_algo!r} (expected one of {SUPPORTED_ALGOS} — vllm-omni adapter may not auto-detect)"
|
||||
)
|
||||
|
||||
if issues:
|
||||
print("[A] WARN — config looks incomplete:")
|
||||
for issue in issues:
|
||||
print(f" - {issue}")
|
||||
return 2, quant_algo
|
||||
print(f"[A] PASS — config looks correct (quant_algo={quant_algo}).")
|
||||
return 0, quant_algo
|
||||
|
||||
|
||||
def _read_safetensors_header(path: Path) -> dict:
|
||||
"""Read the JSON header of a safetensors file. Bypass-safe — doesn't materialize tensors.
|
||||
|
||||
Returns {tensor_name: {'dtype': 'F8_E4M3', 'shape': [...], 'data_offsets': [...]}}.
|
||||
Header dtype strings: F8_E4M3, F8_E5M2, BF16, F16, F32, F64, I8, I16, I32, I64, BOOL, U8, ...
|
||||
"""
|
||||
import struct
|
||||
|
||||
with open(path, "rb") as f:
|
||||
header_len = struct.unpack("<Q", f.read(8))[0]
|
||||
header = json.loads(f.read(header_len))
|
||||
header.pop("__metadata__", None)
|
||||
return header
|
||||
|
||||
|
||||
def _classify_weight_scale_granularity(weight_scale_shapes: list[list[int]]) -> str:
|
||||
"""Infer per-tensor vs per-channel vs per-block from sample weight_scale shapes.
|
||||
|
||||
ModelOpt block-wise produces shapes like `[16, 1, 16, 1]` (broadcasting dims of 1
|
||||
interleaved with block-count dims). We count "meaningful" dims — ones with size > 1 —
|
||||
and classify: 0 meaningful dims = per-tensor (scalar), 1 = per-channel, 2+ = per-block.
|
||||
"""
|
||||
if not weight_scale_shapes:
|
||||
return "no weight_scale tensors found"
|
||||
|
||||
def meaningful_dims(shape: list[int]) -> int:
|
||||
return sum(1 for d in shape if d > 1)
|
||||
|
||||
per_tensor = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 0)
|
||||
per_channel = sum(1 for s in weight_scale_shapes if meaningful_dims(s) == 1)
|
||||
per_block = sum(1 for s in weight_scale_shapes if meaningful_dims(s) >= 2)
|
||||
total = len(weight_scale_shapes)
|
||||
if per_tensor == total:
|
||||
return "per-tensor (all scalar scales)"
|
||||
if per_channel == total:
|
||||
return "per-channel (1 meaningful dim)"
|
||||
if per_block == total:
|
||||
return "per-block (2+ meaningful dims — e.g. [M//bm, 1, N//bn, 1] for tiles)"
|
||||
return f"mixed: per-tensor={per_tensor}, per-channel={per_channel}, per-block={per_block} of {total}"
|
||||
|
||||
|
||||
def _check_safetensors(transformer_dir: Path) -> int:
|
||||
"""Returns 0 on pass, 1 on fail. Reads on-disk dtype from the safetensors header."""
|
||||
files = sorted(transformer_dir.glob("*.safetensors"))
|
||||
if not files:
|
||||
print(f"[FAIL] No *.safetensors in {transformer_dir}.")
|
||||
return 1
|
||||
|
||||
header_dtype_counts: Counter[str] = Counter()
|
||||
sample_quant_weight_keys: list[str] = []
|
||||
sample_scale_keys: list[str] = []
|
||||
weight_scale_shapes: list[list[int]] = []
|
||||
sample_weight_scale_entries: list[tuple[str, list[int]]] = []
|
||||
|
||||
for f in files:
|
||||
try:
|
||||
header = _read_safetensors_header(f)
|
||||
except Exception as exc:
|
||||
print(f"[B] WARN — could not parse header of {f}: {exc}")
|
||||
continue
|
||||
for k, info in header.items():
|
||||
dtype = info.get("dtype", "?")
|
||||
header_dtype_counts[dtype] += 1
|
||||
# FP8 stores weights as F8_E4M3 directly.
|
||||
if dtype.startswith("F8") and k.endswith(".weight") and len(sample_quant_weight_keys) < 5:
|
||||
sample_quant_weight_keys.append(k)
|
||||
if k.endswith(("_scale", ".weight_scale", ".input_scale", "_scale_inv")) and len(sample_scale_keys) < 5:
|
||||
sample_scale_keys.append(k)
|
||||
if k.endswith(".weight_scale"):
|
||||
weight_scale_shapes.append(info.get("shape", []))
|
||||
if len(sample_weight_scale_entries) < 5:
|
||||
sample_weight_scale_entries.append((k, info.get("shape", [])))
|
||||
|
||||
print(f"\n[B] On-disk dtype counts across {len(files)} safetensors file(s) (from header, not get_tensor):")
|
||||
for dtype, count in sorted(header_dtype_counts.items(), key=lambda kv: -kv[1]):
|
||||
marker = " <-- FP8" if dtype.startswith("F8") else ""
|
||||
print(f" {dtype:10s} {count:>6d}{marker}")
|
||||
|
||||
quant_count = sum(c for d, c in header_dtype_counts.items() if d.startswith("F8"))
|
||||
if quant_count == 0:
|
||||
print("[B] FAIL — no FP8 tensors on disk. Calibration likely did not actually quantize the weights.")
|
||||
return 1
|
||||
|
||||
print(f"[B] PASS — {quant_count} FP8 tensors stored on disk.")
|
||||
if sample_quant_weight_keys:
|
||||
print(f" sample quantized weight tensors: {sample_quant_weight_keys[:3]}")
|
||||
if sample_scale_keys:
|
||||
print(f" sample scale tensors: {sample_scale_keys[:3]}")
|
||||
print(" (Note: torch's get_tensor() may return these as bf16 views on some versions —")
|
||||
print(" irrelevant; vLLM's loader uses native FP8 ops.)")
|
||||
|
||||
# Weight-scale granularity — per-tensor (scalar) vs per-channel (1-D) vs per-block (N-D).
|
||||
print(f"\n weight_scale granularity: {_classify_weight_scale_granularity(weight_scale_shapes)}")
|
||||
for key, shape in sample_weight_scale_entries[:3]:
|
||||
print(f" {key}: shape {shape}")
|
||||
return 0
|
||||
|
||||
|
||||
def _disk_size_gib(p: Path) -> float:
|
||||
return sum(f.stat().st_size for f in p.rglob("*") if f.is_file()) / (1024**3)
|
||||
|
||||
|
||||
def _check_size_vs_baseline(transformer_dir: Path, baseline: str | None) -> int:
|
||||
"""Returns 0 always (informational only)."""
|
||||
quant_size = _disk_size_gib(transformer_dir)
|
||||
print(f"\n[C] FP8 transformer disk size: {quant_size:.2f} GiB")
|
||||
|
||||
if baseline is None:
|
||||
print("[C] SKIP — pass --baseline <path or HF id> to compare against BF16.")
|
||||
return 0
|
||||
|
||||
baseline_path = Path(baseline)
|
||||
if not baseline_path.exists():
|
||||
# Try HF download.
|
||||
try:
|
||||
from huggingface_hub import snapshot_download
|
||||
except ImportError:
|
||||
print("[C] SKIP — huggingface_hub not installed and baseline not a local path.")
|
||||
return 0
|
||||
print(f" Downloading baseline transformer from HF: {baseline}")
|
||||
baseline_path = Path(snapshot_download(baseline, allow_patterns=["transformer/*"]))
|
||||
|
||||
bf16_dir = baseline_path / "transformer" if (baseline_path / "transformer").exists() else baseline_path
|
||||
bf16_size = _disk_size_gib(bf16_dir)
|
||||
if bf16_size == 0:
|
||||
print(f"[C] WARN — baseline transformer dir empty: {bf16_dir}")
|
||||
return 0
|
||||
|
||||
# Expected reduction for FP8: ~50%.
|
||||
min_reduction = 30
|
||||
reduction = (1 - quant_size / bf16_size) * 100
|
||||
print(f"[C] BF16 baseline transformer disk size: {bf16_size:.2f} GiB ({bf16_dir})")
|
||||
print(f"[C] Disk reduction: {reduction:.1f}% (FP8 transformer is {quant_size / bf16_size:.0%} of BF16)")
|
||||
if reduction < min_reduction:
|
||||
print(
|
||||
f"[C] WARN — FP8 should typically reduce disk by ~40-50%; <{min_reduction}% suggests partial quantization."
|
||||
)
|
||||
return 0
|
||||
|
||||
|
||||
def main() -> None:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--output", required=True, help="Path to the exported ModelOpt FP8 checkpoint root.")
|
||||
p.add_argument(
|
||||
"--baseline",
|
||||
default=None,
|
||||
help="Optional BF16 baseline (local diffusers dir or HF id) for disk-size comparison.",
|
||||
)
|
||||
args = p.parse_args()
|
||||
|
||||
out_root = Path(args.output).expanduser().resolve()
|
||||
transformer_dir = out_root / "transformer"
|
||||
if not transformer_dir.exists():
|
||||
print(f"[FAIL] {transformer_dir} does not exist.")
|
||||
sys.exit(1)
|
||||
|
||||
print(f"Checking: {out_root}\n")
|
||||
|
||||
fail = 0
|
||||
config_status, _quant_algo = _check_config(transformer_dir)
|
||||
fail |= config_status
|
||||
fail |= _check_safetensors(transformer_dir)
|
||||
_check_size_vs_baseline(transformer_dir, args.baseline)
|
||||
|
||||
print()
|
||||
if fail == 0:
|
||||
print("=" * 60)
|
||||
print("ALL CHECKS PASSED — checkpoint looks ready for vllm-omni serving.")
|
||||
elif fail == 1:
|
||||
print("=" * 60)
|
||||
print("FAILURES detected — calibration may need to be re-run.")
|
||||
sys.exit(1)
|
||||
else:
|
||||
print("=" * 60)
|
||||
print("WARNINGS only — checkpoint may serve but with caveats. See [A] above.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,448 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Quantize HunyuanVideo-1.5 (480p T2V) to a ModelOpt FP8 Hugging Face checkpoint.
|
||||
|
||||
Calibrates the DiT transformer using a small video prompt set and exports a
|
||||
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
|
||||
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
|
||||
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
|
||||
|
||||
Layers kept full precision match the #2728 / #2795 pattern: modulation,
|
||||
AdaLayerNorm, entry/exit projections, embeddings, the token refiner path,
|
||||
and final proj_out. MHA quantizers are off by default; HV-1.5 self-attention
|
||||
empirically degrades under FP8 (see #2920 ablation).
|
||||
|
||||
Example:
|
||||
python examples/quantization/quantize_hunyuanvideo_15_modelopt_fp8.py \\
|
||||
--model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v \\
|
||||
--output ./hv15-480p-modelopt-fp8 \\
|
||||
--overwrite
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
DEFAULT_PROMPTS = [
|
||||
"A dog running across a field of golden wheat.",
|
||||
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
|
||||
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
|
||||
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
|
||||
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
|
||||
"A close-up of a blooming rose covered in morning dew, soft natural light.",
|
||||
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
|
||||
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
|
||||
]
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--model", required=True, help="Input HV-1.5 diffusers directory or HF id.")
|
||||
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
|
||||
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
|
||||
p.add_argument("--height", type=int, default=480)
|
||||
p.add_argument("--width", type=int, default=832)
|
||||
p.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=33,
|
||||
help="Frames per calibration sample. 33 matches the typical short benchmark.",
|
||||
)
|
||||
p.add_argument("--guidance-scale", type=float, default=6.0)
|
||||
p.add_argument(
|
||||
"--calib-steps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
|
||||
)
|
||||
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument(
|
||||
"--prompt",
|
||||
action="append",
|
||||
default=[],
|
||||
help="Custom calibration prompt. Repeat to provide multiple.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--quantize-mha",
|
||||
action="store_true",
|
||||
help="Enable FP8 attention K/V/softmax quantizers. Off by default — empirically degrades HV-1.5 video output.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--weight-block-size",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Per-block weight quantization as 'M,N' (e.g. '128,128' for 128x128 tiles). "
|
||||
"Default: per-tensor (one scale per linear). Block-wise typically gives tighter quality at "
|
||||
"negligible memory cost. Static FP8 is exempt from upstream vLLM's online block-wise gate.",
|
||||
)
|
||||
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
|
||||
return p
|
||||
|
||||
|
||||
def _parse_block_size(spec: str | None) -> list[int] | None:
|
||||
if spec is None:
|
||||
return None
|
||||
parts = [int(x) for x in spec.split(",") if x.strip()]
|
||||
if len(parts) != 2:
|
||||
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
|
||||
return parts
|
||||
|
||||
|
||||
def _require_modelopt() -> Any:
|
||||
try:
|
||||
import modelopt.torch.quantization as mtq
|
||||
except ModuleNotFoundError as exc:
|
||||
raise SystemExit(
|
||||
"NVIDIA ModelOpt is not installed. Install with:\n"
|
||||
" pip install 'nvidia-modelopt[all]'\n"
|
||||
f"Original error: {exc}"
|
||||
) from exc
|
||||
return mtq
|
||||
|
||||
|
||||
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
|
||||
model_path = args.model
|
||||
output_dir = Path(args.output).expanduser().resolve()
|
||||
if output_dir.exists():
|
||||
if not args.overwrite:
|
||||
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
|
||||
shutil.rmtree(output_dir)
|
||||
return model_path, output_dir
|
||||
|
||||
|
||||
def _select_dtype(name: str) -> torch.dtype:
|
||||
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
|
||||
|
||||
|
||||
def _build_prompts(args: argparse.Namespace) -> list[str]:
|
||||
prompts = args.prompt or DEFAULT_PROMPTS
|
||||
if args.calib_size <= 0:
|
||||
raise SystemExit("--calib-size must be positive.")
|
||||
if len(prompts) < args.calib_size:
|
||||
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
|
||||
prompts = (prompts * repeats)[: args.calib_size]
|
||||
return prompts[: args.calib_size]
|
||||
|
||||
|
||||
# Layers to KEEP at full precision (mirror of the #2920 wiring + #2728/#2795 skip pattern).
|
||||
# - x_embedder, image_embedder, context_embedder*, time_embed*, cond_type_embed: entry/embedding
|
||||
# - norm_out, norm1*.linear, norm1_context*.linear, norm2*, norm2_context*: AdaLayerNorm modulation
|
||||
# - proj_out: final output projection
|
||||
# - token_refiner*: text-encoder refinement uses diffusers raw nn.Linear
|
||||
def _filter_func_hv15(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r"(proj_out.*|"
|
||||
r".*(x_embedder|image_embedder|context_embedder|context_embedder_2|"
|
||||
r"time_embed|cond_type_embed|"
|
||||
r"norm_out|norm1\.linear|norm1_context\.linear|norm2|norm2_context|"
|
||||
r"token_refiner).*)"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _mha_filter_func(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
|
||||
if not hasattr(mtq, "disable_quantizer"):
|
||||
return
|
||||
mtq.disable_quantizer(backbone, _filter_func_hv15)
|
||||
if not quantize_mha:
|
||||
mtq.disable_quantizer(backbone, _mha_filter_func)
|
||||
|
||||
|
||||
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
|
||||
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
|
||||
if hasattr(pipe, "set_progress_bar_config"):
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
|
||||
generator = torch.Generator(device="cuda")
|
||||
|
||||
# Try to set guidance on the pipeline's guider object up front (modern
|
||||
# diffusers HV-1.5 uses a Guider abstraction, not a per-call kwarg). Falls
|
||||
# back silently — calibration uses whatever default the pipeline ships with.
|
||||
guider = getattr(pipe, "guider", None)
|
||||
if guider is not None and hasattr(guider, "guidance_scale"):
|
||||
try:
|
||||
guider.guidance_scale = args.guidance_scale
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
base_kwargs = dict(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.calib_steps,
|
||||
output_type="latent",
|
||||
)
|
||||
|
||||
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
|
||||
with torch.inference_mode():
|
||||
for idx, prompt in enumerate(prompts):
|
||||
generator.manual_seed(args.seed + idx)
|
||||
# Try with guidance_scale first; fall back without on TypeError
|
||||
# for pipelines (like HV-1.5) that take CFG via guider config.
|
||||
try:
|
||||
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
|
||||
except TypeError as exc:
|
||||
if "guidance_scale" not in str(exc):
|
||||
raise
|
||||
pipe(prompt=prompt, generator=generator, **base_kwargs)
|
||||
|
||||
return forward_loop
|
||||
|
||||
|
||||
def _summarize_export(output_dir: Path) -> None:
|
||||
cfg_path = output_dir / "transformer" / "config.json"
|
||||
if not cfg_path.exists():
|
||||
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
|
||||
return
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
qc = cfg.get("quantization_config")
|
||||
if not isinstance(qc, dict):
|
||||
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
|
||||
return
|
||||
print("Export summary:")
|
||||
print(f" quant_method: {qc.get('quant_method')}")
|
||||
print(f" quant_algo: {qc.get('quant_algo')}")
|
||||
producer = qc.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
print(f" producer: {producer.get('name')} {producer.get('version')}")
|
||||
print(f" config path: {cfg_path}")
|
||||
|
||||
|
||||
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
|
||||
"""Convert in-memory weights of quantized modules to actual FP8 storage.
|
||||
|
||||
`export_hf_checkpoint` skips this step for unknown model types (HV-1.5 isn't
|
||||
in ModelOpt's recognized-model registry), so we must call the per-weight
|
||||
export helper ourselves. Same workaround as the HunyuanImage-3 calibration
|
||||
helper.
|
||||
"""
|
||||
from modelopt.torch.export.quant_utils import (
|
||||
QUANTIZATION_NONE,
|
||||
get_quantization_format,
|
||||
quantizer_attr_names,
|
||||
weight_attr_names,
|
||||
)
|
||||
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
|
||||
|
||||
exported = 0
|
||||
for name, module in backbone.named_modules():
|
||||
try:
|
||||
quantization_format = get_quantization_format(module)
|
||||
except Exception as exc:
|
||||
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
|
||||
continue
|
||||
if quantization_format == QUANTIZATION_NONE:
|
||||
continue
|
||||
for weight_name in weight_attr_names(module):
|
||||
quantizer_attrs = quantizer_attr_names(weight_name)
|
||||
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
|
||||
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
|
||||
continue
|
||||
_export_quantized_weight(module, dtype, weight_name)
|
||||
exported += 1
|
||||
return exported
|
||||
|
||||
|
||||
def _hv15_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
|
||||
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913).
|
||||
|
||||
Same shape as the HunyuanImage-3 author's _hunyuan_quant_config(). When
|
||||
`weight_block_size` is given, advertise block-wise weight quantization in
|
||||
the saved metadata (so consumers know to expect multi-element scale tensors).
|
||||
"""
|
||||
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
|
||||
if weight_block_size is not None:
|
||||
weights_cfg["strategy"] = "block"
|
||||
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
|
||||
return {
|
||||
"config_groups": {
|
||||
"group_0": {
|
||||
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
|
||||
"weights": weights_cfg,
|
||||
"targets": ["Linear"],
|
||||
}
|
||||
},
|
||||
"ignore": [
|
||||
"context_embedder*",
|
||||
"context_embedder_2*",
|
||||
"cond_type_embed*",
|
||||
"image_embedder*",
|
||||
"norm1.linear*",
|
||||
"norm1_context.linear*",
|
||||
"norm2*",
|
||||
"norm2_context*",
|
||||
"norm_out*",
|
||||
"proj_out*",
|
||||
"time_embed*",
|
||||
"token_refiner*",
|
||||
"x_embedder*",
|
||||
],
|
||||
"producer": {"name": "modelopt"},
|
||||
"quant_algo": "FP8",
|
||||
"quant_method": "modelopt",
|
||||
}
|
||||
|
||||
|
||||
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
|
||||
"""Inject quant_algo: FP8 + config_groups into transformer/config.json so
|
||||
vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8."""
|
||||
cfg_path = output_dir / "transformer" / "config.json"
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
new_qc = _hv15_quant_config_block(weight_block_size=weight_block_size)
|
||||
existing = cfg.get("quantization_config")
|
||||
if isinstance(existing, dict):
|
||||
producer = existing.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
new_qc["producer"] = producer
|
||||
|
||||
cfg["quantization_config"] = new_qc
|
||||
with cfg_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(cfg, f, indent=2)
|
||||
|
||||
|
||||
def _save_pipeline_with_fp8_transformer(
|
||||
pipe: DiffusionPipeline,
|
||||
model_path: str,
|
||||
output_dir: Path,
|
||||
max_shard_size: str = "5GB",
|
||||
) -> None:
|
||||
"""Save the pipeline with the (now FP8) transformer.
|
||||
|
||||
Copies the source directory verbatim except for `transformer/`, then
|
||||
saves the transformer with quantizers hidden so the state dict contains
|
||||
only the FP8 weights + scale tensors.
|
||||
"""
|
||||
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
|
||||
|
||||
src = Path(model_path)
|
||||
if not src.exists():
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
src = Path(snapshot_download(model_path))
|
||||
|
||||
if output_dir.exists():
|
||||
shutil.rmtree(output_dir)
|
||||
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer"))
|
||||
|
||||
transformer_out = output_dir / "transformer"
|
||||
# `hide_quantizers_from_state_dict` walks named_modules(); pass the actual
|
||||
# nn.Module (transformer), not the diffusers Pipeline wrapper.
|
||||
with hide_quantizers_from_state_dict(pipe.transformer):
|
||||
pipe.transformer.save_pretrained(
|
||||
str(transformer_out),
|
||||
safe_serialization=True,
|
||||
max_shard_size=max_shard_size,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _build_parser().parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
|
||||
|
||||
mtq = _require_modelopt()
|
||||
model_path, output_dir = _ensure_paths(args)
|
||||
dtype = _select_dtype(args.dtype)
|
||||
prompts = _build_prompts(args)
|
||||
|
||||
print("Quantization plan:")
|
||||
weight_block_size = _parse_block_size(args.weight_block_size)
|
||||
|
||||
print(f" input: {args.model}")
|
||||
print(f" output: {output_dir}")
|
||||
print(f" dtype: {dtype}")
|
||||
print(f" height/width: {args.height}x{args.width}")
|
||||
print(f" num_frames: {args.num_frames}")
|
||||
print(f" calib_size: {len(prompts)}")
|
||||
print(f" calib_steps: {args.calib_steps}")
|
||||
print(f" quantize_mha: {args.quantize_mha}")
|
||||
print(
|
||||
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
|
||||
)
|
||||
|
||||
pipe = _load_pipeline(model_path, dtype)
|
||||
backbone = pipe.transformer
|
||||
|
||||
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
|
||||
if weight_block_size is not None:
|
||||
# Switch from per-tensor (default) to block-wise weight quantization.
|
||||
# ModelOpt's wildcard "*weight_quantizer" matches every linear's weight quantizer.
|
||||
quant_config["quant_cfg"]["*weight_quantizer"] = {
|
||||
"num_bits": (4, 3), # E4M3 (FP8 weights, same as default)
|
||||
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
|
||||
}
|
||||
print(
|
||||
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
|
||||
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
|
||||
)
|
||||
|
||||
forward_loop = _build_forward_loop(pipe, args, prompts)
|
||||
quantized = mtq.quantize(backbone, quant_config, forward_loop)
|
||||
if quantized is not None:
|
||||
pipe.transformer = quantized
|
||||
backbone = quantized
|
||||
|
||||
_disable_known_problematic_quantizers(mtq, backbone, quantize_mha=args.quantize_mha)
|
||||
|
||||
print("\nForcing FP8 weight serialization (HV-1.5 isn't in ModelOpt's recognized-model registry,")
|
||||
print("so we have to call the per-weight export helper ourselves)...")
|
||||
exported = _force_export_quantized_weights(backbone, dtype)
|
||||
print(f" -> {exported} weights converted to FP8 in memory")
|
||||
if exported == 0:
|
||||
raise SystemExit(
|
||||
"No quantized weights were exported. Calibration may have skipped every layer "
|
||||
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
|
||||
"weight quantizers."
|
||||
)
|
||||
|
||||
print("\nSaving pipeline with FP8 transformer...")
|
||||
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
|
||||
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
|
||||
print(f"Saved to: {output_dir}")
|
||||
_summarize_export(output_dir)
|
||||
|
||||
print("\nNext: validate the checkpoint with vllm-omni:")
|
||||
print(
|
||||
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
|
||||
f" --model {output_dir} \\\n"
|
||||
" --quantization fp8 \\\n"
|
||||
" --prompt 'A dog running across a field of golden wheat.' \\\n"
|
||||
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
|
||||
" --num-inference-steps 30 --guidance-scale 6.0 --seed 42 \\\n"
|
||||
" --output outputs/hv15_modelopt_fp8.mp4 \\\n"
|
||||
" --enforce-eager"
|
||||
)
|
||||
print(
|
||||
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
|
||||
"checkpoint's config.json has modelopt metadata.)"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,471 @@
|
||||
#!/usr/bin/env python3
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Quantize Wan2.2 (TI2V-5B, 704x1280 T2V) to a ModelOpt FP8 Hugging Face checkpoint.
|
||||
|
||||
Calibrates the DiT transformer using a small video prompt set and exports a
|
||||
diffusers-style directory whose transformer carries ModelOpt FP8 metadata.
|
||||
The exported checkpoint is consumable by vllm-omni's ModelOpt FP8 adapter
|
||||
(see vllm_omni/diffusion/model_loader/checkpoint_adapters/modelopt.py).
|
||||
|
||||
Layers kept full precision match the #2728 / #2795 pattern: condition embedder
|
||||
(time/text/image), patch embedding, modulation (scale_shift_table), final
|
||||
norm + proj_out, and sequence-parallel helpers. All attention + FFN linears
|
||||
are quantized — static calibration handles the numerics that online FP8
|
||||
couldn't (see #2920 ablation).
|
||||
|
||||
Default target is `Wan-AI/Wan2.2-TI2V-5B-Diffusers`, the dense 5B variant that
|
||||
fits 80GB BF16. The A14B MoE variants need 2+ GPUs and are out of scope here.
|
||||
|
||||
Example:
|
||||
python examples/quantization/quantize_wan2_2_modelopt_fp8.py \\
|
||||
--model Wan-AI/Wan2.2-TI2V-5B-Diffusers \\
|
||||
--output ./wan22-ti2v-modelopt-fp8 \\
|
||||
--overwrite
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import re
|
||||
import shutil
|
||||
import sys
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from diffusers import DiffusionPipeline
|
||||
|
||||
DEFAULT_PROMPTS = [
|
||||
"A dog running across a field of golden wheat.",
|
||||
"An astronaut riding a horse across the surface of Mars, red dust swirling, cinematic wide shot.",
|
||||
"A hummingbird hovering in front of a vibrant red flower, slow motion, macro shot.",
|
||||
"A crackling campfire at night under a starry sky, sparks rising into the dark.",
|
||||
"An underwater shot of a coral reef with tropical fish swimming by, sun rays piercing the water.",
|
||||
"A close-up of a blooming rose covered in morning dew, soft natural light.",
|
||||
"A peaceful mountain village at dawn, mist rolling over the rooftops, cinematic establishing shot.",
|
||||
"A skateboarder doing a kickflip in an urban plaza, slow motion, golden hour lighting.",
|
||||
]
|
||||
|
||||
|
||||
def _build_parser() -> argparse.ArgumentParser:
|
||||
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
||||
p.add_argument("--model", required=True, help="Input Wan2.2 diffusers directory or HF id.")
|
||||
p.add_argument("--output", required=True, help="Output directory for the ModelOpt FP8 checkpoint.")
|
||||
p.add_argument("--dtype", choices=("bfloat16", "float16"), default="bfloat16")
|
||||
p.add_argument("--height", type=int, default=704, help="Calibration video height (Wan2.2 TI2V-5B native: 704).")
|
||||
p.add_argument("--width", type=int, default=1280, help="Calibration video width (Wan2.2 TI2V-5B native: 1280).")
|
||||
p.add_argument(
|
||||
"--num-frames",
|
||||
type=int,
|
||||
default=49,
|
||||
help="Frames per calibration sample. 49 matches the typical short benchmark; "
|
||||
"use 17 to reduce memory pressure during calibration.",
|
||||
)
|
||||
p.add_argument("--guidance-scale", type=float, default=5.0)
|
||||
p.add_argument(
|
||||
"--calib-steps",
|
||||
type=int,
|
||||
default=10,
|
||||
help="Denoising steps per calibration prompt (10 is enough for amax statistics).",
|
||||
)
|
||||
p.add_argument("--calib-size", type=int, default=8, help="How many prompts to use for calibration.")
|
||||
p.add_argument("--seed", type=int, default=42)
|
||||
p.add_argument(
|
||||
"--prompt",
|
||||
action="append",
|
||||
default=[],
|
||||
help="Custom calibration prompt. Repeat to provide multiple.",
|
||||
)
|
||||
p.add_argument(
|
||||
"--quantize-mha",
|
||||
action="store_true",
|
||||
help="Enable FP8 attention K/V/softmax quantizers. Off by default — Wan2.2's long attention "
|
||||
"sequences amplified FP8 drift in the online ablation (see #2920).",
|
||||
)
|
||||
p.add_argument(
|
||||
"--weight-block-size",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Per-block weight quantization as 'M,N' (e.g. '128,128'). Default per-tensor. "
|
||||
"Note: vllm-omni's ModelOpt adapter may not yet dispatch block-wise scales — check #2924 "
|
||||
"for the HV-1.5 investigation status before relying on this.",
|
||||
)
|
||||
p.add_argument("--overwrite", action="store_true", help="Replace an existing output directory.")
|
||||
return p
|
||||
|
||||
|
||||
def _parse_block_size(spec: str | None) -> list[int] | None:
|
||||
if spec is None:
|
||||
return None
|
||||
parts = [int(x) for x in spec.split(",") if x.strip()]
|
||||
if len(parts) != 2:
|
||||
raise SystemExit(f"--weight-block-size must be 'M,N' (2 ints), got {spec!r}")
|
||||
return parts
|
||||
|
||||
|
||||
def _require_modelopt() -> Any:
|
||||
try:
|
||||
import modelopt.torch.quantization as mtq
|
||||
except ModuleNotFoundError as exc:
|
||||
raise SystemExit(
|
||||
"NVIDIA ModelOpt is not installed. Install with:\n"
|
||||
" pip install 'nvidia-modelopt[all]'\n"
|
||||
f"Original error: {exc}"
|
||||
) from exc
|
||||
return mtq
|
||||
|
||||
|
||||
def _ensure_paths(args: argparse.Namespace) -> tuple[str, Path]:
|
||||
model_path = args.model
|
||||
output_dir = Path(args.output).expanduser().resolve()
|
||||
if output_dir.exists():
|
||||
if not args.overwrite:
|
||||
raise SystemExit(f"Output directory already exists: {output_dir}\nPass --overwrite to replace it.")
|
||||
shutil.rmtree(output_dir)
|
||||
return model_path, output_dir
|
||||
|
||||
|
||||
def _select_dtype(name: str) -> torch.dtype:
|
||||
return {"bfloat16": torch.bfloat16, "float16": torch.float16}[name]
|
||||
|
||||
|
||||
def _build_prompts(args: argparse.Namespace) -> list[str]:
|
||||
prompts = args.prompt or DEFAULT_PROMPTS
|
||||
if args.calib_size <= 0:
|
||||
raise SystemExit("--calib-size must be positive.")
|
||||
if len(prompts) < args.calib_size:
|
||||
repeats = (args.calib_size + len(prompts) - 1) // len(prompts)
|
||||
prompts = (prompts * repeats)[: args.calib_size]
|
||||
return prompts[: args.calib_size]
|
||||
|
||||
|
||||
# Layers to KEEP at full precision. Wan2.2's module naming:
|
||||
# - condition_embedder: time_embedder, time_proj, text_embedder, image_embedder (I2V)
|
||||
# - patch_embedding: Conv3dLayer (already not Linear, belt-and-suspenders skip)
|
||||
# - scale_shift_table: nn.Parameter modulation (not Linear, but pattern guard)
|
||||
# - norm_out: AdaLayerNorm final
|
||||
# - proj_out: final nn.Linear
|
||||
# - timestep_proj_prepare / output_scale_shift_prepare: SP helpers
|
||||
def _filter_func_wan22(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r"(proj_out.*|"
|
||||
r".*(condition_embedder|patch_embedding|"
|
||||
r"norm_out|scale_shift_table|"
|
||||
r"timestep_proj_prepare|output_scale_shift_prepare).*)"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _mha_filter_func(name: str) -> bool:
|
||||
pattern = re.compile(
|
||||
r".*(q_bmm_quantizer|k_bmm_quantizer|v_bmm_quantizer|softmax_quantizer|bmm2_output_quantizer).*"
|
||||
)
|
||||
return pattern.match(name) is not None
|
||||
|
||||
|
||||
def _disable_known_problematic_quantizers(mtq: Any, backbone: torch.nn.Module, *, quantize_mha: bool) -> None:
|
||||
if not hasattr(mtq, "disable_quantizer"):
|
||||
return
|
||||
mtq.disable_quantizer(backbone, _filter_func_wan22)
|
||||
if not quantize_mha:
|
||||
mtq.disable_quantizer(backbone, _mha_filter_func)
|
||||
|
||||
|
||||
def _load_pipeline(model_path: str, dtype: torch.dtype) -> DiffusionPipeline:
|
||||
pipe = DiffusionPipeline.from_pretrained(model_path, torch_dtype=dtype)
|
||||
if hasattr(pipe, "set_progress_bar_config"):
|
||||
pipe.set_progress_bar_config(disable=True)
|
||||
pipe.to("cuda")
|
||||
return pipe
|
||||
|
||||
|
||||
def _build_forward_loop(pipe: DiffusionPipeline, args: argparse.Namespace, prompts: list[str]):
|
||||
generator = torch.Generator(device="cuda")
|
||||
|
||||
# Try setting guidance on the pipeline's guider if present (newer diffusers APIs).
|
||||
guider = getattr(pipe, "guider", None)
|
||||
if guider is not None and hasattr(guider, "guidance_scale"):
|
||||
try:
|
||||
guider.guidance_scale = args.guidance_scale
|
||||
except Exception:
|
||||
pass
|
||||
|
||||
base_kwargs = dict(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.calib_steps,
|
||||
output_type="latent",
|
||||
)
|
||||
|
||||
def forward_loop(*_unused_args, **_unused_kwargs) -> None:
|
||||
with torch.inference_mode():
|
||||
for idx, prompt in enumerate(prompts):
|
||||
generator.manual_seed(args.seed + idx)
|
||||
# Try with guidance_scale first; fall back without on TypeError
|
||||
# for pipelines that take CFG via guider config only.
|
||||
try:
|
||||
pipe(prompt=prompt, generator=generator, guidance_scale=args.guidance_scale, **base_kwargs)
|
||||
except TypeError as exc:
|
||||
if "guidance_scale" not in str(exc):
|
||||
raise
|
||||
pipe(prompt=prompt, generator=generator, **base_kwargs)
|
||||
|
||||
return forward_loop
|
||||
|
||||
|
||||
def _summarize_export(output_dir: Path) -> None:
|
||||
cfg_path = output_dir / "transformer" / "config.json"
|
||||
if not cfg_path.exists():
|
||||
print(f"[warn] {cfg_path} missing.", file=sys.stderr)
|
||||
return
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
qc = cfg.get("quantization_config")
|
||||
if not isinstance(qc, dict):
|
||||
print("[warn] No quantization_config in transformer/config.json.", file=sys.stderr)
|
||||
return
|
||||
print("Export summary:")
|
||||
print(f" quant_method: {qc.get('quant_method')}")
|
||||
print(f" quant_algo: {qc.get('quant_algo')}")
|
||||
producer = qc.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
print(f" producer: {producer.get('name')} {producer.get('version')}")
|
||||
print(f" config path: {cfg_path}")
|
||||
|
||||
|
||||
def _force_export_quantized_weights(backbone: torch.nn.Module, dtype: torch.dtype) -> int:
|
||||
"""Convert in-memory weights of quantized modules to actual FP8 storage.
|
||||
|
||||
`export_hf_checkpoint` skips this step for unknown model types (Wan2.2 isn't
|
||||
in ModelOpt's recognized-model registry), so we must call the per-weight
|
||||
export helper ourselves. Same workaround as the HunyuanVideo-1.5 / HunyuanImage-3
|
||||
calibration helpers.
|
||||
"""
|
||||
from modelopt.torch.export.quant_utils import (
|
||||
QUANTIZATION_NONE,
|
||||
get_quantization_format,
|
||||
quantizer_attr_names,
|
||||
weight_attr_names,
|
||||
)
|
||||
from modelopt.torch.export.unified_export_hf import _export_quantized_weight
|
||||
|
||||
exported = 0
|
||||
for name, module in backbone.named_modules():
|
||||
try:
|
||||
quantization_format = get_quantization_format(module)
|
||||
except Exception as exc:
|
||||
print(f"[warn] Could not inspect quantization format for {name}: {exc}", file=sys.stderr)
|
||||
continue
|
||||
if quantization_format == QUANTIZATION_NONE:
|
||||
continue
|
||||
for weight_name in weight_attr_names(module):
|
||||
quantizer_attrs = quantizer_attr_names(weight_name)
|
||||
weight_quantizer = getattr(module, quantizer_attrs.weight_quantizer, None)
|
||||
if weight_quantizer is None or not getattr(weight_quantizer, "is_enabled", False):
|
||||
continue
|
||||
_export_quantized_weight(module, dtype, weight_name)
|
||||
exported += 1
|
||||
return exported
|
||||
|
||||
|
||||
def _wan22_quant_config_block(weight_block_size: list[int] | None = None) -> dict:
|
||||
"""Mirror ModelOpt FP8 metadata expected by vllm-omni's adapter (#2913)."""
|
||||
weights_cfg: dict = {"dynamic": False, "num_bits": 8, "type": "float"}
|
||||
if weight_block_size is not None:
|
||||
weights_cfg["strategy"] = "block"
|
||||
weights_cfg["block_structure"] = f"{weight_block_size[0]}x{weight_block_size[1]}"
|
||||
return {
|
||||
"config_groups": {
|
||||
"group_0": {
|
||||
"input_activations": {"dynamic": False, "num_bits": 8, "type": "float"},
|
||||
"weights": weights_cfg,
|
||||
"targets": ["Linear"],
|
||||
}
|
||||
},
|
||||
"ignore": [
|
||||
"condition_embedder*",
|
||||
"norm_out*",
|
||||
"output_scale_shift_prepare*",
|
||||
"patch_embedding*",
|
||||
"proj_out*",
|
||||
"scale_shift_table*",
|
||||
"timestep_proj_prepare*",
|
||||
],
|
||||
"producer": {"name": "modelopt"},
|
||||
"quant_algo": "FP8",
|
||||
"quant_method": "modelopt",
|
||||
}
|
||||
|
||||
|
||||
def _patch_quant_config(output_dir: Path, weight_block_size: list[int] | None = None) -> None:
|
||||
"""Inject quant_algo: FP8 + config_groups into each transformer's config.json
|
||||
so vllm-omni's adapter (#2913) recognises the checkpoint as ModelOpt FP8.
|
||||
|
||||
Patches both ``transformer`` and ``transformer_2`` (A14B) when present."""
|
||||
for sub in ("transformer", "transformer_2"):
|
||||
cfg_path = output_dir / sub / "config.json"
|
||||
if not cfg_path.exists():
|
||||
continue
|
||||
with cfg_path.open(encoding="utf-8") as f:
|
||||
cfg = json.load(f)
|
||||
|
||||
new_qc = _wan22_quant_config_block(weight_block_size=weight_block_size)
|
||||
existing = cfg.get("quantization_config")
|
||||
if isinstance(existing, dict):
|
||||
producer = existing.get("producer")
|
||||
if isinstance(producer, dict):
|
||||
new_qc["producer"] = producer
|
||||
|
||||
cfg["quantization_config"] = new_qc
|
||||
with cfg_path.open("w", encoding="utf-8") as f:
|
||||
json.dump(cfg, f, indent=2)
|
||||
|
||||
|
||||
def _list_transformers(pipe: DiffusionPipeline) -> list[tuple[str, torch.nn.Module]]:
|
||||
"""Return [(attr_name, module)] for every DiT on the pipeline.
|
||||
|
||||
Wan2.2 A14B has both ``transformer`` (low-noise) and ``transformer_2``
|
||||
(high-noise). TI2V-5B has only ``transformer``.
|
||||
"""
|
||||
out = []
|
||||
for attr in ("transformer", "transformer_2"):
|
||||
mod = getattr(pipe, attr, None)
|
||||
if mod is not None:
|
||||
out.append((attr, mod))
|
||||
return out
|
||||
|
||||
|
||||
def _save_pipeline_with_fp8_transformer(
|
||||
pipe: DiffusionPipeline,
|
||||
model_path: str,
|
||||
output_dir: Path,
|
||||
max_shard_size: str = "5GB",
|
||||
) -> None:
|
||||
"""Copy source dir verbatim minus transformer dirs, then save each quantized DiT.
|
||||
|
||||
Handles A14B's dual transformer (``transformer`` + ``transformer_2``) as well
|
||||
as the single-DiT TI2V-5B.
|
||||
"""
|
||||
from modelopt.torch.export.diffusers_utils import hide_quantizers_from_state_dict
|
||||
|
||||
src = Path(model_path)
|
||||
if not src.exists():
|
||||
from huggingface_hub import snapshot_download
|
||||
|
||||
src = Path(snapshot_download(model_path))
|
||||
|
||||
if output_dir.exists():
|
||||
shutil.rmtree(output_dir)
|
||||
shutil.copytree(src, output_dir, ignore=shutil.ignore_patterns("transformer", "transformer_2"))
|
||||
|
||||
for attr, backbone in _list_transformers(pipe):
|
||||
transformer_out = output_dir / attr
|
||||
# Pass the nn.Module (transformer), not the Pipeline wrapper.
|
||||
with hide_quantizers_from_state_dict(backbone):
|
||||
backbone.save_pretrained(
|
||||
str(transformer_out),
|
||||
safe_serialization=True,
|
||||
max_shard_size=max_shard_size,
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _build_parser().parse_args()
|
||||
if not torch.cuda.is_available():
|
||||
raise SystemExit("CUDA is required for ModelOpt FP8 quantization.")
|
||||
|
||||
mtq = _require_modelopt()
|
||||
model_path, output_dir = _ensure_paths(args)
|
||||
dtype = _select_dtype(args.dtype)
|
||||
prompts = _build_prompts(args)
|
||||
weight_block_size = _parse_block_size(args.weight_block_size)
|
||||
|
||||
print("Quantization plan:")
|
||||
print(f" input: {args.model}")
|
||||
print(f" output: {output_dir}")
|
||||
print(f" dtype: {dtype}")
|
||||
print(f" height/width: {args.height}x{args.width}")
|
||||
print(f" num_frames: {args.num_frames}")
|
||||
print(f" calib_size: {len(prompts)}")
|
||||
print(f" calib_steps: {args.calib_steps}")
|
||||
print(f" quantize_mha: {args.quantize_mha}")
|
||||
print(
|
||||
f" weight strategy: {'block-wise ' + str(weight_block_size) if weight_block_size else 'per-tensor (default)'}"
|
||||
)
|
||||
|
||||
pipe = _load_pipeline(model_path, dtype)
|
||||
transformers = _list_transformers(pipe)
|
||||
if not transformers:
|
||||
raise SystemExit("Pipeline has no transformer or transformer_2 attribute.")
|
||||
print(f" found {len(transformers)} transformer(s): {', '.join(a for a, _ in transformers)}")
|
||||
|
||||
quant_config = copy.deepcopy(mtq.FP8_DEFAULT_CFG)
|
||||
if weight_block_size is not None:
|
||||
quant_config["quant_cfg"]["*weight_quantizer"] = {
|
||||
"num_bits": (4, 3),
|
||||
"block_sizes": {-1: weight_block_size[1], -2: weight_block_size[0]},
|
||||
}
|
||||
print(
|
||||
f" -> overriding weight quantizer with block_sizes={weight_block_size} "
|
||||
f"({weight_block_size[0]}x{weight_block_size[1]} tiles)"
|
||||
)
|
||||
|
||||
forward_loop = _build_forward_loop(pipe, args, prompts)
|
||||
|
||||
# Quantize each DiT. The forward_loop runs the full pipeline, so quantizers
|
||||
# on every attached transformer (A14B: transformer + transformer_2) get
|
||||
# calibrated during the same passes.
|
||||
for attr, backbone in transformers:
|
||||
print(f"\nQuantizing {attr}...")
|
||||
quantized = mtq.quantize(backbone, quant_config, forward_loop)
|
||||
if quantized is not None:
|
||||
setattr(pipe, attr, quantized)
|
||||
_disable_known_problematic_quantizers(mtq, getattr(pipe, attr), quantize_mha=args.quantize_mha)
|
||||
|
||||
# Export weights to FP8 only AFTER all transformers are calibrated. Exporting
|
||||
# inside the loop would convert transformer's weights to real FP8, and the
|
||||
# next transformer's forward_loop (full pipeline) would then fake-quantize an
|
||||
# already-FP8 weight -> "fake_e4m3fy not implemented for Float8_e4m3fn".
|
||||
print("\nForcing FP8 weight serialization (Wan2.2 isn't in ModelOpt's recognized-model registry,")
|
||||
print("so we have to call the per-weight export helper ourselves)...")
|
||||
total_exported = 0
|
||||
for attr, backbone in _list_transformers(pipe):
|
||||
exported = _force_export_quantized_weights(backbone, dtype)
|
||||
print(f" -> {exported} weights converted to FP8 in {attr}")
|
||||
total_exported += exported
|
||||
|
||||
if total_exported == 0:
|
||||
raise SystemExit(
|
||||
"No quantized weights were exported. Calibration may have skipped every layer "
|
||||
"(check the disable_quantizer regex) or `mtq.quantize` did not actually wrap any "
|
||||
"weight quantizers."
|
||||
)
|
||||
|
||||
print("\nSaving pipeline with FP8 transformer(s)...")
|
||||
_save_pipeline_with_fp8_transformer(pipe, model_path, output_dir)
|
||||
_patch_quant_config(output_dir, weight_block_size=weight_block_size)
|
||||
print(f"Saved to: {output_dir}")
|
||||
_summarize_export(output_dir)
|
||||
|
||||
print("\nNext: validate the checkpoint with vllm-omni:")
|
||||
print(
|
||||
" python examples/offline_inference/text_to_video/text_to_video.py \\\n"
|
||||
f" --model {output_dir} \\\n"
|
||||
" --quantization fp8 \\\n"
|
||||
" --prompt 'A dog running across a field of golden wheat.' \\\n"
|
||||
f" --height {args.height} --width {args.width} --num-frames {args.num_frames} \\\n"
|
||||
" --num-inference-steps 30 --guidance-scale 5.0 --seed 42 \\\n"
|
||||
" --output outputs/wan22_modelopt_fp8.mp4"
|
||||
)
|
||||
print(
|
||||
"\n (--quantization fp8 is auto-upgraded to ModelOpt FP8 at runtime because the "
|
||||
"checkpoint's config.json has modelopt metadata.)"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
+12
-1
@@ -66,7 +66,8 @@ dev = [
|
||||
"voxcpm>=2.0.2", # for VoxCPM2 path-locating in tests
|
||||
"torchdiffeq>=0.2.5", # for covo_audio Test
|
||||
"FlagEmbedding", # for hunyuan_image3 tests
|
||||
"kernels<0.15.0" # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
||||
"kernels<0.15.0", # for accuracy tests comparison with diffusers output (0.15.0 incompatible with diffusers 0.38.0)
|
||||
"s3tokenizer>=0.3.0" # for tts test
|
||||
]
|
||||
|
||||
demo = [
|
||||
@@ -90,6 +91,12 @@ minicpmo = [
|
||||
"stepaudio2-minicpmo",
|
||||
]
|
||||
|
||||
# Optional Blackwell-only fused-bias FP8 GEMM; package is `quack-kernels` (imported
|
||||
# as `quack`), not the unrelated `quack`. See docs/user_guide/quantization/fp8.md.
|
||||
quack = [
|
||||
"quack-kernels>=0.3.11",
|
||||
]
|
||||
|
||||
docs = [
|
||||
"mkdocs>=1.5.0",
|
||||
"mkdocs-api-autonav",
|
||||
@@ -286,5 +293,9 @@ fullset = "fullset"
|
||||
Vai = "Vai"
|
||||
tockens = "tockens"
|
||||
CANN = "CANN"
|
||||
# Step-Audio reasoning parser streams `</think>` tokens split across chunks
|
||||
# (e.g. "</" + "thi" + "nk>"). The fragment "thi" appears literally in its
|
||||
# parser logic and tests and is not a typo of "the"/"this".
|
||||
thi = "thi"
|
||||
ND = "ND"
|
||||
byt5 = "byt5"
|
||||
@@ -0,0 +1,131 @@
|
||||
# Higgs-Audio V3 TTS
|
||||
|
||||
> Multilingual text-to-speech with voice cloning on 1×H100
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Boson AI
|
||||
- Model: `bosonai/higgs-audio-v3-tts-4b`
|
||||
- Task: Text-to-speech synthesis with optional voice cloning (100+ languages)
|
||||
- Mode: Online serving with the OpenAI-compatible `/v1/audio/speech` API; offline batch inference
|
||||
- Maintainer: @yuekaiz
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to serve `bosonai/higgs-audio-v3-tts-4b` for high-quality
|
||||
multilingual TTS. The model generates 24 kHz speech, supports zero-shot voice
|
||||
cloning from a reference clip, and handles 100+ languages with inline control
|
||||
tokens for emotion, style, and prosody. The architecture is a ~4B Qwen3
|
||||
backbone with fused multi-codebook embedding/head (8 codebooks × 1026 vocab,
|
||||
MusicGen-style delay pattern).
|
||||
|
||||
## References
|
||||
|
||||
- Model card: [bosonai/higgs-audio-v3-tts-4b](https://huggingface.co/bosonai/higgs-audio-v3-tts-4b)
|
||||
- Offline example: [`examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py`](../../examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py)
|
||||
- Online example: [`examples/online_serving/text_to_speech/higgs_audio_v3/`](../../examples/online_serving/text_to_speech/higgs_audio_v3/)
|
||||
- Benchmark results: see Performance section below
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## GPU
|
||||
|
||||
### 1×H100 80GB
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Linux
|
||||
- Python: 3.12+
|
||||
- CUDA: 12.x
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: `36e048fd` (branch `higgs-v3`)
|
||||
|
||||
#### Command
|
||||
|
||||
**Online serving:**
|
||||
|
||||
```bash
|
||||
vllm-omni serve bosonai/higgs-audio-v3-tts-4b \
|
||||
--host 0.0.0.0 --port 8095 \
|
||||
--trust-remote-code --omni
|
||||
```
|
||||
|
||||
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:**
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
|
||||
--texts "Hello world." "The quick brown fox jumps over the lazy dog." \
|
||||
--output-dir results/higgs_v3_wavs
|
||||
```
|
||||
|
||||
**Offline voice clone:**
|
||||
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/higgs_audio_v3/end2end.py \
|
||||
--texts "Text to synthesize in the cloned voice." \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference clip." \
|
||||
--output-dir results/higgs_v3_clone
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
Basic TTS via curl:
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8095/v1/audio/speech \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "bosonai/higgs-audio-v3-tts-4b",
|
||||
"input": "Hello, how are you?"
|
||||
}' --output hello.wav
|
||||
```
|
||||
|
||||
Voice clone via Python client:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/higgs_audio_v3/batch_speech_client.py \
|
||||
--base-url http://localhost:8095 \
|
||||
--model bosonai/higgs-audio-v3-tts-4b \
|
||||
--ref-audio path/to/reference.wav \
|
||||
--ref-text "Transcript of the reference." \
|
||||
--prompts "Text to clone."
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Stage 0 (talker, ~4B) uses ~60% GPU memory; Stage 1 (codec decoder) uses ~25%.
|
||||
- Key flags: `--trust-remote-code` and `--omni` are required.
|
||||
- Output: 24 kHz mono WAV.
|
||||
- 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.
|
||||
- Known limitations:
|
||||
- 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.
|
||||
+4
-2
@@ -33,11 +33,12 @@ 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 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`cosmos3/Cosmos3-Nano.md`](./cosmos3/Cosmos3-Nano.md) | Text-to-image, text-to-video, image-to-video, video-to-video generation, text to video with sound, action policy | 1x H200 141GB / B300 |
|
||||
| [`cosmos3/Cosmos3-Super.md`](./cosmos3/Cosmos3-Super.md) | 64B T2I / T2V / I2V / V2V generation (+ optional audio) / Action policy | 8x H200/H100/A100 / 2x H200 / B300 |
|
||||
| [`OpenBMB/MiniCPM-o-4_5.md`](./OpenBMB/MiniCPM-o-4_5.md) | Online serving for omni multimodal chat (text / image / audio / video → text + 24 kHz speech) | 2x A100/H100 80GB / 3x mid-tier GPU / 8x RTX 4090 24GB |
|
||||
| [`OpenBMB/VoxCPM2.md`](./OpenBMB/VoxCPM2.md) | Online + offline TTS with native AR pipeline (48 kHz, 30+ languages) | 1x RTX 4090 24GB |
|
||||
| [`Qwen/Qwen-Image.md`](./Qwen/Qwen-Image.md) | Text-to-image serving with step-wise continuous batching replay and ModelOpt mixed FP8/NVFP4 | 1x A100 80GB / 2x B200 |
|
||||
@@ -50,6 +51,7 @@ recipes/
|
||||
| [`Tencent/HunyuanImage-3.0-Instruct.md`](./Tencent/HunyuanImage-3.0-Instruct.md) | DiT-only text-to-image serving and benchmark, including ModelOpt mixed FP8/NVFP4 | 4x H100/H800 80GB / 2x B200 |
|
||||
| [`Wan-AI/Wan2.2-I2V.md`](./Wan-AI/Wan2.2-I2V.md) | Image-to-video serving (Wan2.2 14B) | 8x Ascend NPU (A2/A3) |
|
||||
| [`Wan-AI/Wan2.2-S2V.md`](./Wan-AI/Wan2.2-S2V.md) | Speech-to-video serving (Wan2.2 14B) | 2x A100/H100 80GB |
|
||||
| [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB |
|
||||
| [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB |
|
||||
| [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB |
|
||||
|
|
||||
|
||||
@@ -0,0 +1,139 @@
|
||||
# Stable Diffusion 3.5
|
||||
|
||||
> Text-to-image serving with Stability AI's SD 3.5 family
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: Stability AI
|
||||
- Model: `stabilityai/stable-diffusion-3.5-medium`, `stabilityai/stable-diffusion-3.5-large`
|
||||
- Task: Text-to-image generation
|
||||
- Mode: Online serving with the OpenAI-compatible images API
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe when you want a lightweight, fast text-to-image model that
|
||||
fits on a single consumer or workstation GPU. SD 3.5-medium requires only
|
||||
~17 GiB VRAM and generates 1024×1024 images in ~4 seconds, making it one of
|
||||
the most accessible diffusion models in vLLM-Omni. SD 3.5-large produces
|
||||
noticeably sharper images at the cost of ~34 GiB VRAM and ~21 seconds per
|
||||
image.
|
||||
|
||||
## References
|
||||
|
||||
- Model (medium): <https://huggingface.co/stabilityai/stable-diffusion-3.5-medium>
|
||||
- Model (large): <https://huggingface.co/stabilityai/stable-diffusion-3.5-large>
|
||||
- Related example under `examples/`:
|
||||
[`examples/online_serving/text_to_image/README.md`](../../examples/online_serving/text_to_image/README.md)
|
||||
- User guide:
|
||||
[`docs/user_guide/examples/online_serving/text_to_image.md`](../../docs/user_guide/examples/online_serving/text_to_image.md)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
This recipe documents two CUDA GPU configurations on a single RTX A6000 48 GB,
|
||||
one for each model variant. Extend it with more hardware sections as community
|
||||
validation lands.
|
||||
|
||||
## GPU
|
||||
|
||||
### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-medium
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04
|
||||
- Python: 3.12
|
||||
- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve stabilityai/stable-diffusion-3.5-medium --omni --port 8091
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "stabilityai/stable-diffusion-3.5-medium",
|
||||
"prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
|
||||
"negative_prompt": "blurry, out of focus",
|
||||
"guidance_scale": 4.5,
|
||||
"n": 1,
|
||||
"size": "1024x1024"
|
||||
}' | python3 -c "
|
||||
import json, base64, sys
|
||||
data = json.load(sys.stdin)
|
||||
with open('output.png', 'wb') as f:
|
||||
f.write(base64.b64decode(data['data'][0]['b64_json']))
|
||||
print('saved output.png')
|
||||
"
|
||||
```
|
||||
|
||||
Or use the example client:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_image/openai_chat_client.py \
|
||||
--server http://localhost:8091 \
|
||||
--prompt "a cute cat sitting on a windowsill" \
|
||||
--output output.png
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Model weights ~15.6 GiB, peak VRAM ~17.0 GiB (no CFG), ~20.9 GiB with `guidance_scale > 1` (CFG enabled).
|
||||
- Generation time: ~4 seconds without CFG, ~8 seconds with `guidance_scale=4.5`.
|
||||
- Model loading time: ~55 seconds (including weight download on first run).
|
||||
- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
|
||||
|
||||
### 1x NVIDIA RTX A6000 48GB — stable-diffusion-3.5-large
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04
|
||||
- Python: 3.12
|
||||
- Driver / runtime: NVIDIA driver 580.126.09, CUDA 13.0
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: 0.22.0rc2.dev8+g41a795b5d
|
||||
|
||||
#### Command
|
||||
|
||||
```bash
|
||||
vllm serve stabilityai/stable-diffusion-3.5-large --omni --port 8091
|
||||
```
|
||||
|
||||
> **Note**: `stable-diffusion-3.5-large` is a gated model. Accept the license
|
||||
> on [Hugging Face](https://huggingface.co/stabilityai/stable-diffusion-3.5-large)
|
||||
> and set `HF_TOKEN` before running.
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
curl -X POST http://localhost:8091/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"model": "stabilityai/stable-diffusion-3.5-large",
|
||||
"prompt": "a cute cat sitting on a windowsill, highly detailed, sharp focus",
|
||||
"negative_prompt": "blurry, out of focus",
|
||||
"guidance_scale": 4.5,
|
||||
"n": 1,
|
||||
"size": "1024x1024"
|
||||
}' | python3 -c "
|
||||
import json, base64, sys
|
||||
data = json.load(sys.stdin)
|
||||
with open('output.png', 'wb') as f:
|
||||
f.write(base64.b64decode(data['data'][0]['b64_json']))
|
||||
print('saved output.png')
|
||||
"
|
||||
```
|
||||
|
||||
#### Notes
|
||||
|
||||
- Memory usage: Peak VRAM ~31.6 GiB at 1024×1024, ~45.5 GiB at 2048×2048 with `guidance_scale=4.5`.
|
||||
- Generation time: ~19 seconds at 1024×1024, ~102 seconds at 2048×2048.
|
||||
- Image quality is noticeably better than medium, with sharper details.
|
||||
- `negative_prompt` requires `guidance_scale > 1` to take effect; SD3 defaults to 1.0 (no CFG).
|
||||
- 2048×2048 fits within 48 GB but leaves little headroom (~45.5 GiB used out of 49 GiB).
|
||||
@@ -171,6 +171,34 @@ python examples/offline_inference/speech_to_video/speech_to_video.py \
|
||||
```
|
||||
|
||||
|
||||
### Online Serving
|
||||
|
||||
#### Server
|
||||
|
||||
```bash
|
||||
vllm serve Wan-AI/Wan2.2-S2V-14B --omni \
|
||||
--model-class-name WanS2VPipeline \
|
||||
--tensor-parallel-size 2 \
|
||||
--flow-shift 3.0 \
|
||||
--vae-use-slicing --vae-use-tiling \
|
||||
--port 8091
|
||||
```
|
||||
|
||||
#### Client
|
||||
|
||||
```bash
|
||||
no_proxy=127.0.0.1 \
|
||||
curl -X POST http://127.0.0.1:8091/v1/videos/sync \
|
||||
-F "prompt=A person singing" \
|
||||
-F 'image_reference={"image_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.png"}' \
|
||||
-F 'audio_reference={"audio_url": "https://raw.githubusercontent.com/Wan-Video/Wan2.2/main/examples/Five%20Hundred%20Miles.MP3"}' \
|
||||
-F "width=832" -F "height=480" \
|
||||
-F "num_inference_steps=40" \
|
||||
-F "guidance_scale=4.5" \
|
||||
-F "fps=16" \
|
||||
--output s2v_480p_serve.mp4
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
The output video should show a talking/singing person matching the reference
|
||||
|
||||
@@ -1,34 +1,43 @@
|
||||
# Cosmos3-Nano
|
||||
|
||||
> Text-to-image, text-to-video, and image-to-video serving
|
||||
> Text-to-image, text-to-video, image-to-video, and video-to-video serving
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: `nvidia/Cosmos3-Nano`
|
||||
- Task: Text-to-image (T2I), text-to-video (T2V), and image-to-video (I2V) generation, with optional synchronized audio (video + sound), action policy
|
||||
- Task: Text-to-image (T2I), text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) generation, with optional synchronized audio (video + sound), action policy
|
||||
- Mode: Online serving with the OpenAI-compatible image/video APIs, plus offline generation via the `Omni` API
|
||||
- Maintainer: Community
|
||||
|
||||
## When to use this recipe
|
||||
|
||||
Use this recipe to deploy `nvidia/Cosmos3-Nano` for image and video generation.
|
||||
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves all three modes;
|
||||
the mode is selected per request:
|
||||
A single pipeline class (`Cosmos3OmniDiffusersPipeline`) serves these modes; the
|
||||
mode is selected per request:
|
||||
|
||||
- **T2I** — `POST /v1/images/generations` (or a prompt carrying `modalities=["image"]`).
|
||||
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image.
|
||||
- **T2V** — `POST /v1/videos/sync` with `num_frames > 1` and no reference image/video.
|
||||
- **I2V** — `POST /v1/videos/sync` with a reference image (`input_reference` file
|
||||
upload, or `image_reference` JSON).
|
||||
- **V2V** — `POST /v1/videos/sync` with a reference video (`input_reference`
|
||||
video upload, or `video_reference` JSON). Cosmos3 conditions on selected
|
||||
reference-video latent frames; use `extra_params.condition_frame_indexes_vision`
|
||||
and `extra_params.condition_video_keep` to choose which prefix/tail frames guide
|
||||
generation.
|
||||
- **T2VS / I2VS** — add `generate_sound=true` (and optional `sound_duration`) to a
|
||||
T2V/I2V `/v1/videos/sync` request to also generate synchronized audio, muxed into
|
||||
the mp4 as AAC 48 kHz stereo. See the official model card's "Video + Audio" examples.
|
||||
- **Action** — pass `extra_params={"action_mode": ...}` to drive Physical-AI tasks:
|
||||
- `forward_dynamics` — given a first frame **and** an action trajectory, roll out
|
||||
the resulting video. Synchronous: `POST /v1/videos/sync`.
|
||||
- `policy` — given a first frame and a language instruction, **predict** the action
|
||||
trajectory (and a rollout video). Use the async `POST /v1/videos` endpoint and
|
||||
read the predicted action from the top-level `action` field
|
||||
- `forward_dynamics` — given a first frame or video **and** an action trajectory,
|
||||
roll out the resulting video. Synchronous: `POST /v1/videos/sync`.
|
||||
- `policy` — given a first frame or video and a language instruction,
|
||||
**predict** the action trajectory (and a rollout video). Use the async
|
||||
`POST /v1/videos` endpoint and read the predicted action from the top-level
|
||||
`action` field.
|
||||
- `inverse_dynamics` — given a video, **recover** the action trajectory. Use
|
||||
the async `POST /v1/videos` endpoint and read the recovered action from
|
||||
the top-level `action` field
|
||||
(`{data, shape, dtype, raw_action_dim, domain_id}`).
|
||||
|
||||
Action requests also take `domain_name` (e.g. `av`, `bridge_orig_lerobot`,
|
||||
@@ -38,8 +47,9 @@ the mode is selected per request:
|
||||
**`nvidia/Cosmos3-Nano-Policy-DROID`** is served the same way
|
||||
(`domain_name=droid_lerobot`).
|
||||
|
||||
`inverse_dynamics` (recover the action from a given video) is supported by the
|
||||
pipeline; **online inference of inverse dynamics will be added in a follow-up MR.**
|
||||
Action requests can use `input_reference` or `video_reference` for video input.
|
||||
`policy` and `forward_dynamics` can also use an image reference; `inverse_dynamics`
|
||||
requires a video reference.
|
||||
|
||||
## References
|
||||
|
||||
@@ -142,6 +152,36 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_i2v.mp4
|
||||
|
||||
# Video-to-video -> /v1/videos/sync with an uploaded reference video.
|
||||
# By default Cosmos3 conditions on latent indexes [0, 1]. For the default
|
||||
# temporal VAE stride this decodes only the first 5 input frames.
|
||||
# The model works best when the prompt describes the actual situation happening in the video.
|
||||
# Generic prompts may create sub-standard generations.
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Nano" \
|
||||
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
|
||||
-F "negative_prompt=blurry, distorted, low quality, jittery, deformed" \
|
||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
|
||||
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
|
||||
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
||||
-F 'extra_params={"use_resolution_template":false,"use_duration_template":false,"guardrails":true,"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
|
||||
-F "seed=2222" \
|
||||
-F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
|
||||
-o cosmos3_v2v.mp4
|
||||
|
||||
# V2V can also use a JSON-safe URL/data-URL video reference. Do not combine
|
||||
# video_reference with input_reference or image_reference.
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
-H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Nano" \
|
||||
-F "prompt=Continue the same scene with smooth natural motion and consistent subjects." \
|
||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" \
|
||||
-F "num_inference_steps=35" -F "guidance_scale=6.0" \
|
||||
-F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
||||
-F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"last"}' \
|
||||
-F 'video_reference={"video_url":"https://example.com/reference.mp4"}' \
|
||||
-o cosmos3_v2v_from_url.mp4
|
||||
|
||||
# Text-to-video-with-sound
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync \
|
||||
@@ -190,6 +230,21 @@ VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
|
||||
# poll until status == completed, then:
|
||||
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
|
||||
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_policy.mp4
|
||||
|
||||
# Action — inverse dynamics (video -> recovered action trajectory).
|
||||
# Asynchronous: use the job metadata to read the recovered action.
|
||||
VIDEO_ID=$(curl -sS -X POST http://localhost:8000/v1/videos \
|
||||
-H "Accept: application/json" \
|
||||
--form-string "model=nvidia/Cosmos3-Nano" \
|
||||
--form-string "prompt=Recover the robot action trajectory from this clip." \
|
||||
-F "input_reference=@motion_clip.mp4;type=video/mp4" \
|
||||
-F "size=640x480" -F "num_frames=17" -F "fps=5" \
|
||||
-F "num_inference_steps=30" -F "guidance_scale=1.0" -F "flow_shift=5.0" \
|
||||
--form-string 'extra_params={"action_mode":"inverse_dynamics","domain_name":"bridge_orig_lerobot","raw_action_dim":10,"action_chunk_size":16}' \
|
||||
-F "seed=0" | jq -r '.id')
|
||||
# poll until status == completed, then:
|
||||
curl -sS "http://localhost:8000/v1/videos/$VIDEO_ID" | jq '.action | {shape, dtype, raw_action_dim, domain_id}'
|
||||
curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_inverse_dynamics.mp4
|
||||
```
|
||||
|
||||
#### Notes
|
||||
@@ -206,24 +261,27 @@ curl -sS -L "http://localhost:8000/v1/videos/$VIDEO_ID/content" -o cosmos3_polic
|
||||
- **Determinism:** identical seed reproduces identical output on the same
|
||||
hardware; outputs are not bit-identical across different GPU types.
|
||||
- **Supported sizes (per model card):** 256p / 480p / 720p at 16:9, 4:3, 1:1,
|
||||
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V 1280×720,
|
||||
189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
|
||||
3:4, 9:16. Defaults: T2I 1024², 50 steps, guidance 7.0; T2V/I2V/V2V
|
||||
1280×720, 189 frames, 35 steps, guidance 6.0, `flow_shift=10.0`.
|
||||
- **Key flags / params:** `--no-guardrails` (server) or
|
||||
`extra_params={"guardrails":false}` (per request) toggles safety. The
|
||||
per-request flag only takes effect when the server was launched **with**
|
||||
guardrails enabled (it cannot re-enable them on a `--no-guardrails` server).
|
||||
`use_resolution_template` / `use_duration_template` are off by default and only
|
||||
needed when not using upsampled prompts that already encode resolution/duration.
|
||||
For V2V, `condition_frame_indexes_vision` selects the clean conditioned latent
|
||||
frame indexes (default `[0, 1]`), and `condition_video_keep` selects whether the
|
||||
API decodes the first or last needed reference frames (`"first"` by default).
|
||||
- **Known limitations:**
|
||||
- Guardrails-on requires `cosmos-guardrail` **and** access to the gated
|
||||
`nvidia/Cosmos-1.0-Guardrail` repo (accept license + `HF_TOKEN`); otherwise
|
||||
the server fails at pipeline build with a gated-repo / safety-checker error.
|
||||
- A guardrail-blocked prompt currently returns HTTP 500
|
||||
(`"Guardrail blocked prompt"`).
|
||||
- Action `forward_dynamics` (sync `/v1/videos/sync`) and `policy` (async
|
||||
`/v1/videos`, returns the predicted action under the top-level `action`
|
||||
field) are supported online. **Online inference of inverse dynamics will be
|
||||
added in a follow-up MR.**
|
||||
- Action `forward_dynamics`, `policy`, and `inverse_dynamics` are supported
|
||||
online. Use async `POST /v1/videos` when you need the predicted/recovered
|
||||
action payload under the top-level `action` field; sync `/v1/videos/sync`
|
||||
returns raw MP4 bytes and does not expose action metadata in the response body.
|
||||
|
||||
### 1x GPU (Offline generation)
|
||||
|
||||
@@ -287,8 +345,9 @@ python -c "from PIL import Image; im=Image.open('cosmos3_t2i.png'); print('image
|
||||
|
||||
- Same `Cosmos3OmniDiffusersPipeline` as online; mode is chosen by
|
||||
`prompt["modalities"]` (`["image"]` → T2I, `["video"]` → T2V) plus
|
||||
`num_frames`/`fps`, and `multi_modal_data={"image": ...}` for I2V. For video,
|
||||
frames are returned in `outputs[0].request_output.images` as an
|
||||
`num_frames`/`fps`, `multi_modal_data={"image": ...}` for I2V, and
|
||||
`multi_modal_data={"video": [<PIL frames>]}` or a video tensor/array for V2V.
|
||||
For video, frames are returned in `outputs[0].request_output.images` as an
|
||||
`(B, F, H, W, 3)` array.
|
||||
- The offline entry must be guarded by `if __name__ == "__main__":` — the engine
|
||||
spawns workers with the `spawn` start method.
|
||||
@@ -1,12 +1,12 @@
|
||||
# Cosmos3-Super
|
||||
|
||||
> Frontier 64B world model: text-to-image, text-to-video, image-to-video (+ optional audio)
|
||||
> Frontier 64B world model: text-to-image, text-to-video, image-to-video, video-to-video (+ optional audio)
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: NVIDIA
|
||||
- Model: `nvidia/Cosmos3-Super` (64B; also `Cosmos3-Super-Text2Image`, `Cosmos3-Super-Image2Video`)
|
||||
- Task: T2I, T2V, I2V generation, with optional synchronized audio (video + sound)
|
||||
- Task: T2I, T2V, I2V, V2V generation, with optional synchronized audio (video + sound)
|
||||
- Mode: Online serving with the OpenAI-compatible image/video APIs
|
||||
- Maintainer: Community
|
||||
|
||||
@@ -16,8 +16,8 @@ Use this recipe to deploy the 64B `nvidia/Cosmos3-Super` for the highest-quality
|
||||
Cosmos3 generation. It shares the same `Cosmos3OmniDiffusersPipeline` and request
|
||||
formats as [Cosmos3-Nano](./Cosmos3-Nano.md) — only the checkpoint size and the
|
||||
recommended parallelism differ. Mode is selected per request (T2I →
|
||||
`/v1/images/generations`; T2V/I2V → `/v1/videos/sync`; add `generate_sound=true`
|
||||
for audio).
|
||||
`/v1/images/generations`; T2V/I2V/V2V → `/v1/videos/sync`; add
|
||||
`generate_sound=true` for audio).
|
||||
|
||||
## References
|
||||
|
||||
@@ -62,8 +62,9 @@ disable. `--enable-layerwise-offload` reduces VRAM on smaller GPUs.
|
||||
#### Verification
|
||||
|
||||
Requests are identical to Nano (see [`Cosmos3-Nano.md`](./Cosmos3-Nano.md) for full
|
||||
T2I/T2V/I2V/T2VS curls); official params: `size=1280x720, num_frames=189, fps=24,
|
||||
num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0, max_sequence_length=4096`.
|
||||
T2I/T2V/I2V/V2V/T2VS curls); official params: `size=1280x720, num_frames=189,
|
||||
fps=24, num_inference_steps=35, guidance_scale=6.0, flow_shift=10.0,
|
||||
max_sequence_length=4096`.
|
||||
|
||||
```bash
|
||||
curl http://localhost:8000/v1/models
|
||||
@@ -84,6 +85,15 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "seed=1111" -F "input_reference=@/path/to/reference.jpg;type=image/jpeg" \
|
||||
-o cosmos3_super_i2v.mp4
|
||||
|
||||
# V2V — add an uploaded reference video. condition_video_keep can be "first" or "last".
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=Continue the same scene with smooth natural motion." \
|
||||
-F "size=1280x720" -F "num_frames=189" -F "fps=24" -F "num_inference_steps=35" \
|
||||
-F "guidance_scale=6.0" -F "max_sequence_length=4096" -F "flow_shift=10.0" \
|
||||
-F 'extra_params={"condition_frame_indexes_vision":[0,1],"condition_video_keep":"first"}' \
|
||||
-F "seed=2222" -F "input_reference=@/path/to/reference.mp4;type=video/mp4" \
|
||||
-o cosmos3_super_v2v.mp4
|
||||
|
||||
# T2V + sound — add generate_sound/sound_duration (output muxes AAC 48 kHz stereo)
|
||||
curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
-F "model=nvidia/Cosmos3-Super" -F "prompt=A robot arm is cleaning a plate in the kitchen" \
|
||||
@@ -103,10 +113,12 @@ curl -sS -X POST http://localhost:8000/v1/videos/sync -H "Accept: video/mp4" \
|
||||
- T2V + sound (189 frames, 35 steps) → **~198 s**, output muxes **AAC 48 kHz stereo**
|
||||
- (NVIDIA's reference: 8×H200 @ 50 steps ≈ 55 s/video; 2×H200 @ 35 steps ≈ 3 min/video.)
|
||||
- **Memory:** ~61.5 GiB per GPU when sharded across 2 GPUs (HSDP shard 2); repo ~135 GB on disk.
|
||||
- Same generation defaults, supported sizes, and `generate_sound`/`sound_duration`
|
||||
semantics as Nano, including the **action** modality: `forward_dynamics`
|
||||
(sync `/v1/videos/sync`) and `policy` (async `/v1/videos`, predicted action under
|
||||
the top-level `action` field) — see the Cosmos3-Nano recipe for the request shape.
|
||||
Online inference of inverse dynamics will be added in a follow-up MR. Verified on
|
||||
the 64B Super under `--cfg-parallel-size 2`: async `policy` returns the predicted
|
||||
action (`[16, 10]`) and the rollout video reliably.
|
||||
- Same generation defaults, supported sizes, V2V reference-video controls
|
||||
(`condition_frame_indexes_vision`, `condition_video_keep`), and
|
||||
`generate_sound`/`sound_duration`
|
||||
semantics as Nano, including the **action** modality: `forward_dynamics`,
|
||||
`policy`, and `inverse_dynamics` — see the Cosmos3-Nano recipe for the request
|
||||
shapes. Use async `/v1/videos` when you need predicted/recovered action metadata
|
||||
under the top-level `action` field. Verified on the 64B Super under
|
||||
`--cfg-parallel-size 2`: async `policy` returns the predicted action (`[16, 10]`)
|
||||
and the rollout video reliably.
|
||||
@@ -0,0 +1,82 @@
|
||||
# Ming-omni-tts 0.5B
|
||||
|
||||
> Offline and online dense Ming text-to-speech on ROCm
|
||||
|
||||
## Summary
|
||||
|
||||
- Vendor: inclusionAI
|
||||
- Model: `inclusionAI/Ming-omni-tts-0.5B`
|
||||
- Task: Text-to-speech with style, dialect, cloning, and multi-speaker controls
|
||||
- Mode: Online serving via the OpenAI-compatible `/v1/audio/speech` API; offline inference
|
||||
- Maintainer: Community
|
||||
|
||||
## References
|
||||
|
||||
- [Model card](https://huggingface.co/inclusionAI/Ming-omni-tts-0.5B)
|
||||
- [Upstream repository](https://github.com/inclusionAI/Ming-omni-tts)
|
||||
- [Offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
|
||||
- [Online example](../../examples/online_serving/text_to_speech/ming_tts/)
|
||||
|
||||
## Hardware Support
|
||||
|
||||
## ROCm
|
||||
|
||||
### 1x AMD `gfx942`
|
||||
|
||||
#### Environment
|
||||
|
||||
- OS: Ubuntu 22.04.5 LTS, x86_64
|
||||
- Python: 3.12.13
|
||||
- ROCm / HIP: 7.2.53211
|
||||
- vLLM version: 0.22.0
|
||||
- vLLM-Omni version or commit: 0.1.dev1873 / `99c35c410`
|
||||
- Docker image: `vllm/vllm-omni-rocm:v0.22.0`
|
||||
|
||||
#### Command
|
||||
|
||||
From the vLLM-Omni repository root:
|
||||
|
||||
```bash
|
||||
docker run --rm \
|
||||
--group-add=video \
|
||||
--ipc=host \
|
||||
--cap-add=SYS_PTRACE \
|
||||
--security-opt seccomp=unconfined \
|
||||
--device /dev/kfd \
|
||||
--device /dev/dri \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
-v "$PWD":/app/vllm-omni \
|
||||
-w /app/vllm-omni \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.22.0 \
|
||||
--model inclusionAI/Ming-omni-tts-0.5B \
|
||||
--deploy-config vllm_omni/deploy/ming_tts.yaml \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--enforce-eager
|
||||
```
|
||||
|
||||
#### Verification
|
||||
|
||||
```bash
|
||||
python examples/online_serving/text_to_speech/ming_tts/openai_speech_client.py \
|
||||
--text "我觉得社会企业同个人都有责任" \
|
||||
--instruction-json '{"方言":"广粤话"}' \
|
||||
--ref-audio /path/to/yue_prompt.wav \
|
||||
--max-new-tokens 200 \
|
||||
--output dialect.wav
|
||||
```
|
||||
|
||||
`--ref-audio` matches upstream `use_spk_emb=True`; do not add `--ref-text`
|
||||
for the dialect case.
|
||||
|
||||
## Notes
|
||||
|
||||
- The official ROCm image includes the platform dependencies.
|
||||
- See the [ROCm installation guide](../../docs/getting_started/installation/gpu.md)
|
||||
for interactive and source-build workflows.
|
||||
- The [offline example](../../examples/offline_inference/text_to_speech/ming_tts/)
|
||||
can be run from an interactive container.
|
||||
- The tested environment uses `--enforce-eager`.
|
||||
- Output is mono 44.1 kHz audio.
|
||||
@@ -22,4 +22,3 @@ prettytable>=3.8.0
|
||||
aenum==3.1.16
|
||||
pyzmq>=25.0.0
|
||||
janus>=1.0.0
|
||||
pydub
|
||||
@@ -1,2 +1,3 @@
|
||||
-r common.txt
|
||||
onnxruntime>=1.23.2
|
||||
auto-round-lib
|
||||
@@ -20,6 +20,7 @@ from vllm.sampling_params import SamplingParams
|
||||
from vllm.v1.engine import EngineCoreEventType
|
||||
from vllm.v1.request import Request, RequestStatus, StreamingUpdate
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
from vllm_omni.core.sched.omni_scheduler_mixin import OmniSchedulerMixin
|
||||
|
||||
# isort: on
|
||||
|
||||
@@ -142,3 +143,300 @@ class TestReplaceSessionWithStreamingUpdate:
|
||||
|
||||
assert session.events
|
||||
assert session.events[-1].type == EngineCoreEventType.QUEUED
|
||||
|
||||
|
||||
class _RealignSchedulerStub(OmniSchedulerMixin):
|
||||
"""Minimal scheduler surface for ``_realign_request_status_to_queues``.
|
||||
|
||||
Pins ``self.requests`` / ``self.running`` / ``self.waiting`` to the
|
||||
same shapes the upstream ``Scheduler.finish_requests`` reads, so the
|
||||
helper sees realistic state without spinning a real scheduler.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
*,
|
||||
requests: dict,
|
||||
running: list,
|
||||
waiting: list,
|
||||
) -> None:
|
||||
self.requests = requests
|
||||
self.running = running
|
||||
self.waiting = waiting
|
||||
|
||||
|
||||
class TestRealignRequestStatusToQueues:
|
||||
"""Regression for the residual hang described in
|
||||
https://github.com/vllm-project/vllm-omni/pull/3774 -- chunk-transfer
|
||||
adapter's ``requests_origin_status`` table goes stale on the
|
||||
``waiting → running`` admit transition, and an abort that lands
|
||||
before the next deque round-trip stomps stale ``WAITING`` onto a
|
||||
request that actually lives in ``self.running``. After
|
||||
``max_num_seqs`` such aborts every ``input_batch`` slot is leaked and
|
||||
new requests hang at ``chunks=0``.
|
||||
"""
|
||||
|
||||
def test_running_with_stale_waiting_status_is_realigned_to_running(
|
||||
self,
|
||||
) -> None:
|
||||
"""The exact race the hang reproduces: a request lives in
|
||||
``self.running`` but its ``status`` is ``WAITING``. After
|
||||
realign, status must be ``RUNNING`` so upstream
|
||||
``Scheduler.finish_requests`` removes it from ``self.running``.
|
||||
"""
|
||||
req = _make_request(request_id="req-stale")
|
||||
req.status = RequestStatus.WAITING # stale -- actually in running
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[req],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.RUNNING
|
||||
|
||||
def test_waiting_with_stale_running_status_is_realigned_to_waiting(
|
||||
self,
|
||||
) -> None:
|
||||
"""Symmetric case: request lives in ``self.waiting`` but status
|
||||
is ``RUNNING``. Upstream's else branch should run, so realign to
|
||||
``WAITING``.
|
||||
"""
|
||||
req = _make_request(request_id="req-stale-2")
|
||||
req.status = RequestStatus.RUNNING # stale -- actually in waiting
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[],
|
||||
waiting=[req],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.WAITING
|
||||
|
||||
def test_already_aligned_status_is_left_unchanged(self) -> None:
|
||||
"""Healthy case: status matches actual queue. No status mutation
|
||||
means no spurious side effects.
|
||||
"""
|
||||
req_running = _make_request(request_id="req-r")
|
||||
req_running.status = RequestStatus.RUNNING
|
||||
req_waiting = _make_request(request_id="req-w")
|
||||
req_waiting.status = RequestStatus.WAITING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
req_running.request_id: req_running,
|
||||
req_waiting.request_id: req_waiting,
|
||||
},
|
||||
running=[req_running],
|
||||
waiting=[req_waiting],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req_running.request_id, req_waiting.request_id])
|
||||
|
||||
assert req_running.status == RequestStatus.RUNNING
|
||||
assert req_waiting.status == RequestStatus.WAITING
|
||||
|
||||
def test_unknown_request_id_is_skipped_silently(self) -> None:
|
||||
"""Aligning ids that aren't in ``self.requests`` (already freed
|
||||
upstream, or never existed) must be a no-op.
|
||||
"""
|
||||
sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
|
||||
sched._realign_request_status_to_queues(["never-existed"]) # no raise
|
||||
|
||||
def test_request_in_neither_queue_is_left_unchanged(self) -> None:
|
||||
"""If a tracked request is in neither queue (e.g. parked in a
|
||||
chunk-transfer deque), realign must not invent a status. The
|
||||
adapter / deque purge owns that surface; the realign helper is
|
||||
only for the admit-transition staleness.
|
||||
"""
|
||||
req = _make_request(request_id="req-parked")
|
||||
req.status = RequestStatus.WAITING_FOR_CHUNK
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
|
||||
def test_request_ids_str_is_treated_as_single_id(self) -> None:
|
||||
"""Match upstream ``Scheduler.finish_requests`` resolution: a
|
||||
bare string is treated as one id, not iterated as characters.
|
||||
"""
|
||||
req = _make_request(request_id="req-s")
|
||||
req.status = RequestStatus.WAITING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[req],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues(req.request_id)
|
||||
|
||||
assert req.status == RequestStatus.RUNNING
|
||||
|
||||
def test_request_ids_none_aligns_every_known_request(self) -> None:
|
||||
"""``request_ids=None`` matches upstream's "all requests" path.
|
||||
Realign must walk every entry in ``self.requests`` and fix any
|
||||
stale status it finds.
|
||||
"""
|
||||
stale = _make_request(request_id="req-stale-none")
|
||||
stale.status = RequestStatus.WAITING # but actually in running
|
||||
|
||||
clean = _make_request(request_id="req-clean-none")
|
||||
clean.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
stale.request_id: stale,
|
||||
clean.request_id: clean,
|
||||
},
|
||||
running=[stale, clean],
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues(None)
|
||||
|
||||
assert stale.status == RequestStatus.RUNNING
|
||||
assert clean.status == RequestStatus.RUNNING
|
||||
|
||||
def test_finished_request_is_skipped(self) -> None:
|
||||
"""Already-finished requests must not be touched -- they may
|
||||
have legitimate finished statuses (FINISHED_STOPPED etc.) that
|
||||
a status flip would corrupt.
|
||||
"""
|
||||
req = _make_request(request_id="req-finished")
|
||||
req.status = RequestStatus.FINISHED_STOPPED
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={req.request_id: req},
|
||||
running=[req], # not realistic, but we want to prove the guard fires
|
||||
waiting=[],
|
||||
)
|
||||
sched._realign_request_status_to_queues([req.request_id])
|
||||
|
||||
assert req.status == RequestStatus.FINISHED_STOPPED
|
||||
|
||||
|
||||
class TestPurgeFinishedFromRunning:
|
||||
"""Regression for the residual ``self.running`` slot leak surface
|
||||
paired with ``_realign_request_status_to_queues``: even after
|
||||
realign + ``super().finish_requests`` runs, corner cases can leave
|
||||
already-finished or untracked entries in ``self.running`` -- e.g.
|
||||
a connector cleanup that pops from ``self.requests`` without
|
||||
unwinding ``self.running``, or a ``status`` mid-transition when
|
||||
finish ran. The defensive post-finish purge sweeps those residues
|
||||
so the worker's ``input_batch`` slot never pins a freed request.
|
||||
|
||||
See https://github.com/vllm-project/vllm-omni/pull/3774 discussion
|
||||
and the residual-hang reproduction in that PR.
|
||||
"""
|
||||
|
||||
def test_finished_request_is_purged_from_running(self) -> None:
|
||||
"""``is_finished()`` request lingering in ``self.running`` must
|
||||
be swept so its ``input_batch`` slot is freed."""
|
||||
finished = _make_request(request_id="req-finished")
|
||||
finished.status = RequestStatus.FINISHED_STOPPED # is_finished() True
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={finished.request_id: finished},
|
||||
running=[finished],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == []
|
||||
|
||||
def test_untracked_request_is_purged_from_running(self) -> None:
|
||||
"""Request lingering in ``self.running`` but no longer present
|
||||
in ``self.requests`` (already freed by upstream / connector)
|
||||
must be swept."""
|
||||
untracked = _make_request(request_id="req-untracked")
|
||||
untracked.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={}, # already deleted from self.requests
|
||||
running=[untracked],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == []
|
||||
|
||||
def test_healthy_running_is_left_unchanged(self) -> None:
|
||||
"""Live tracked running requests must not be touched -- the
|
||||
purge is defensive, not aggressive."""
|
||||
alive = _make_request(request_id="req-alive")
|
||||
alive.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={alive.request_id: alive},
|
||||
running=[alive],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == [alive]
|
||||
|
||||
def test_empty_running_is_noop(self) -> None:
|
||||
"""Empty ``self.running`` must short-circuit cleanly."""
|
||||
sched = _RealignSchedulerStub(requests={}, running=[], waiting=[])
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == []
|
||||
|
||||
def test_mixed_alive_and_dead_keeps_only_alive_in_order(self) -> None:
|
||||
"""Mixed ``self.running`` -- some alive, some finished, some
|
||||
untracked. Sweep keeps only the alive ones and preserves their
|
||||
relative order."""
|
||||
alive_a = _make_request(request_id="req-alive-a")
|
||||
alive_a.status = RequestStatus.RUNNING
|
||||
finished_b = _make_request(request_id="req-finished-b")
|
||||
finished_b.status = RequestStatus.FINISHED_STOPPED
|
||||
untracked_c = _make_request(request_id="req-untracked-c")
|
||||
untracked_c.status = RequestStatus.RUNNING
|
||||
alive_d = _make_request(request_id="req-alive-d")
|
||||
alive_d.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
alive_a.request_id: alive_a,
|
||||
finished_b.request_id: finished_b, # tracked but is_finished
|
||||
# untracked_c is NOT in self.requests
|
||||
alive_d.request_id: alive_d,
|
||||
},
|
||||
running=[alive_a, finished_b, untracked_c, alive_d],
|
||||
waiting=[],
|
||||
)
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert sched.running == [alive_a, alive_d]
|
||||
|
||||
def test_in_place_mutation_keeps_helper_local_list_identity(self) -> None:
|
||||
"""The slice-assign form (``self.running[:] = ...``) avoids an
|
||||
extra rebind inside this helper. Note that across a full
|
||||
``finish_requests`` call upstream ``Scheduler.finish_requests``
|
||||
already rebinds ``self.running``, so global identity is not a
|
||||
contract -- this test only pins the local helper behavior so a
|
||||
future refactor away from in-place mutation is intentional.
|
||||
"""
|
||||
finished = _make_request(request_id="req-finished")
|
||||
finished.status = RequestStatus.FINISHED_STOPPED
|
||||
alive = _make_request(request_id="req-alive")
|
||||
alive.status = RequestStatus.RUNNING
|
||||
|
||||
sched = _RealignSchedulerStub(
|
||||
requests={
|
||||
finished.request_id: finished,
|
||||
alive.request_id: alive,
|
||||
},
|
||||
running=[finished, alive],
|
||||
waiting=[],
|
||||
)
|
||||
held_ref = sched.running
|
||||
sched._purge_finished_from_running()
|
||||
|
||||
assert held_ref is sched.running
|
||||
assert held_ref == [alive]
|
||||
@@ -0,0 +1,158 @@
|
||||
"""Integration tests for the post-finish defensive purge of
|
||||
``self.running`` placed in
|
||||
``OmniARScheduler.finish_requests`` /
|
||||
``OmniGenerationScheduler.finish_requests``.
|
||||
|
||||
Pins *placement* — that ``_purge_finished_from_running`` runs AFTER
|
||||
``super().finish_requests`` so it catches entries upstream's removal
|
||||
missed (e.g. status mid-transition, connector cleanup popping
|
||||
``self.requests`` without unwinding ``self.running``). Direct helper
|
||||
unit tests live in ``test_omni_scheduler_mixin.py``; these tests
|
||||
exercise the full ``finish_requests`` path end-to-end.
|
||||
|
||||
Regression for the residual self.running slot leak surface paired
|
||||
with ``_realign_request_status_to_queues`` (vllm-omni#3774).
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import pytest
|
||||
from vllm.v1.request import RequestStatus
|
||||
|
||||
import vllm_omni.core.sched.omni_ar_scheduler as ar_sched_mod
|
||||
import vllm_omni.core.sched.omni_generation_scheduler as gen_sched_mod
|
||||
from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler
|
||||
from vllm_omni.core.sched.omni_generation_scheduler import OmniGenerationScheduler
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
class _StubRequest:
|
||||
"""Minimal Request stub with the surface the helpers exercise."""
|
||||
|
||||
def __init__(self, request_id: str, status: RequestStatus) -> None:
|
||||
self.request_id = request_id
|
||||
self.status = status
|
||||
|
||||
def is_finished(self) -> bool:
|
||||
return RequestStatus.is_finished(self.status)
|
||||
|
||||
|
||||
_SCHEDULER_PARAMS = [
|
||||
pytest.param(OmniARScheduler, ar_sched_mod, id="ar"),
|
||||
pytest.param(OmniGenerationScheduler, gen_sched_mod, id="generation"),
|
||||
]
|
||||
|
||||
|
||||
def _make_scheduler(scheduler_cls, *, requests, running, waiting):
|
||||
"""Bare scheduler with just the surface ``finish_requests`` reads."""
|
||||
scheduler = scheduler_cls.__new__(scheduler_cls)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = None
|
||||
scheduler.requests = requests
|
||||
scheduler.running = running
|
||||
scheduler.waiting = waiting
|
||||
return scheduler
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_purges_finished_request_from_running(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""A request whose ``status`` is a finished value but still lives
|
||||
in ``self.running`` after ``super().finish_requests`` returns must
|
||||
be swept by the post-finish purge. Mocks ``super()`` to simulate
|
||||
the corner case where upstream's status-driven removal missed the
|
||||
entry.
|
||||
"""
|
||||
finished = _StubRequest("req-finished", RequestStatus.FINISHED_STOPPED)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={finished.request_id: finished},
|
||||
running=[finished],
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
# super() returning the abort tuple but leaving self.running untouched
|
||||
return [(finished.request_id, 0)]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
result = scheduler_cls.finish_requests(
|
||||
scheduler,
|
||||
[finished.request_id],
|
||||
RequestStatus.FINISHED_STOPPED,
|
||||
)
|
||||
|
||||
assert scheduler.running == []
|
||||
assert result == [(finished.request_id, 0)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_purges_untracked_request_from_running(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""A request that has been deleted from ``self.requests`` (e.g.
|
||||
after upstream's ``_free_request`` ran but a status mismatch
|
||||
caused the ``self.running`` removal branch to no-op) must be
|
||||
swept by the post-finish purge.
|
||||
"""
|
||||
untracked = _StubRequest("req-untracked", RequestStatus.RUNNING)
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={}, # already deleted from self.requests
|
||||
running=[untracked],
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
return [(untracked.request_id, 0)]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
result = scheduler_cls.finish_requests(
|
||||
scheduler,
|
||||
[untracked.request_id],
|
||||
RequestStatus.FINISHED_STOPPED,
|
||||
)
|
||||
|
||||
assert scheduler.running == []
|
||||
assert result == [(untracked.request_id, 0)]
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("scheduler_cls", "scheduler_mod"), _SCHEDULER_PARAMS)
|
||||
def test_finish_requests_leaves_healthy_running_intact(
|
||||
monkeypatch: pytest.MonkeyPatch,
|
||||
scheduler_cls,
|
||||
scheduler_mod,
|
||||
) -> None:
|
||||
"""A live ``RUNNING`` request that is not in the finish set must
|
||||
not be swept by the defensive purge -- this pins that the predicate
|
||||
is conservative and won't aggressively drop healthy entries.
|
||||
"""
|
||||
alive = _StubRequest("req-alive", RequestStatus.RUNNING)
|
||||
leaving_id = "req-leaving"
|
||||
scheduler = _make_scheduler(
|
||||
scheduler_cls,
|
||||
requests={alive.request_id: alive},
|
||||
running=[alive], # leaving already removed by super() (typical happy path)
|
||||
waiting=[],
|
||||
)
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
return [(leaving_id, 0)]
|
||||
|
||||
monkeypatch.setattr(scheduler_mod.VLLMScheduler, "finish_requests", fake_finish_requests)
|
||||
|
||||
scheduler_cls.finish_requests(
|
||||
scheduler,
|
||||
[leaving_id],
|
||||
RequestStatus.FINISHED_STOPPED,
|
||||
)
|
||||
|
||||
assert scheduler.running == [alive]
|
||||
@@ -46,6 +46,9 @@ def test_finish_requests_cleans_input_coordinator_for_finished_ids(
|
||||
scheduler = scheduler_cls.__new__(scheduler_cls)
|
||||
scheduler.chunk_transfer_adapter = None
|
||||
scheduler.input_coordinator = coordinator
|
||||
scheduler.requests = {}
|
||||
scheduler.running = []
|
||||
scheduler.waiting = []
|
||||
|
||||
def fake_finish_requests(self, request_ids, finished_status):
|
||||
assert request_ids == ["req-a", "req-b"]
|
||||
|
||||
@@ -0,0 +1,74 @@
|
||||
[
|
||||
{
|
||||
"test_name": "test_higgs_audio_v3",
|
||||
"server_params": {
|
||||
"model": "bosonai/higgs-audio-v3-tts-4b",
|
||||
"stage_overrides": {
|
||||
"0": {
|
||||
"gpu_memory_utilization": 0.3
|
||||
},
|
||||
"1": {
|
||||
"gpu_memory_utilization": 0.2
|
||||
}
|
||||
},
|
||||
"extra_cli_args": [
|
||||
"--deploy-config",
|
||||
"vllm_omni/deploy/higgs_multimodal_qwen3.yaml",
|
||||
"--trust-remote-code"
|
||||
]
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "latency",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [
|
||||
20
|
||||
],
|
||||
"max_concurrency": [
|
||||
1
|
||||
],
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
3000
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
0.5
|
||||
]
|
||||
}
|
||||
},
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "throughput",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [
|
||||
500
|
||||
],
|
||||
"max_concurrency": [
|
||||
16
|
||||
],
|
||||
"seed_tts_locale": "en",
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [
|
||||
5000
|
||||
],
|
||||
"median_audio_rtf": [
|
||||
1.2
|
||||
],
|
||||
"audio_throughput": [
|
||||
30
|
||||
]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -339,5 +339,66 @@
|
||||
}
|
||||
}
|
||||
]
|
||||
},
|
||||
{
|
||||
"test_name": "test_voxtral_tts",
|
||||
"server_params": {
|
||||
"model": "mistralai/Voxtral-4B-TTS-2603"
|
||||
},
|
||||
"benchmark_params": [
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "latency",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [20],
|
||||
"max_concurrency": [1],
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {"voice": "casual_male"},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [150],
|
||||
"median_audio_rtf": [0.15]
|
||||
}
|
||||
},
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "throughput",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [80],
|
||||
"max_concurrency": [8],
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {"voice": "casual_male"},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [1500],
|
||||
"median_audio_rtf": [0.30],
|
||||
"audio_throughput": [30.0]
|
||||
}
|
||||
},
|
||||
{
|
||||
"task": "default_voice",
|
||||
"eval_phase": "stress",
|
||||
"dataset_name": "seed-tts-text",
|
||||
"backend": "openai-audio-speech",
|
||||
"endpoint": "/v1/audio/speech",
|
||||
"dataset_path": "benchmarks/build_dataset/seed_tts_smoke",
|
||||
"num_prompts": [100],
|
||||
"request_rate": [2.0],
|
||||
"seed_tts_locale": "en",
|
||||
"extra_body": {"voice": "casual_male"},
|
||||
"percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration",
|
||||
"baseline": {
|
||||
"median_audio_ttfp_ms": [3000],
|
||||
"median_audio_rtf": [0.25],
|
||||
"audio_throughput": [4.0]
|
||||
}
|
||||
}
|
||||
]
|
||||
}
|
||||
]
|
||||
@@ -198,19 +198,78 @@ def test_video_chat_stream_invalid_requests(
|
||||
# WS /v1/realtime
|
||||
# ─────────────────────────────────────────────────────────────────────────────
|
||||
|
||||
# Resolved in ``test_realtime_invalid_requests`` (skip ``session.created`` handshake).
|
||||
_REALTIME_WS_MODEL_MISMATCH = object()
|
||||
_REALTIME_WS_INVALID_AUDIO_APPEND = object()
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2)
|
||||
@pytest.mark.parametrize(
|
||||
"send_frames_spec, ws_error_code, err_message",
|
||||
[
|
||||
pytest.param("{not-json", "invalid_json", "Invalid JSON", id="invalid_json"),
|
||||
pytest.param(
|
||||
json.dumps({"type": "session.update"}),
|
||||
"invalid_event",
|
||||
("model", "Missing required field"),
|
||||
id="session_update_missing_model",
|
||||
),
|
||||
pytest.param(
|
||||
_REALTIME_WS_MODEL_MISMATCH,
|
||||
"model_not_found",
|
||||
("model", "does not exist"),
|
||||
id="session_update_model_mismatch",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "input_audio_buffer.commit", "final": False}),
|
||||
"model_not_validated",
|
||||
("session.update", "validate the model"),
|
||||
id="commit_without_model_validation",
|
||||
),
|
||||
pytest.param(
|
||||
json.dumps({"type": "unknown.event"}),
|
||||
"unknown_event",
|
||||
"Unknown event type",
|
||||
id="unknown_event_type",
|
||||
),
|
||||
pytest.param(
|
||||
_REALTIME_WS_INVALID_AUDIO_APPEND,
|
||||
"invalid_audio",
|
||||
"Invalid audio data",
|
||||
id="invalid_audio_append",
|
||||
),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _QWEN3_OMNI_SERVER, indirect=True)
|
||||
def test_realtime_rejected_when_async_chunk_enabled(
|
||||
def test_realtime_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
send_frames_spec: Any,
|
||||
ws_error_code: str,
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
if send_frames_spec is _REALTIME_WS_MODEL_MISMATCH:
|
||||
send_frames = json.dumps(
|
||||
{
|
||||
"type": "session.update",
|
||||
"model": "this-model-is-not-served-on-this-instance",
|
||||
}
|
||||
)
|
||||
elif send_frames_spec is _REALTIME_WS_INVALID_AUDIO_APPEND:
|
||||
send_frames = [
|
||||
json.dumps({"type": "session.update", "model": omni_server.model}),
|
||||
json.dumps({"type": "input_audio_buffer.append", "audio": "not-valid-base64!!!"}),
|
||||
]
|
||||
else:
|
||||
send_frames = send_frames_spec
|
||||
openai_client.send_realtime_ws_request(
|
||||
{
|
||||
"send_frames": send_frames,
|
||||
"timeout": 120,
|
||||
"ws_max_size": None,
|
||||
"ws_skip_types": ["session.created"],
|
||||
"ws_json_type": "error",
|
||||
"ws_error_code": "unsupported",
|
||||
"err_message": "async_chunk",
|
||||
"ws_error_code": ws_error_code,
|
||||
"err_message": err_message,
|
||||
}
|
||||
)
|
||||
@@ -319,6 +319,7 @@ QWEN_PARAMS = create_reliability_omni_server_params(RELIABILITY_SCENARIOS, DEPLO
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
|
||||
@@ -387,6 +388,7 @@ def test_reliability_fault_gpu_oom_error_contract_consistent_chat_speech(
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_function, openai_client_function) -> None:
|
||||
@@ -429,6 +431,7 @@ def test_reliability_fault_gpu_oom_chat_large_payload_failure(omni_server_functi
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_concurrent_pressure_failure(omni_server_function, openai_client_function) -> None:
|
||||
@@ -771,7 +774,7 @@ def test_reliability_fault_process_kill_tree_with_load_fast_fail_and_cleanup(
|
||||
|
||||
|
||||
@pytest.mark.slow
|
||||
@pytest.mark.skip(reason="OOM injection causes the serving instance to exit")
|
||||
@pytest.mark.skip(reason="issue#4285")
|
||||
@hardware_test(res={"cuda": "H100"}, num_cards=2)
|
||||
@pytest.mark.parametrize("omni_server_function", QWEN_PARAMS, indirect=True)
|
||||
def test_reliability_fault_gpu_oom_state_converges_after_fault_removed(
|
||||
|
||||
@@ -1,19 +1,46 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Tests for the DiffusersPipelineLoader.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from huggingface_hub import snapshot_download
|
||||
from vllm.config.load import LoadConfig
|
||||
|
||||
from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
|
||||
from vllm_omni.diffusion.data import OmniDiffusionConfig
|
||||
from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader
|
||||
from vllm_omni.diffusion.model_loader.gguf_adapters import get_gguf_adapter
|
||||
from vllm_omni.diffusion.models.helios import HeliosPipeline
|
||||
from vllm_omni.diffusion.registry import initialize_model
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
model_path = "hf-internal-testing/tiny-helios-modular-pipe"
|
||||
|
||||
|
||||
@pytest.fixture(scope="module")
|
||||
def prefetch_helios_model():
|
||||
"""Downloads the tiny helios model prior to running a test."""
|
||||
snapshot_download(model_path)
|
||||
|
||||
|
||||
@pytest.fixture(scope="function")
|
||||
def mock_tp_group(mocker):
|
||||
"""Mocks the tensor parallel group; this is needed to initialize the Helios model."""
|
||||
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1)
|
||||
mocker.patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_rank", return_value=0)
|
||||
mock_group = mocker.MagicMock()
|
||||
mock_group.world_size = 1
|
||||
mock_group.rank_in_group = 0
|
||||
mocker.patch("vllm.distributed.parallel_state.get_tp_group", return_value=mock_group)
|
||||
|
||||
|
||||
class _DummyPipelineModel(nn.Module):
|
||||
def __init__(self, *, source_prefix: str):
|
||||
@@ -42,7 +69,13 @@ class _DummyPipelineModel(nn.Module):
|
||||
|
||||
|
||||
def _make_loader_with_weights(weight_names: list[str]) -> DiffusersPipelineLoader:
|
||||
loader = object.__new__(DiffusersPipelineLoader)
|
||||
od_config = SimpleNamespace(
|
||||
dtype=torch.float32,
|
||||
parallel_config=SimpleNamespace(use_hsdp=False),
|
||||
quantization_config=None,
|
||||
)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
loader.counter_after_loading_weights = 0.0
|
||||
|
||||
@@ -150,16 +183,15 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
|
||||
quantization_config=None,
|
||||
)
|
||||
|
||||
loader = object.__new__(DiffusersPipelineLoader)
|
||||
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
|
||||
loader.load_weights = lambda model: None # type: ignore[assignment]
|
||||
loader._process_weights_after_loading = lambda model, target_device: None # type: ignore[assignment]
|
||||
loader._is_gguf_quantization = lambda _od_config: False # type: ignore[assignment]
|
||||
loader._is_gguf_quantization = lambda: False # type: ignore[assignment]
|
||||
|
||||
monkeypatch.setattr(loader_mod, "resolve_obj_by_qualname", lambda _name: _ConfigAwareModel)
|
||||
monkeypatch.setattr(loader_mod.torch, "device", lambda _name: _DeviceContext("cpu"))
|
||||
|
||||
model = loader.load_model(
|
||||
od_config,
|
||||
load_device="cpu",
|
||||
load_format="custom_pipeline",
|
||||
custom_pipeline_name="tests.dummy.ConfigAwarePipeline",
|
||||
@@ -168,3 +200,33 @@ def test_load_model_custom_pipeline_sets_current_diffusion_config(monkeypatch):
|
||||
assert model.captured_config is od_config
|
||||
assert model.seen_config_during_init is od_config
|
||||
assert get_current_diffusion_config_or_none() is None
|
||||
|
||||
|
||||
def test_get_all_weights(prefetch_helios_model, mock_tp_group):
|
||||
"""Ensure that get all weights on a tiny model resolves to nonempty weights."""
|
||||
od_config = OmniDiffusionConfig(
|
||||
model_class_name="HeliosPipeline",
|
||||
model=model_path,
|
||||
)
|
||||
loader = DiffusersPipelineLoader(
|
||||
load_config=LoadConfig(),
|
||||
od_config=od_config,
|
||||
)
|
||||
pipeline = HeliosPipeline(od_config=od_config)
|
||||
|
||||
weights = list(loader.get_all_weights(pipeline))
|
||||
assert len(weights) > 0
|
||||
|
||||
|
||||
def test_load_model(prefetch_helios_model, mock_tp_group):
|
||||
"""Ensure that load model creates an instance of the expected pipeline class."""
|
||||
od_config = OmniDiffusionConfig(
|
||||
model_class_name="HeliosPipeline",
|
||||
model=model_path,
|
||||
)
|
||||
loader = DiffusersPipelineLoader(
|
||||
load_config=LoadConfig(),
|
||||
od_config=od_config,
|
||||
)
|
||||
model = loader.load_model(load_device="cpu")
|
||||
assert isinstance(model, HeliosPipeline)
|
||||
@@ -62,6 +62,7 @@ def _make_loader() -> DiffusersPipelineLoader:
|
||||
ignore_patterns=["*.tmp"],
|
||||
)
|
||||
loader.od_config = None
|
||||
loader.quant_config = None
|
||||
loader.counter_before_loading_weights = 0.0
|
||||
loader.counter_after_loading_weights = 0.0
|
||||
return loader
|
||||
@@ -69,53 +70,47 @@ def _make_loader() -> DiffusersPipelineLoader:
|
||||
|
||||
def test_is_gguf_quantization_accepts_dict_config():
|
||||
loader = _make_loader()
|
||||
od_config = SimpleNamespace(
|
||||
quantization_config={
|
||||
"method": "gguf",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
)
|
||||
loader.quant_config = {
|
||||
"method": "gguf",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
|
||||
assert loader._is_gguf_quantization(od_config) is True
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_rejects_non_gguf_dict_config():
|
||||
loader = _make_loader()
|
||||
od_config = SimpleNamespace(
|
||||
quantization_config={
|
||||
"method": "fp8",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
)
|
||||
loader.quant_config = {
|
||||
"method": "fp8",
|
||||
"gguf_model": "weights.gguf",
|
||||
}
|
||||
|
||||
assert loader._is_gguf_quantization(od_config) is False
|
||||
assert loader._is_gguf_quantization() is False
|
||||
|
||||
|
||||
def test_is_gguf_quantization_requires_gguf_model_for_dict_config():
|
||||
loader = _make_loader()
|
||||
od_config = SimpleNamespace(quantization_config={"method": "gguf"})
|
||||
loader.quant_config = {"method": "gguf"}
|
||||
|
||||
with pytest.raises(ValueError, match="gguf_model"):
|
||||
loader._is_gguf_quantization(od_config)
|
||||
loader._is_gguf_quantization()
|
||||
|
||||
|
||||
def test_is_gguf_quantization_accepts_object_config():
|
||||
loader = _make_loader()
|
||||
od_config = SimpleNamespace(
|
||||
quantization_config=SimpleNamespace(
|
||||
get_name=lambda: "gguf",
|
||||
gguf_model="weights.gguf",
|
||||
)
|
||||
loader.quant_config = SimpleNamespace(
|
||||
get_name=lambda: "gguf",
|
||||
gguf_model="weights.gguf",
|
||||
)
|
||||
|
||||
assert loader._is_gguf_quantization(od_config) is True
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_is_gguf_quantization_uses_fallback_object_without_get_name():
|
||||
loader = _make_loader()
|
||||
od_config = SimpleNamespace(quantization_config=SimpleNamespace(gguf_model="weights.gguf"))
|
||||
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
|
||||
|
||||
assert loader._is_gguf_quantization(od_config) is True
|
||||
assert loader._is_gguf_quantization() is True
|
||||
|
||||
|
||||
def test_get_model_loadable_names_collects_parameters_and_buffers():
|
||||
@@ -171,16 +166,14 @@ def test_resolve_gguf_model_path_downloads_by_quant_type(monkeypatch: pytest.Mon
|
||||
|
||||
def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
loader.quant_config = SimpleNamespace(gguf_model="weights.gguf")
|
||||
loader.od_config = SimpleNamespace(revision="main")
|
||||
source = DiffusersPipelineLoader.ComponentSource(
|
||||
model_or_path="dummy",
|
||||
subfolder="transformer",
|
||||
revision=None,
|
||||
prefix="transformer.",
|
||||
)
|
||||
od_config = SimpleNamespace(
|
||||
revision="main",
|
||||
quantization_config=SimpleNamespace(gguf_model="weights.gguf"),
|
||||
)
|
||||
|
||||
class _Adapter:
|
||||
def weights_iterator(self):
|
||||
@@ -190,7 +183,7 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
|
||||
monkeypatch.setattr(loader, "_resolve_gguf_model_path", lambda gguf_model, revision: "resolved.gguf")
|
||||
monkeypatch.setattr(loader_module, "get_gguf_adapter", lambda gguf_file, model, source, od_config: _Adapter())
|
||||
|
||||
weights = list(loader._get_gguf_weights_iterator(source, object(), od_config))
|
||||
weights = list(loader._get_gguf_weights_iterator(source, object()))
|
||||
|
||||
assert weights[0][0] == "transformer.weight"
|
||||
assert torch.equal(weights[0][1], torch.ones((2, 2)))
|
||||
@@ -201,13 +194,12 @@ def test_get_gguf_weights_iterator_prefixes_source_names(monkeypatch: pytest.Mon
|
||||
def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
model = _DummyModel()
|
||||
od_config = SimpleNamespace()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model, od_config: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
|
||||
lambda source, model: iter([("transformer.weight", torch.ones((2, 2), dtype=torch.float32))]),
|
||||
)
|
||||
|
||||
def _hf_iter(source):
|
||||
@@ -222,7 +214,7 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
|
||||
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
|
||||
)
|
||||
|
||||
loaded = loader._load_weights_with_gguf(model, od_config)
|
||||
loaded = loader._load_weights_with_gguf(model)
|
||||
|
||||
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
|
||||
assert model.calls[0] == ["transformer.weight"]
|
||||
@@ -233,13 +225,12 @@ def test_load_weights_with_gguf_falls_back_only_for_missing_transformer_weights(
|
||||
def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete(monkeypatch: pytest.MonkeyPatch):
|
||||
loader = _make_loader()
|
||||
model = _DummyModel()
|
||||
od_config = SimpleNamespace()
|
||||
|
||||
monkeypatch.setattr(loader, "_get_weight_sources", lambda _model: tuple(model.weights_sources))
|
||||
monkeypatch.setattr(
|
||||
loader,
|
||||
"_get_gguf_weights_iterator",
|
||||
lambda source, model, od_config: iter(
|
||||
lambda source, model: iter(
|
||||
[
|
||||
("transformer.weight", torch.ones((2, 2), dtype=torch.float32)),
|
||||
("transformer.bias", torch.zeros(2, dtype=torch.float32)),
|
||||
@@ -262,7 +253,7 @@ def test_load_weights_with_gguf_skips_transformer_fallback_when_gguf_is_complete
|
||||
lambda _model: {"transformer.weight", "transformer.bias", "vae.weight"},
|
||||
)
|
||||
|
||||
loaded = loader._load_weights_with_gguf(model, od_config)
|
||||
loaded = loader._load_weights_with_gguf(model)
|
||||
|
||||
assert loaded == {"transformer.weight", "transformer.bias", "vae.weight"}
|
||||
assert hf_calls == ["vae"]
|
||||
@@ -374,6 +374,19 @@ def test_preprocess_i2v_image_and_action_video_inputs() -> None:
|
||||
assert tuple(additional["preprocessed_image"].shape) == (1, 3, 16, 32)
|
||||
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 3, 16, 32)
|
||||
|
||||
frames = [Image.new("RGB", (8, 4), color) for color in ("red", "green", "blue", "yellow", "purple", "black")]
|
||||
v2v = SimpleNamespace(
|
||||
prompts=[{"prompt": "Continue.", "multi_modal_data": {"video": frames}}],
|
||||
sampling_params=SimpleNamespace(
|
||||
height=16,
|
||||
width=32,
|
||||
extra_args={"condition_frame_indexes_vision": [0, 1], "condition_video_keep": "last"},
|
||||
),
|
||||
)
|
||||
additional = preprocess(v2v).prompts[0]["additional_information"]
|
||||
assert tuple(additional["preprocessed_video"].shape) == (1, 3, 5, 16, 32)
|
||||
assert additional["condition_frame_indexes_vision"] == [0, 1]
|
||||
|
||||
|
||||
def test_postprocess_handles_image_video_audio_and_validation() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.pipeline_cosmos3 import get_cosmos3_post_process_func
|
||||
@@ -464,6 +477,19 @@ def test_prepare_latents_for_video_image_sound_and_action(make_cosmos3_pipeline)
|
||||
assert velocity_mask.tolist() == [[[[[0.0]], [[1.0]]]]]
|
||||
assert image_latent.shape == (1, 2, 1, 2, 3)
|
||||
|
||||
pipeline._encode_video_tensor = lambda *args, **kwargs: torch.full((1, 2, 3, 2, 3), 6.0)
|
||||
v2v_latents, v2v_velocity_mask, v2v_condition = pipeline._prepare_latents_v2v(
|
||||
torch.zeros(1, 3, 5, 16, 24),
|
||||
16,
|
||||
24,
|
||||
9,
|
||||
torch.Generator(device="cpu").manual_seed(0),
|
||||
[0, 1],
|
||||
)
|
||||
torch.testing.assert_close(v2v_latents[:, :, 0:2], torch.full((1, 2, 2, 2, 3), 6.0))
|
||||
assert v2v_velocity_mask.tolist() == [[[[[0.0]], [[0.0]], [[1.0]]]]]
|
||||
assert v2v_condition.shape == (1, 2, 3, 2, 3)
|
||||
|
||||
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
||||
pipeline._sound_tokenizer = SimpleNamespace(
|
||||
sample_rate=10,
|
||||
@@ -670,6 +696,34 @@ class TestForwardRouting:
|
||||
)
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is velocity_mask
|
||||
|
||||
video_tensor = torch.zeros(1, 3, 5, 16, 16)
|
||||
v2v_condition = torch.full((1, 2, 2, 1, 1), 4.0)
|
||||
v2v_mask = torch.tensor([[[[[0.0]], [[1.0]]]]])
|
||||
pipeline._prepare_latents_v2v = lambda *args, **kwargs: (
|
||||
torch.zeros(1, 2, 2, 1, 1),
|
||||
v2v_mask,
|
||||
v2v_condition,
|
||||
)
|
||||
pipeline.forward(
|
||||
SimpleNamespace(
|
||||
prompts=[
|
||||
{
|
||||
"prompt": "continue",
|
||||
"modalities": ["video"],
|
||||
"additional_information": {
|
||||
"preprocessed_video": video_tensor,
|
||||
"condition_frame_indexes_vision": [0],
|
||||
},
|
||||
}
|
||||
],
|
||||
sampling_params=make_sampling_params(height=16, width=16, num_frames=5),
|
||||
)
|
||||
)
|
||||
assert captured["flow_shifts"][-1] == 10.0
|
||||
assert captured["format"]["negative_prompt"] == ""
|
||||
assert captured["diffuse_calls"][-1]["shared_kwargs"]["noisy_frame_mask"] is v2v_mask
|
||||
assert captured["diffuse_calls"][-1]["condition_latents"] is v2v_condition
|
||||
|
||||
pipeline.transformer = pipeline.transformer.__class__(latent_channel_size=2, sound_gen=True, sound_dim=3)
|
||||
sound_latents = torch.zeros(1, 3, 4)
|
||||
pipeline._resolve_sound_target_samples = lambda *args: (20, 2.0, 10)
|
||||
|
||||
@@ -131,6 +131,38 @@ def test_forward_returns_video_prediction(monkeypatch: pytest.MonkeyPatch) -> No
|
||||
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
||||
|
||||
|
||||
def test_forward_gathers_gen_tokens_before_unpatchify(monkeypatch: pytest.MonkeyPatch) -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3 import transformer_cosmos3
|
||||
|
||||
class RecordingGather(nn.Module):
|
||||
def __init__(self) -> None:
|
||||
super().__init__()
|
||||
self.calls = 0
|
||||
|
||||
def forward(self, hidden_gen: torch.Tensor) -> torch.Tensor:
|
||||
self.calls += 1
|
||||
return hidden_gen
|
||||
|
||||
monkeypatch.setattr(transformer_cosmos3, "_get_ulysses_state", lambda: (1, 0, None))
|
||||
model = transformer_cosmos3.Cosmos3VFMTransformer(
|
||||
SimpleNamespace(tf_model_config=_tiny_cosmos3_config(), dtype=torch.float32)
|
||||
)
|
||||
gather = RecordingGather()
|
||||
model.gen_sp_gather = gather
|
||||
|
||||
output = model(
|
||||
hidden_states=torch.zeros(1, 2, 1, 2, 2),
|
||||
timestep=torch.tensor([1.0]),
|
||||
text_ids=torch.tensor([[1, 2]], dtype=torch.long),
|
||||
text_mask=torch.ones(1, 2, dtype=torch.long),
|
||||
video_shape=(1, 2, 2),
|
||||
fps=24.0,
|
||||
)
|
||||
|
||||
assert gather.calls == 1
|
||||
assert tuple(output.shape) == (1, 2, 1, 2, 2)
|
||||
|
||||
|
||||
def test_sound_and_action_modules_follow_config() -> None:
|
||||
from vllm_omni.diffusion.models.cosmos3.transformer_cosmos3 import Cosmos3VFMTransformer
|
||||
|
||||
|
||||
@@ -0,0 +1,728 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Tests for multi-resolution support in HunyuanImage3.
|
||||
|
||||
Covers three areas:
|
||||
|
||||
1. ``ResolutionGroup._calc_by_step`` — training-bucket generation,
|
||||
verifying that each resolution has correct width×height and that
|
||||
indices match expected ratio ordering.
|
||||
|
||||
2. ``HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS`` — extra resolution matching
|
||||
via ``Resolution.append`` and new-entry insertion in
|
||||
``ResolutionGroup.__init__``.
|
||||
|
||||
3. ``get_cached_resolution_group`` and ``ar2diffusion`` — caching
|
||||
behaviour and correct ``ResolutionGroup`` construction for the
|
||||
AR→DiT stage transition.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import math
|
||||
from typing import Any
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import (
|
||||
HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS,
|
||||
Resolution,
|
||||
ResolutionGroup,
|
||||
get_cached_resolution_group,
|
||||
)
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# helpers
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
def _make_reso(desc: str | tuple[int, int]) -> Resolution:
|
||||
"""Build a ``Resolution`` from a ``"HxW"`` string or ``(h, w)`` tuple."""
|
||||
if isinstance(desc, str):
|
||||
h, w = desc.split("x")
|
||||
return Resolution(int(h), int(w))
|
||||
return Resolution(desc[0], desc[1])
|
||||
|
||||
|
||||
def _reso_set(rg: ResolutionGroup) -> set[tuple[int, int]]:
|
||||
"""Return all ``(h, w)`` tuples reachable through a ``ResolutionGroup``.
|
||||
|
||||
For entries that carry ``extra_res`` this includes every alternative
|
||||
resolution in that set.
|
||||
"""
|
||||
out: set[tuple[int, int]] = set()
|
||||
for r in rg.data:
|
||||
out.add((r.h, r.w))
|
||||
out |= r.extra_res
|
||||
return out
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 1. _calc_by_step — training-bucket resolution parsing
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_BASE = 1024
|
||||
_STEP = _BASE // 16 # 64
|
||||
_ALIGN = 16
|
||||
_MIN_SIDE = _BASE // 2 # 512
|
||||
_MAX_SIDE = _BASE * 2 # 2048
|
||||
|
||||
|
||||
def _build_plain_group(
|
||||
base_size: int = _BASE,
|
||||
step: int | None = None,
|
||||
align: int = _ALIGN,
|
||||
) -> ResolutionGroup:
|
||||
"""Return a ``ResolutionGroup`` with *no* extra resolutions."""
|
||||
return ResolutionGroup(base_size=base_size, step=step, align=align, extra_resolutions=None)
|
||||
|
||||
|
||||
class TestCalcByStep:
|
||||
"""Verify ``_calc_by_step`` generates the expected training buckets."""
|
||||
|
||||
# -- count & range -------------------------------------------------------
|
||||
|
||||
def test_count_for_base_1024(self):
|
||||
"""33 buckets: 1 centre + 2×16 from the stepping loops."""
|
||||
rg = _build_plain_group()
|
||||
# 1 centre + 16 portrait + 16 landscape
|
||||
assert len(rg.data) == 33, f"expected 33 training buckets, got {len(rg.data)}"
|
||||
|
||||
def test_all_sizes_within_bounds(self):
|
||||
rg = _build_plain_group()
|
||||
for r in rg.data:
|
||||
assert _MIN_SIDE <= r.h <= _MAX_SIDE, f"height {r.h} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
|
||||
assert _MIN_SIDE <= r.w <= _MAX_SIDE, f"width {r.w} outside [{_MIN_SIDE}, {_MAX_SIDE}]"
|
||||
|
||||
def test_all_sizes_aligned(self):
|
||||
rg = _build_plain_group()
|
||||
for r in rg.data:
|
||||
assert r.h % _ALIGN == 0, f"height {r.h} not {_ALIGN}-aligned"
|
||||
assert r.w % _ALIGN == 0, f"width {r.w} not {_ALIGN}-aligned"
|
||||
|
||||
# -- ratio ordering ------------------------------------------------------
|
||||
|
||||
def test_sorted_by_ratio_ascending(self):
|
||||
rg = _build_plain_group()
|
||||
ratios = [r.ratio for r in rg.data]
|
||||
for i in range(1, len(ratios)):
|
||||
assert ratios[i - 1] <= ratios[i], f"ratio[{i - 1}]={ratios[i - 1]:.4f} > ratio[{i}]={ratios[i]:.4f}"
|
||||
|
||||
def test_ratio_index_matches_numpy_array(self):
|
||||
rg = _build_plain_group()
|
||||
for i, r in enumerate(rg.data):
|
||||
assert math.isclose(rg.ratio[i], r.ratio, rel_tol=1e-9), (
|
||||
f"rg.ratio[{i}]={rg.ratio[i]:.6f} != data[{i}].ratio={r.ratio:.6f}"
|
||||
)
|
||||
|
||||
# -- key index positions -------------------------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"idx,expected_hw",
|
||||
[
|
||||
# First (most portrait) → 1:4 ratio bucket
|
||||
(0, (512, 2048)),
|
||||
# Centre square bucket
|
||||
(16, (1024, 1024)),
|
||||
# Most landscape bucket → 4:1 ratio
|
||||
(32, (2048, 512)),
|
||||
],
|
||||
)
|
||||
def test_landmark_index(self, idx, expected_hw):
|
||||
rg = _build_plain_group()
|
||||
r = rg.data[idx]
|
||||
assert (r.h, r.w) == expected_hw, f"idx {idx}: expected {expected_hw}, got ({r.h}, {r.w})"
|
||||
|
||||
def test_center_is_square(self):
|
||||
rg = _build_plain_group()
|
||||
r = rg.data[16]
|
||||
assert r.h == r.w == _BASE
|
||||
|
||||
def test_first_and_last_are_extremes(self):
|
||||
rg = _build_plain_group()
|
||||
# first → 512×2048 (ratio 0.25)
|
||||
assert rg.data[0].ratio < 0.3
|
||||
# last → 2048×512 (ratio 4.0)
|
||||
assert rg.data[-1].ratio > 3.9
|
||||
|
||||
# -- width × height matching ---------------------------------------------
|
||||
|
||||
def test_every_resolution_has_same_area_product(self):
|
||||
"""Symmetry: stepping loop produces the same product set in both
|
||||
directions, so (h × w) tuples should appear symmetrically."""
|
||||
rg = _build_plain_group()
|
||||
products = {(r.h, r.w) for r in rg.data}
|
||||
# For every (h, w) there should exist (w, h) or a nearby equivalent
|
||||
for h, w in products:
|
||||
if h == w:
|
||||
continue # square — trivially symmetric
|
||||
assert (w, h) in products or any(abs(rr.h - w) <= _STEP and abs(rr.w - h) <= _STEP for rr in rg.data), (
|
||||
f"no symmetric counterpart for {h}×{w}"
|
||||
)
|
||||
|
||||
def test_product_of_wh_divisible_by_256(self):
|
||||
"""Each bucket's area should be a multiple of 256 (16×16),
|
||||
which is the VAE down-sample grid."""
|
||||
rg = _build_plain_group()
|
||||
for r in rg.data:
|
||||
assert (r.h * r.w) % 256 == 0, f"{r.h}×{r.w} = {r.h * r.w} not divisible by 256"
|
||||
|
||||
# -- custom base_size / step / align -------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("base_size", [512, 768, 1280])
|
||||
def test_count_scales_with_base_size(self, base_size):
|
||||
step = base_size // 16
|
||||
rg = ResolutionGroup(base_size=base_size, step=step, align=16, extra_resolutions=None)
|
||||
expected = 1 + 2 * (base_size // step) # centre + 2×phases
|
||||
assert len(rg.data) == expected, f"base_size={base_size}: expected {expected}, got {len(rg.data)}"
|
||||
|
||||
def test_custom_step(self):
|
||||
"""A smaller step produces more buckets."""
|
||||
rg_small = ResolutionGroup(base_size=1024, step=32, align=16, extra_resolutions=None)
|
||||
rg_default = _build_plain_group()
|
||||
assert len(rg_small.data) > len(rg_default.data), "smaller step should produce more buckets"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 2. HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS matching
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
_EXTRA_SQUARES = {"512x512", "640x640", "768x768", "896x896"}
|
||||
_EXTRA_NON_SQUARES = {"1024x768", "1280x720", "768x1024", "720x1280"}
|
||||
|
||||
|
||||
def _build_extra_group() -> ResolutionGroup:
|
||||
return ResolutionGroup(
|
||||
base_size=_BASE,
|
||||
extra_resolutions=[_make_reso(s) for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS],
|
||||
)
|
||||
|
||||
|
||||
class TestExtraResolutionMatching:
|
||||
"""Verify extra resolutions are correctly integrated into
|
||||
the ``ResolutionGroup``."""
|
||||
|
||||
# -- square extras (ratio = 1.0) -----------------------------------------
|
||||
|
||||
def test_square_extras_appended_to_base_square(self):
|
||||
"""512², 640², 768², 896² all have ratio 1.0 and should be
|
||||
appended to the (1024, 1024) Resolution via ``append()``."""
|
||||
rg = _build_extra_group()
|
||||
# Find the 1024×1024 entry
|
||||
square_reso = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
|
||||
expected_extra = {(512, 512), (640, 640), (768, 768), (896, 896)}
|
||||
assert square_reso.extra_res == expected_extra, f"square extra_res: {square_reso.extra_res}"
|
||||
|
||||
def test_base_1024_unchanged(self):
|
||||
"""The 1024×1024 entry itself still reports h=1024, w=1024."""
|
||||
rg = _build_extra_group()
|
||||
r = next(r for r in rg.data if r.h == 1024 and r.w == 1024)
|
||||
assert r.h == 1024 and r.w == 1024
|
||||
|
||||
def test_squares_not_standalone_entries(self):
|
||||
"""Square extras must NOT appear as top-level ``Resolution`` entries;
|
||||
they live inside ``extra_res`` only."""
|
||||
rg = _build_extra_group()
|
||||
standalone = {(r.h, r.w) for r in rg.data}
|
||||
for square in [(512, 512), (640, 640), (768, 768), (896, 896)]:
|
||||
assert square not in standalone, f"{square} should be in extra_res, not a standalone entry"
|
||||
|
||||
# -- non-square extras (ratio ≠ 1.0) ------------------------------------
|
||||
|
||||
@pytest.mark.parametrize("reso_str", sorted(_EXTRA_NON_SQUARES))
|
||||
def test_non_square_extra_is_standalone_entry(self, reso_str):
|
||||
"""1024×768, 1280×720, 768×1024, 720×1280 are appended as
|
||||
new ``Resolution`` entries because no training bucket shares
|
||||
their exact ratio."""
|
||||
r = _make_reso(reso_str)
|
||||
rg = _build_extra_group()
|
||||
matching = [d for d in rg.data if d.h == r.h and d.w == r.w]
|
||||
assert len(matching) == 1, f"{reso_str} should be in rg.data exactly once"
|
||||
assert matching[0].extra_res == set(), f"{reso_str} should have empty extra_res, got {matching[0].extra_res}"
|
||||
|
||||
def test_non_square_extras_empty_extra_res(self):
|
||||
"""Non-square extras have no own extra_res."""
|
||||
rg = _build_extra_group()
|
||||
for r in rg.data:
|
||||
if (r.h, r.w) in {(1024, 768), (1280, 720), (768, 1024), (720, 1280)}:
|
||||
assert r.extra_res == set()
|
||||
|
||||
# -- total count ---------------------------------------------------------
|
||||
|
||||
def test_total_count(self):
|
||||
"""33 training + 4 non-square extras = 37."""
|
||||
rg = _build_extra_group()
|
||||
assert len(rg.data) == 37, f"expected 37, got {len(rg.data)}"
|
||||
|
||||
def test_extra_count(self):
|
||||
"""Exactly 8 extra resolutions (4 square + 4 non-square)."""
|
||||
rg = _build_extra_group()
|
||||
reso_set = _reso_set(rg)
|
||||
extra_in_set = reso_set - {(r.h, r.w) for r in _build_plain_group().data}
|
||||
assert len(extra_in_set) == 8, f"expected 8 extra, got {len(extra_in_set)}: {extra_in_set}"
|
||||
|
||||
# -- index stability -----------------------------------------------------
|
||||
|
||||
def test_indices_0_32_unchanged(self):
|
||||
"""Indices 0–32 remain the training buckets (sorted by ratio)."""
|
||||
plain = _build_plain_group()
|
||||
rg = _build_extra_group()
|
||||
for i in range(33):
|
||||
assert (rg.data[i].h, rg.data[i].w) == (plain.data[i].h, plain.data[i].w), (
|
||||
f"idx {i}: extra group has ({rg.data[i].h},{rg.data[i].w}), "
|
||||
f"expected ({plain.data[i].h},{plain.data[i].w})"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"reso_str,expected_idx",
|
||||
[
|
||||
("1024x768", 33),
|
||||
("1280x720", 34),
|
||||
("768x1024", 35),
|
||||
("720x1280", 36),
|
||||
],
|
||||
)
|
||||
def test_non_square_extra_indices(self, reso_str, expected_idx):
|
||||
"""Non-square extras are appended in order at indices 33-36."""
|
||||
r = _make_reso(reso_str)
|
||||
rg = _build_extra_group()
|
||||
idx = next(i for i, d in enumerate(rg.data) if d.h == r.h and d.w == r.w)
|
||||
assert idx == expected_idx, f"{reso_str}: expected idx {expected_idx}, got {idx}"
|
||||
|
||||
# -- ratio lookup via get_base_size_and_ratio_index ----------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"h,w,expected_idx",
|
||||
[
|
||||
(1024, 768, 33),
|
||||
(1280, 720, 34),
|
||||
(768, 1024, 35),
|
||||
(720, 1280, 36),
|
||||
(1024, 1024, 16),
|
||||
(512, 512, 16),
|
||||
(640, 640, 16),
|
||||
(768, 768, 16),
|
||||
(896, 896, 16),
|
||||
],
|
||||
)
|
||||
def test_ratio_index_for_resolutions(self, h, w, expected_idx):
|
||||
"""``get_base_size_and_ratio_index`` must return the correct
|
||||
ratio index for both training-bucket and extra resolutions."""
|
||||
rg = _build_extra_group()
|
||||
_, idx = rg.get_base_size_and_ratio_index(w, h)
|
||||
assert idx == expected_idx, f"({w},{h}): expected idx {expected_idx}, got {idx}"
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Resolution.match — area-based matching within a ratio group
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestResolutionMatch:
|
||||
"""Verify ``Resolution.match`` picks the best-fit alternative from
|
||||
``extra_res`` based on area proximity."""
|
||||
|
||||
def test_no_extra_returns_self(self):
|
||||
r = _make_reso("1024x768")
|
||||
assert r.match(800, 600) == (768, 1024)
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"target_w,target_h,expected",
|
||||
[
|
||||
# Exact area match in extra_res
|
||||
(512, 512, (512, 512)),
|
||||
(640, 640, (640, 640)),
|
||||
(768, 768, (768, 768)),
|
||||
(896, 896, (896, 896)),
|
||||
# Closer to 1024² than to any extra square
|
||||
(1024, 1024, (1024, 1024)),
|
||||
(1000, 1000, (1024, 1024)),
|
||||
# Small square picks 512²
|
||||
(256, 256, (512, 512)),
|
||||
(400, 400, (512, 512)),
|
||||
# 700² → closest to 640² (diff 129600) vs 768² (diff 99856)
|
||||
# 768*768 - 700*700 = 589824 - 490000 = 99824
|
||||
# 640*640 - 700*700 = 409600 - 490000 = -80400 → abs=80400
|
||||
# So 640² is closer in area
|
||||
(700, 700, (640, 640)),
|
||||
# 800² → closest to 768² (diff 49984) vs 896² (diff 162944)
|
||||
(800, 800, (768, 768)),
|
||||
],
|
||||
)
|
||||
def test_square_extra_match(self, target_w, target_h, expected):
|
||||
"""Build a Resolution like the 1024² entry with square extras,
|
||||
then verify area-based matching."""
|
||||
r = Resolution(1024, 1024)
|
||||
for s in [(512, 512), (640, 640), (768, 768), (896, 896)]:
|
||||
r.extra_res.add(s)
|
||||
result = r.match(target_w, target_h)
|
||||
assert result == expected, f"match({target_w},{target_h}) = {result}, expected {expected}"
|
||||
|
||||
def test_match_returns_tuple(self):
|
||||
r = Resolution(1024, 1024)
|
||||
r.extra_res.add((512, 512))
|
||||
result = r.match(512, 512)
|
||||
assert isinstance(result, tuple)
|
||||
assert len(result) == 2
|
||||
assert all(isinstance(v, int) for v in result)
|
||||
|
||||
def test_match_does_not_mutate_resolution(self):
|
||||
r = Resolution(1024, 1024)
|
||||
r.extra_res.add((512, 512))
|
||||
extra_before = set(r.extra_res)
|
||||
r.match(300, 300)
|
||||
assert r.extra_res == extra_before
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ResolutionGroup.get_target_size — end-to-end resolution selection
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetTargetSize:
|
||||
"""Verify ``get_target_size`` uses ratio-then-area logic correctly."""
|
||||
|
||||
@pytest.fixture(scope="class")
|
||||
def rg(self) -> ResolutionGroup:
|
||||
return _build_extra_group()
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"w,h,expected_min_area,expected_max_area",
|
||||
[
|
||||
# Square targets should land on square variants
|
||||
(1024, 1024, 1024 * 1024, 1024 * 1024),
|
||||
(512, 512, 512 * 512, 512 * 512),
|
||||
(640, 640, 640 * 640, 640 * 640),
|
||||
(768, 768, 768 * 768, 768 * 768),
|
||||
(896, 896, 896 * 896, 896 * 896),
|
||||
],
|
||||
)
|
||||
def test_square_target_returns_square(self, rg, w, h, expected_min_area, expected_max_area):
|
||||
rw, rh = rg.get_target_size(w, h)
|
||||
area = rw * rh
|
||||
assert expected_min_area <= area <= expected_max_area, f"({w},{h}) → ({rw},{rh}) area={area}"
|
||||
|
||||
def test_training_bucket_self_maps(self, rg):
|
||||
"""A training bucket should map to itself (or an extra_res variant
|
||||
with the same ratio)."""
|
||||
for r in rg.data:
|
||||
rw, rh = rg.get_target_size(r.w, r.h)
|
||||
# Either exact match or the ratio must be preserved
|
||||
assert abs(rh / rw - r.ratio) < 0.01, f"({r.w},{r.h}) → ({rw},{rh}) ratio drifted from {r.ratio:.4f}"
|
||||
|
||||
def test_wild_ratio_still_finds_bucket(self, rg):
|
||||
"""Even an unusual aspect ratio should map to some valid bucket."""
|
||||
for w, h in [(800, 600), (1200, 900), (1600, 400), (400, 1600)]:
|
||||
rw, rh = rg.get_target_size(w, h)
|
||||
assert rw > 0 and rh > 0
|
||||
assert rw % 16 == 0 and rh % 16 == 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# 3. get_cached_resolution_group & ar2diffusion integration
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestGetCachedResolutionGroup:
|
||||
"""Verify ``get_cached_resolution_group`` returns a correctly
|
||||
configured ``ResolutionGroup`` and respects the LRU cache."""
|
||||
|
||||
def test_returns_resolution_group(self):
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
assert isinstance(rg, ResolutionGroup)
|
||||
|
||||
def test_base_size_propagated(self):
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
assert rg.base_size == 1024
|
||||
|
||||
def test_has_extra_resolutions(self):
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
reso_set = _reso_set(rg)
|
||||
for s in HUNYUAN_IMAGE3_EXTRA_RESOLUTIONS:
|
||||
r = _make_reso(s)
|
||||
assert (r.h, r.w) in reso_set, f"{s} missing from cached group"
|
||||
|
||||
def test_cache_same_base_size_same_object(self):
|
||||
"""LRU cache: same base_size → same object identity."""
|
||||
# Clear any prior cache entries
|
||||
get_cached_resolution_group.cache_clear()
|
||||
rg1 = get_cached_resolution_group(base_size=1024)
|
||||
rg2 = get_cached_resolution_group(base_size=1024)
|
||||
assert rg1 is rg2, "same base_size must return cached object"
|
||||
|
||||
def test_cache_different_base_size_different_object(self):
|
||||
get_cached_resolution_group.cache_clear()
|
||||
rg1 = get_cached_resolution_group(base_size=1024)
|
||||
rg2 = get_cached_resolution_group(base_size=768)
|
||||
assert rg1 is not rg2, "different base_size must return different objects"
|
||||
|
||||
def test_cache_size_limit(self):
|
||||
"""Cache is bounded at 4 entries (``maxsize=4``)."""
|
||||
get_cached_resolution_group.cache_clear()
|
||||
for bs in [512, 768, 1024, 1280, 1536]:
|
||||
get_cached_resolution_group(base_size=bs)
|
||||
info = get_cached_resolution_group.cache_info()
|
||||
assert info.currsize <= 4, f"cache grew to {info.currsize}, max is 4"
|
||||
# All 5 unique base_size values caused cache misses
|
||||
assert info.misses >= 5, f"expected >=5 misses for 5 unique base_size values, got {info.misses}"
|
||||
|
||||
def test_resolution_indexing_works(self):
|
||||
"""Indices 0..len(rg)-1 must be accessible."""
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
for i in range(len(rg)):
|
||||
r = rg[i]
|
||||
assert isinstance(r, Resolution)
|
||||
assert r.h > 0 and r.w > 0
|
||||
|
||||
@pytest.mark.parametrize("base_size", [512, 1024])
|
||||
def test_ratio_index_roundtrip_via_group_length(self, base_size):
|
||||
"""Every ratio index in [0, len(rg)) must resolve to a valid
|
||||
Resolution with height/width within bounds."""
|
||||
rg = get_cached_resolution_group(base_size=base_size)
|
||||
for i in range(len(rg)):
|
||||
r = rg[i]
|
||||
assert isinstance(r.height, int) and r.height > 0
|
||||
assert isinstance(r.width, int) and r.width > 0
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# ar2diffusion — integration tests
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
# Token IDs matching the real prompt_utils values
|
||||
_RATIO_0 = 128044
|
||||
_RATIO_32 = 128076
|
||||
_RATIO_33 = 130103
|
||||
_RATIO_36 = 130106
|
||||
|
||||
|
||||
def _fake_ar_output(
|
||||
*,
|
||||
cumulative_token_ids: list[int] | None = None,
|
||||
text: str = "",
|
||||
) -> MagicMock:
|
||||
"""Minimal stub for a vLLM ``RequestOutput`` with one completion."""
|
||||
choice = MagicMock()
|
||||
choice.cumulative_token_ids = cumulative_token_ids or []
|
||||
choice.text = text
|
||||
output = MagicMock()
|
||||
output.outputs = [choice]
|
||||
return output
|
||||
|
||||
|
||||
class TestAr2DiffusionIntegration:
|
||||
"""Verify that ``ar2diffusion`` calls ``get_cached_resolution_group``
|
||||
and obtains a correctly configured ``ResolutionGroup`` whose indices map
|
||||
to the expected training-bucket or extra-resolution sizes."""
|
||||
|
||||
def _call_ar2diffusion(
|
||||
self,
|
||||
token_ids: list[int],
|
||||
prompt_kwargs: dict[str, Any] | None = None,
|
||||
) -> tuple[Any, Any]:
|
||||
"""Invoke ``ar2diffusion`` with a single fake AR output.
|
||||
|
||||
Returns ``(diffusion_inputs, ar_predicted_sizes)`` so callers can
|
||||
assert on the resolved height/width.
|
||||
"""
|
||||
from vllm_omni.model_executor.stage_input_processors.hunyuan_image3 import (
|
||||
ar2diffusion,
|
||||
)
|
||||
|
||||
pk = dict(prompt_kwargs or {})
|
||||
pk.setdefault("prompt", "a beautiful sunset")
|
||||
pk.setdefault("height", 1024)
|
||||
pk.setdefault("width", 1024)
|
||||
pk.setdefault("image_base_size", 1024)
|
||||
|
||||
# Build a prompt-like object that behaves like the real OmniTokensPrompt
|
||||
class _FakePrompt:
|
||||
def __init__(self, d):
|
||||
self._d = d
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self._d[key]
|
||||
|
||||
def get(self, key, default=None):
|
||||
return self._d.get(key, default)
|
||||
|
||||
def _asdict(self):
|
||||
return dict(self._d)
|
||||
|
||||
prompt = _FakePrompt(pk)
|
||||
ar_output = _fake_ar_output(cumulative_token_ids=token_ids)
|
||||
result = ar2diffusion(
|
||||
source_outputs=[ar_output],
|
||||
prompt=[prompt],
|
||||
requires_multimodal_data=False,
|
||||
)
|
||||
return result
|
||||
|
||||
# -- ratio token ↔ resolution resolution ---------------------------------
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"ratio_idx,expected_hw",
|
||||
[
|
||||
# Training buckets
|
||||
(0, (512, 2048)),
|
||||
(16, (1024, 1024)),
|
||||
(32, (2048, 512)),
|
||||
# Non-square extras
|
||||
(33, (1024, 768)),
|
||||
(34, (1280, 720)),
|
||||
(35, (768, 1024)),
|
||||
(36, (720, 1280)),
|
||||
],
|
||||
)
|
||||
def test_ratio_index_resolves_to_correct_size(self, ratio_idx, expected_hw):
|
||||
"""When the AR emits ``<img_ratio_N>``, the DiT input must carry
|
||||
the training-bucket width/height for that index."""
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
reso = rg[ratio_idx]
|
||||
assert (reso.height, reso.width) == expected_hw, (
|
||||
f"ratio_idx {ratio_idx}: expected {expected_hw}, got ({reso.height}, {reso.width})"
|
||||
)
|
||||
|
||||
@pytest.mark.parametrize("ratio_idx", [0, 16, 32, 33, 34, 35, 36])
|
||||
def test_ar2diffusion_extracts_height_width(self, ratio_idx):
|
||||
"""End-to-end: simulate AR output containing ``<img_ratio_{idx}>``
|
||||
and assert the DiT inputs carry the expected size."""
|
||||
rg = get_cached_resolution_group(base_size=1024)
|
||||
expected_h = rg[ratio_idx].height
|
||||
expected_w = rg[ratio_idx].width
|
||||
|
||||
token_id = _RATIO_0 + ratio_idx if ratio_idx <= 32 else _RATIO_33 + (ratio_idx - 33)
|
||||
result = self._call_ar2diffusion(
|
||||
token_ids=[150000, 150001, token_id],
|
||||
prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
diffusion_input = result[0]
|
||||
assert diffusion_input["height"] == expected_h, (
|
||||
f"ratio_idx={ratio_idx}: expected height {expected_h}, got {diffusion_input['height']}"
|
||||
)
|
||||
assert diffusion_input["width"] == expected_w, (
|
||||
f"ratio_idx={ratio_idx}: expected width {expected_w}, got {diffusion_input['width']}"
|
||||
)
|
||||
|
||||
def test_no_ratio_token_falls_back_to_prompt_size(self):
|
||||
"""If no ratio token is present, height/width from the prompt
|
||||
are kept unchanged."""
|
||||
result = self._call_ar2diffusion(
|
||||
token_ids=[150000, 150001], # no ratio tokens
|
||||
prompt_kwargs={"height": 768, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
diffusion_input = result[0]
|
||||
assert diffusion_input["height"] == 768
|
||||
assert diffusion_input["width"] == 1024
|
||||
|
||||
def test_out_of_range_ratio_index_falls_back(self):
|
||||
"""If ratio_idx ≥ len(reso_group), ar2diffusion must keep the
|
||||
prompt size (fallback path)."""
|
||||
with patch(
|
||||
"vllm_omni.model_executor.stage_input_processors.hunyuan_image3._extract_ratio_index",
|
||||
return_value=999, # out of range
|
||||
):
|
||||
result = self._call_ar2diffusion(
|
||||
token_ids=[150000],
|
||||
prompt_kwargs={"height": 512, "width": 512, "image_base_size": 1024},
|
||||
)
|
||||
assert len(result) == 1
|
||||
assert result[0]["height"] == 512
|
||||
assert result[0]["width"] == 512
|
||||
|
||||
# -- cache behaviour -----------------------------------------------------
|
||||
|
||||
def test_ar2diffusion_uses_cached_group(self):
|
||||
"""Each invocation of ``ar2diffusion`` calls
|
||||
``get_cached_resolution_group``, which returns a shared cached object."""
|
||||
get_cached_resolution_group.cache_clear()
|
||||
rg_before = get_cached_resolution_group(base_size=1024)
|
||||
|
||||
token_id = _RATIO_0 + 16 # <img_ratio_16> = 1024×1024
|
||||
self._call_ar2diffusion(
|
||||
token_ids=[token_id],
|
||||
prompt_kwargs={"height": 1024, "width": 1024, "image_base_size": 1024},
|
||||
)
|
||||
|
||||
rg_after = get_cached_resolution_group(base_size=1024)
|
||||
assert rg_before is rg_after, "get_cached_resolution_group should return the same cached object"
|
||||
|
||||
def test_correct_reso_group_used_for_base_size(self):
|
||||
"""Verify ``ar2diffusion`` constructs or retrieves the
|
||||
ResolutionGroup for the correct ``image_base_size``."""
|
||||
get_cached_resolution_group.cache_clear()
|
||||
|
||||
# First request with a non-default base_size
|
||||
token_id = _RATIO_0 # <img_ratio_0>
|
||||
self._call_ar2diffusion(
|
||||
token_ids=[token_id],
|
||||
prompt_kwargs={"height": 512, "width": 512, "image_base_size": 512},
|
||||
)
|
||||
|
||||
# The cached group for base_size=512 should exist
|
||||
rg512 = get_cached_resolution_group(base_size=512)
|
||||
assert rg512.base_size == 512
|
||||
|
||||
# And base_size=1024 should also be separately accessible
|
||||
rg1024 = get_cached_resolution_group(base_size=1024)
|
||||
assert rg1024.base_size == 1024
|
||||
assert rg512 is not rg1024
|
||||
|
||||
|
||||
# ---------------------------------------------------------------------------
|
||||
# Edge cases
|
||||
# ---------------------------------------------------------------------------
|
||||
|
||||
|
||||
class TestEdgeCases:
|
||||
def test_resolution_string_parsing(self):
|
||||
"""Resolution accepts both "HxW" strings and (h, w) tuples."""
|
||||
r1 = Resolution("512x2048")
|
||||
assert r1.h == 512 and r1.w == 2048
|
||||
|
||||
r2 = Resolution(768, 1024)
|
||||
assert r2.h == 768 and r2.w == 1024
|
||||
|
||||
r3 = Resolution(1024) # single int → square
|
||||
assert r3.h == 1024 and r3.w == 1024
|
||||
|
||||
def test_resolution_extra_res_initially_empty(self):
|
||||
r = Resolution(1024, 1024)
|
||||
assert r.extra_res == set()
|
||||
assert r.match(512, 512) == (1024, 1024)
|
||||
|
||||
def test_resolution_append_deduplicates(self):
|
||||
"""Appending the same resolution twice is a no-op on the set."""
|
||||
r = Resolution(1024, 1024)
|
||||
r.append(Resolution(512, 512))
|
||||
r.append(Resolution(512, 512))
|
||||
assert len(r.extra_res) == 1
|
||||
|
||||
def test_all_final_resolutions_are_16_aligned(self):
|
||||
"""Every reachable resolution (including extras) must be
|
||||
divisible by the VAE down-sample factor of 16."""
|
||||
rg = _build_extra_group()
|
||||
for h, w in _reso_set(rg):
|
||||
assert h % 16 == 0, f"height {h} not 16-aligned in {h}×{w}"
|
||||
assert w % 16 == 0, f"width {w} not 16-aligned in {h}×{w}"
|
||||
|
||||
def test_ratio_array_length_matches_data(self):
|
||||
rg = _build_extra_group()
|
||||
assert len(rg.ratio) == len(rg.data)
|
||||
|
||||
def test_resolution_repr(self):
|
||||
r = Resolution(1024, 1024)
|
||||
rep = repr(r)
|
||||
assert "1024x1024" in rep
|
||||
|
||||
r2 = Resolution(1024, 1024)
|
||||
r2.extra_res.add((512, 512))
|
||||
rep2 = repr(r2)
|
||||
assert "512x512" in rep2
|
||||
@@ -0,0 +1,2 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
@@ -0,0 +1,160 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for HunyuanVideo-1.5 quant_config propagation through transformer creation.
|
||||
|
||||
Tests cover:
|
||||
- HunyuanVideo15Pipeline passes quant_config to HunyuanVideo15Transformer3DModel
|
||||
- HunyuanVideo15I2VPipeline passes quant_config to HunyuanVideo15Transformer3DModel
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5 as t2v_module
|
||||
import vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v as i2v_module
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu, pytest.mark.diffusion]
|
||||
|
||||
|
||||
class TestHunyuanVideoQuantConfigPropagation:
|
||||
"""Verify quant_config is propagated to the transformer model in HunyuanVideo-1.5 pipelines."""
|
||||
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_local_device",
|
||||
return_value="cpu",
|
||||
)
|
||||
@patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.prefetch_subfolders")
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.from_pretrained_with_prefetch",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.Qwen2Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.ByT5Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.AutoConfig.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.T5EncoderModel",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.FlowMatchEulerDiscreteScheduler.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5.get_transformer_config_kwargs",
|
||||
return_value={},
|
||||
)
|
||||
def test_t2v_quant_config_passed_through(
|
||||
self,
|
||||
mock_get_kwargs,
|
||||
mock_scheduler,
|
||||
mock_t5_encoder,
|
||||
mock_auto_config,
|
||||
mock_byt5_tok,
|
||||
mock_qwen2_tok,
|
||||
mock_from_pretrained,
|
||||
mock_prefetch,
|
||||
mock_get_device,
|
||||
):
|
||||
captured_kwargs = {}
|
||||
|
||||
class FakeTransformer:
|
||||
def __init__(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
with patch.object(t2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
|
||||
fake_qc = MagicMock()
|
||||
od_config = SimpleNamespace(
|
||||
model="fake-t2v-model",
|
||||
dtype=torch.bfloat16,
|
||||
flow_shift=None,
|
||||
tf_model_config=SimpleNamespace(),
|
||||
quantization_config=fake_qc,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
|
||||
t2v_module.HunyuanVideo15Pipeline(od_config=od_config)
|
||||
|
||||
assert captured_kwargs.get("quant_config") is fake_qc
|
||||
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_local_device",
|
||||
return_value="cpu",
|
||||
)
|
||||
@patch("vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.prefetch_subfolders")
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.from_pretrained_with_prefetch",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.Qwen2Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.ByT5Tokenizer.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.AutoConfig.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.T5EncoderModel",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.SiglipImageProcessor.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.FlowMatchEulerDiscreteScheduler.from_pretrained",
|
||||
return_value=MagicMock(),
|
||||
)
|
||||
@patch(
|
||||
"vllm_omni.diffusion.models.hunyuan_video.pipeline_hunyuan_video_1_5_i2v.get_transformer_config_kwargs",
|
||||
return_value={},
|
||||
)
|
||||
def test_i2v_quant_config_passed_through(
|
||||
self,
|
||||
mock_get_kwargs,
|
||||
mock_scheduler,
|
||||
mock_siglip_processor,
|
||||
mock_t5_encoder,
|
||||
mock_auto_config,
|
||||
mock_byt5_tok,
|
||||
mock_qwen2_tok,
|
||||
mock_from_pretrained,
|
||||
mock_prefetch,
|
||||
mock_get_device,
|
||||
):
|
||||
captured_kwargs = {}
|
||||
|
||||
class FakeTransformer:
|
||||
def __init__(self, **kwargs):
|
||||
captured_kwargs.update(kwargs)
|
||||
|
||||
with patch.object(i2v_module, "HunyuanVideo15Transformer3DModel", FakeTransformer):
|
||||
fake_qc = MagicMock()
|
||||
od_config = SimpleNamespace(
|
||||
model="fake-i2v-model",
|
||||
dtype=torch.bfloat16,
|
||||
flow_shift=None,
|
||||
tf_model_config=SimpleNamespace(),
|
||||
quantization_config=fake_qc,
|
||||
enable_diffusion_pipeline_profiler=False,
|
||||
)
|
||||
|
||||
i2v_module.HunyuanVideo15I2VPipeline(od_config=od_config)
|
||||
|
||||
assert captured_kwargs.get("quant_config") is fake_qc
|
||||
@@ -72,6 +72,285 @@ class TestPipelineIndependence:
|
||||
|
||||
assert issubclass(LTX23Pipeline, CFGParallelMixin)
|
||||
|
||||
def test_ltx23_pipeline_declares_offload_components(self):
|
||||
"""LTX23Pipeline must expose LTX-2.3-specific modules to offload discovery."""
|
||||
from vllm_omni.diffusion.models.ltx2.pipeline_ltx2_3 import LTX23Pipeline
|
||||
from vllm_omni.diffusion.offloader.module_collector import ModuleDiscovery
|
||||
|
||||
pipe = object.__new__(LTX23Pipeline)
|
||||
torch.nn.Module.__init__(pipe)
|
||||
pipe.transformer = torch.nn.Linear(1, 1)
|
||||
pipe.text_encoder = torch.nn.Linear(1, 1)
|
||||
pipe.connectors = torch.nn.Linear(1, 1)
|
||||
pipe.vae = torch.nn.Linear(1, 1)
|
||||
pipe.audio_vae = torch.nn.Linear(1, 1)
|
||||
pipe.vocoder = torch.nn.Linear(1, 1)
|
||||
|
||||
modules = ModuleDiscovery.discover(pipe)
|
||||
|
||||
assert LTX23Pipeline._dit_modules == ["transformer"]
|
||||
assert LTX23Pipeline._encoder_modules == ["text_encoder", "connectors"]
|
||||
assert LTX23Pipeline._vae_modules == ["vae", "audio_vae"]
|
||||
assert LTX23Pipeline._resident_modules == ["vocoder"]
|
||||
assert modules.dit_names == ["transformer"]
|
||||
assert modules.encoder_names == ["text_encoder", "connectors"]
|
||||
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."""
|
||||
|
||||
@@ -0,0 +1,22 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
from torch import nn
|
||||
|
||||
from vllm_omni.diffusion.models.ltx2.ltx2_transformer import _make_rms_norm
|
||||
|
||||
|
||||
def test_ltx_rms_norm_no_affine_identity_weight_is_non_persistent_buffer():
|
||||
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=False)
|
||||
|
||||
assert "weight" not in dict(norm.named_parameters())
|
||||
assert "weight" in dict(norm.named_buffers())
|
||||
assert "weight" not in norm.state_dict()
|
||||
|
||||
|
||||
def test_ltx_rms_norm_affine_weight_remains_parameter():
|
||||
norm = _make_rms_norm(8, eps=1e-6, elementwise_affine=True)
|
||||
|
||||
assert isinstance(dict(norm.named_parameters())["weight"], nn.Parameter)
|
||||
assert "weight" not in dict(norm.named_buffers())
|
||||
assert "weight" in norm.state_dict()
|
||||
@@ -6,13 +6,11 @@ End-to-end tests for the unified quantization framework (PR #1764).
|
||||
Validates FP8 quantization works correctly for all supported model types:
|
||||
- Single-stage diffusion models (FLUX.1-dev, Qwen-Image, Z-Image-Turbo)
|
||||
- Multi-stage models (BAGEL: LLM + Diffusion)
|
||||
- MoT gen-only models (SenseNova-U1: gen-path FP8, und-path BF16)
|
||||
|
||||
Tests verify:
|
||||
1. FP8 quantization produces valid images
|
||||
2. Memory usage is lower than BF16 baseline
|
||||
3. Multi-stage models only quantize the diffusion stage (not the LLM stage)
|
||||
4. MoT models only quantize gen-path layers (mot_gen)
|
||||
|
||||
Usage:
|
||||
# Run all FP8 quantization tests
|
||||
@@ -438,92 +436,3 @@ def test_single_stage_quantization_config_key():
|
||||
quantization_config="fp8",
|
||||
)
|
||||
assert len(images) >= 1
|
||||
|
||||
|
||||
# ─── SenseNova-U1 gen-only FP8 tests ─────────────────────────────────────────
|
||||
|
||||
_SENSENOVA_MODEL = "SenseNova/SenseNova-U1-8B-MoT"
|
||||
|
||||
|
||||
def _generate_sensenova_u1_image(
|
||||
quantization_config: str | None = None,
|
||||
num_inference_steps: int = 20,
|
||||
) -> tuple[Any, float]:
|
||||
"""Generate an image with SenseNova-U1 (single-stage MoT, gen-only FP8).
|
||||
|
||||
Returns (generated_image, peak_memory_gib).
|
||||
"""
|
||||
omni_kwargs: dict[str, Any] = {}
|
||||
if quantization_config:
|
||||
omni_kwargs["quantization_config"] = quantization_config
|
||||
|
||||
with OmniRunner(_SENSENOVA_MODEL, **omni_kwargs) as runner:
|
||||
torch.accelerator.reset_peak_memory_stats()
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=1024,
|
||||
width=1024,
|
||||
seed=42,
|
||||
num_inference_steps=num_inference_steps,
|
||||
extra_args={
|
||||
"cfg_scale": 4.0,
|
||||
"cfg_norm": "none",
|
||||
"timestep_shift": 3.0,
|
||||
},
|
||||
)
|
||||
omni_outputs = list(
|
||||
runner.omni.generate(
|
||||
prompts={"prompt": "A cat sitting on a windowsill", "modalities": ["image"]},
|
||||
sampling_params_list=sampling_params,
|
||||
)
|
||||
)
|
||||
|
||||
peak_mem = torch.accelerator.max_memory_allocated() / (1024**3)
|
||||
|
||||
generated_image = None
|
||||
for req_output in omni_outputs:
|
||||
if images := getattr(req_output, "images", None):
|
||||
generated_image = images[0]
|
||||
break
|
||||
|
||||
assert generated_image is not None, "No images generated from SenseNova-U1"
|
||||
assert generated_image.size == (1024, 1024), f"Expected 1024x1024, got {generated_image.size}"
|
||||
|
||||
return generated_image, peak_mem
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_sensenova_u1_fp8_generates_image():
|
||||
"""SenseNova-U1 with gen-only FP8 generates a valid image.
|
||||
|
||||
FP8 is applied only to gen-path (mot_gen) layers; understanding-path
|
||||
layers stay in BF16.
|
||||
"""
|
||||
image, _ = _generate_sensenova_u1_image(quantization_config="fp8")
|
||||
image.save("test_sensenova_u1_fp8.png")
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_sensenova_u1_bf16_generates_image():
|
||||
"""SenseNova-U1 without quantization generates a valid image (baseline)."""
|
||||
image, _ = _generate_sensenova_u1_image(quantization_config=None)
|
||||
image.save("test_sensenova_u1_bf16.png")
|
||||
|
||||
|
||||
@hardware_test(res={"cuda": "H100"})
|
||||
def test_sensenova_u1_fp8_uses_less_memory():
|
||||
"""Gen-only FP8 should use less peak memory than BF16 for SenseNova-U1."""
|
||||
_, mem_bf16 = _generate_sensenova_u1_image(
|
||||
quantization_config=None,
|
||||
num_inference_steps=10,
|
||||
)
|
||||
torch.accelerator.empty_cache()
|
||||
|
||||
_, mem_fp8 = _generate_sensenova_u1_image(
|
||||
quantization_config="fp8",
|
||||
num_inference_steps=10,
|
||||
)
|
||||
|
||||
print(f"SenseNova-U1 BF16 peak memory: {mem_bf16:.2f} GiB")
|
||||
print(f"SenseNova-U1 FP8 peak memory: {mem_fp8:.2f} GiB")
|
||||
assert mem_fp8 < mem_bf16, f"FP8 ({mem_fp8:.2f} GiB) should use less memory than BF16 ({mem_bf16:.2f} GiB)"
|
||||
@@ -1,175 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""CPU tests for SenseNova-U1 gen-only FP8 quantization.
|
||||
|
||||
Validates that the ``_GenOnlyQuantConfig`` wrapper correctly routes FP8
|
||||
quantization to only the gen-path (``*_mot_gen``) layers while leaving
|
||||
understanding-path layers in BF16.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
from vllm.distributed import parallel_state
|
||||
from vllm.model_executor.layers.linear import LinearBase
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu]
|
||||
|
||||
|
||||
@pytest.fixture(autouse=True)
|
||||
def _fake_tp_group():
|
||||
"""Provide a fake TP group so vLLM parallel layers can be instantiated.
|
||||
|
||||
Note: a shared ``init_fake_tp_group`` fixture exists in
|
||||
``tests/helpers/fixtures/env.py`` but requires pytest-mock.
|
||||
"""
|
||||
mock_tp = MagicMock()
|
||||
mock_tp.world_size = 1
|
||||
mock_tp.rank_in_group = 0
|
||||
old = parallel_state._TP
|
||||
parallel_state._TP = mock_tp
|
||||
yield
|
||||
parallel_state._TP = old
|
||||
|
||||
|
||||
def _make_minimal_llm_config(num_layers: int = 2) -> SimpleNamespace:
|
||||
"""Create a minimal SenseNova-U1 LLM config for unit testing."""
|
||||
return SimpleNamespace(
|
||||
hidden_size=256,
|
||||
intermediate_size=512,
|
||||
num_hidden_layers=num_layers,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=4,
|
||||
head_dim=64,
|
||||
hidden_act="silu",
|
||||
rms_norm_eps=1e-6,
|
||||
vocab_size=1024,
|
||||
rope_theta=1000000.0,
|
||||
rope_theta_hw=10000.0,
|
||||
max_position_embeddings=4096,
|
||||
max_position_embeddings_hw=1024,
|
||||
attention_bias=True,
|
||||
layer_types=["full_attention"] * num_layers,
|
||||
tie_word_embeddings=False,
|
||||
)
|
||||
|
||||
|
||||
def _build_model_with_gen_only_quant(quant_method: str | None = "fp8"):
|
||||
"""Construct SenseNovaU1ForCausalLM with gen-only FP8 wrapping."""
|
||||
from unittest.mock import MagicMock, patch
|
||||
|
||||
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
|
||||
_GenOnlyQuantConfig,
|
||||
)
|
||||
from vllm_omni.diffusion.models.sensenova_u1.sensenova_u1_transformer import (
|
||||
SenseNovaU1ForCausalLM,
|
||||
)
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
inner_config = build_quant_config(quant_method)
|
||||
if inner_config is not None:
|
||||
quant_config = _GenOnlyQuantConfig(inner_config)
|
||||
else:
|
||||
quant_config = None
|
||||
|
||||
llm_cfg = _make_minimal_llm_config(num_layers=2)
|
||||
|
||||
mock_vllm_config = MagicMock()
|
||||
mock_vllm_config.model_config.dtype = torch.bfloat16
|
||||
mock_vllm_config.model_config.quantization = quant_method
|
||||
|
||||
with (
|
||||
torch.device("meta"),
|
||||
patch("vllm.model_executor.layers.quantization.fp8.get_current_vllm_config", return_value=mock_vllm_config),
|
||||
):
|
||||
model = SenseNovaU1ForCausalLM(
|
||||
llm_cfg,
|
||||
quant_config=quant_config,
|
||||
prefix="language_model",
|
||||
)
|
||||
return model
|
||||
|
||||
|
||||
def test_mot_gen_layers_have_fp8_quant_method():
|
||||
"""Gen-path (mot_gen) Linear layers should have a real FP8 quant_method."""
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
model = _build_model_with_gen_only_quant("fp8")
|
||||
|
||||
mot_gen_linears = [
|
||||
(name, mod) for name, mod in model.named_modules() if isinstance(mod, LinearBase) and "mot_gen" in name
|
||||
]
|
||||
|
||||
assert len(mot_gen_linears) > 0, "No mot_gen LinearBase modules found"
|
||||
|
||||
for name, mod in mot_gen_linears:
|
||||
assert mod.quant_method is not None, f"mot_gen layer {name!r} should have quant_method with FP8"
|
||||
assert not isinstance(mod.quant_method, UnquantizedLinearMethod), (
|
||||
f"mot_gen layer {name!r} should NOT use UnquantizedLinearMethod"
|
||||
)
|
||||
|
||||
|
||||
def test_und_layers_have_no_quant_method():
|
||||
"""Understanding-path Linear layers should use UnquantizedLinearMethod (stay BF16)."""
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
model = _build_model_with_gen_only_quant("fp8")
|
||||
|
||||
und_linears = [
|
||||
(name, mod) for name, mod in model.named_modules() if isinstance(mod, LinearBase) and "mot_gen" not in name
|
||||
]
|
||||
|
||||
assert len(und_linears) > 0, "No und LinearBase modules found"
|
||||
|
||||
for name, mod in und_linears:
|
||||
assert isinstance(mod.quant_method, UnquantizedLinearMethod), (
|
||||
f"und layer {name!r} should use UnquantizedLinearMethod, got {type(mod.quant_method).__name__}"
|
||||
)
|
||||
|
||||
|
||||
def test_no_quant_method_without_quantization():
|
||||
"""Without quantization, all layers should use UnquantizedLinearMethod."""
|
||||
from vllm.model_executor.layers.linear import UnquantizedLinearMethod
|
||||
|
||||
model = _build_model_with_gen_only_quant(None)
|
||||
|
||||
for name, mod in model.named_modules():
|
||||
if isinstance(mod, LinearBase):
|
||||
assert isinstance(mod.quant_method, UnquantizedLinearMethod), (
|
||||
f"LinearBase module {name!r} has quant_method="
|
||||
f"{type(mod.quant_method).__name__} "
|
||||
"but no quantization was configured"
|
||||
)
|
||||
|
||||
|
||||
def test_und_and_gen_layers_both_present():
|
||||
"""Sanity check: the model has both und and gen layers."""
|
||||
model = _build_model_with_gen_only_quant(None)
|
||||
|
||||
layer_names = [name for name, _ in model.named_modules()]
|
||||
|
||||
und_qkv = [n for n in layer_names if "qkv_proj" in n and "mot_gen" not in n]
|
||||
gen_qkv = [n for n in layer_names if "qkv_proj_mot_gen" in n]
|
||||
und_mlp = [n for n in layer_names if ".mlp." in n and "mot_gen" not in n]
|
||||
gen_mlp = [n for n in layer_names if "mlp_mot_gen" in n]
|
||||
|
||||
assert len(und_qkv) > 0, "No und qkv_proj layers found"
|
||||
assert len(gen_qkv) > 0, "No gen qkv_proj_mot_gen layers found"
|
||||
assert len(und_mlp) > 0, "No und mlp layers found"
|
||||
assert len(gen_mlp) > 0, "No gen mlp_mot_gen layers found"
|
||||
|
||||
|
||||
def test_gen_only_wrapper_delegates_attributes():
|
||||
"""_GenOnlyQuantConfig should delegate attribute access to the inner config."""
|
||||
from vllm_omni.diffusion.models.sensenova_u1.pipeline_sensenova_u1 import (
|
||||
_GenOnlyQuantConfig,
|
||||
)
|
||||
from vllm_omni.quantization import build_quant_config
|
||||
|
||||
inner = build_quant_config("fp8")
|
||||
wrapper = _GenOnlyQuantConfig(inner)
|
||||
|
||||
assert wrapper.get_name() == "fp8"
|
||||
assert wrapper.activation_scheme == inner.activation_scheme
|
||||
@@ -41,7 +41,7 @@ def mock_od_config(mocker: MockerFixture):
|
||||
config.cache_backend = None
|
||||
config.cache_config = None
|
||||
config.model = "test-model"
|
||||
config.diffusion_load_format = None
|
||||
config.diffusion_load_format = "default"
|
||||
config.dtype = "float32"
|
||||
config.max_cpu_loras = 0
|
||||
config.lora_path = None
|
||||
|
||||
@@ -145,7 +145,7 @@ def test_shm_connector_flow(mocker: MockerFixture):
|
||||
Using real SharedMemoryConnector (inline mode for simplicity).
|
||||
"""
|
||||
# 1. Setup Connector
|
||||
config = {"shm_threshold_bytes": 1024} # Large threshold to use inline
|
||||
config = {"shm_threshold_bytes": 1024, "inline_small_payloads": True}
|
||||
connector = SharedMemoryConnector(config)
|
||||
connectors_map = {("0", "1"): connector}
|
||||
|
||||
|
||||
@@ -75,18 +75,21 @@ def test_create_unknown_connector():
|
||||
|
||||
@pytest.fixture
|
||||
def shm_connector():
|
||||
config = {"shm_threshold_bytes": 100} # Small threshold for testing
|
||||
config = {"shm_threshold_bytes": 100, "inline_small_payloads": True}
|
||||
return SharedMemoryConnector(config)
|
||||
|
||||
|
||||
def test_put_get_inline(shm_connector):
|
||||
"""Test transfer for small data (currently always uses SHM)."""
|
||||
"""Test inline transfer for small data."""
|
||||
data = {"small": "data"}
|
||||
|
||||
success, size, metadata = shm_connector.put("stage_0", "stage_1", "req_1", data)
|
||||
assert success is True
|
||||
assert "shm" in metadata
|
||||
assert "inline_bytes" in metadata
|
||||
assert "shm" not in metadata
|
||||
assert "size" in metadata
|
||||
assert shm_connector._metrics["inline_writes"] == 1
|
||||
assert shm_connector._metrics["shm_writes"] == 0
|
||||
|
||||
# Retrieve
|
||||
retrieved_data, ret_size = shm_connector.get("stage_0", "stage_1", "req_1", metadata)
|
||||
|
||||
@@ -387,14 +387,15 @@ def test_process_and_restore_queues(build_adapter):
|
||||
running_req = _req("r1", RequestStatus.RUNNING)
|
||||
waiting_queue = DummyWaitingQueue([waiting_req])
|
||||
running_queue = [running_req]
|
||||
scheduler_requests = {waiting_req.request_id: waiting_req, running_req.request_id: running_req}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue)
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert waiting_req.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert running_req.status == RequestStatus.WAITING_FOR_CHUNK
|
||||
assert waiting_queue == []
|
||||
assert running_queue == []
|
||||
|
||||
adapter.restore_queues(waiting_queue, running_queue)
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert waiting_queue == [waiting_req]
|
||||
assert running_queue == [running_req]
|
||||
assert adapter.waiting_for_chunk_waiting_requests == deque()
|
||||
@@ -897,6 +898,8 @@ def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture):
|
||||
sched = OmniARScheduler.__new__(OmniARScheduler)
|
||||
sched.chunk_transfer_adapter = adapter
|
||||
sched.requests = {}
|
||||
sched.running = []
|
||||
sched.waiting = []
|
||||
|
||||
with patch.object(VLLMScheduler, "finish_requests", _super_finish):
|
||||
OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED)
|
||||
@@ -1170,3 +1173,164 @@ def test_deferred_finish_not_finished_still_emits_output(mocker: MockerFixture):
|
||||
assert eco.outputs[0].finish_reason == "stop"
|
||||
assert eco.outputs[0].kv_transfer_params is None
|
||||
assert scheduler._pending_finish_reqs == []
|
||||
|
||||
|
||||
# ---------------------------------------------------------------
|
||||
# Zombie purge tests (regression for vllm-project/vllm-omni#3736)
|
||||
# ---------------------------------------------------------------
|
||||
|
||||
|
||||
def test_process_pending_chunks_purges_zombies_in_running_deque(
|
||||
build_adapter,
|
||||
):
|
||||
"""A request aborted while parked in ``waiting_for_chunk_running_requests``
|
||||
is no longer tracked by ``scheduler.requests`` once
|
||||
``Scheduler._free_request`` runs. ``process_pending_chunks`` must drop
|
||||
that zombie *before* ``restore_queues`` would re-inject it onto the
|
||||
running queue.
|
||||
|
||||
Regression for https://github.com/vllm-project/vllm-omni/issues/3736
|
||||
(engine-core ``KeyError`` on aborts under chunk-transfer pressure).
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
live_req = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie_req = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
adapter.waiting_for_chunk_running_requests.append(live_req)
|
||||
adapter.waiting_for_chunk_running_requests.append(zombie_req)
|
||||
# Mirror state the adapter would carry for an in-flight request so we can
|
||||
# later assert ``cleanup_receiver`` actually fired against the zombie.
|
||||
adapter.requests_origin_status[zombie_req.request_id] = RequestStatus.RUNNING
|
||||
adapter.requests_origin_status[live_req.request_id] = RequestStatus.RUNNING
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
# Simulate the post-abort scheduler: only ``live-1`` is still tracked.
|
||||
scheduler_requests = {live_req.request_id: live_req}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
|
||||
# 1. Zombie was popped from the deque.
|
||||
assert zombie_req not in adapter.waiting_for_chunk_running_requests
|
||||
# 2. Live request is still in the deque.
|
||||
assert live_req in adapter.waiting_for_chunk_running_requests
|
||||
# 3. ``cleanup_receiver`` ran for the zombie (drops origin-status mapping
|
||||
# and registers the id as cancelled so a late load/poll is dropped too).
|
||||
assert zombie_req.request_id not in adapter.requests_origin_status
|
||||
assert zombie_req.request_id in adapter._cancelled_load_reqs
|
||||
# 4. Live request's bookkeeping is untouched.
|
||||
assert adapter.requests_origin_status[live_req.request_id] == RequestStatus.RUNNING
|
||||
assert live_req.request_id not in adapter._cancelled_load_reqs
|
||||
|
||||
# 5. ``restore_queues`` (which the scheduler runs in its ``finally``
|
||||
# clause) now only re-injects the live request -- the zombie is
|
||||
# gone, so the worker's ``_update_states`` cannot crash on it.
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert running_queue == [live_req]
|
||||
|
||||
|
||||
def test_process_pending_chunks_purges_zombies_in_waiting_deque(build_adapter):
|
||||
"""Zombie purge applies symmetrically to the waiting-side deque.
|
||||
|
||||
Regression for https://github.com/vllm-project/vllm-omni/issues/3736 --
|
||||
aborted requests that landed in ``waiting_for_chunk_waiting_requests``
|
||||
must also be removed before ``restore_queues`` re-injects them, since
|
||||
``restore_queues`` uses ``waiting_queue.add_request`` rather than
|
||||
``running_queue.extend`` for that path.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
live_req = _req("live-w", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie_req = _req("zombie-w", RequestStatus.WAITING_FOR_CHUNK)
|
||||
adapter.waiting_for_chunk_waiting_requests.append(live_req)
|
||||
adapter.waiting_for_chunk_waiting_requests.append(zombie_req)
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
scheduler_requests = {live_req.request_id: live_req}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
|
||||
assert live_req in adapter.waiting_for_chunk_waiting_requests
|
||||
assert zombie_req not in adapter.waiting_for_chunk_waiting_requests
|
||||
assert zombie_req.request_id in adapter._cancelled_load_reqs
|
||||
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert waiting_queue == [live_req]
|
||||
|
||||
|
||||
def test_purge_preserves_live_order_with_interleaved_zombies(build_adapter):
|
||||
"""Interleaved live/zombie ordering -- pin the ``popleft`` + append-live
|
||||
in-place filter at chunk_transfer_adapter._purge_untracked_chunk_requests:
|
||||
survivor order must match insertion order, and ``cleanup_receiver`` must
|
||||
run exactly once per zombie.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
live1 = _req("live-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie1 = _req("zombie-1", RequestStatus.WAITING_FOR_CHUNK)
|
||||
live2 = _req("live-2", RequestStatus.WAITING_FOR_CHUNK)
|
||||
zombie2 = _req("zombie-2", RequestStatus.WAITING_FOR_CHUNK)
|
||||
for req in (live1, zombie1, live2, zombie2):
|
||||
adapter.waiting_for_chunk_running_requests.append(req)
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
scheduler_requests = {live1.request_id: live1, live2.request_id: live2}
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
|
||||
assert list(adapter.waiting_for_chunk_running_requests) == [live1, live2]
|
||||
assert zombie1.request_id in adapter._cancelled_load_reqs
|
||||
assert zombie2.request_id in adapter._cancelled_load_reqs
|
||||
|
||||
|
||||
def test_restore_queues_purges_late_aborts_after_process_pending_chunks(
|
||||
build_adapter,
|
||||
):
|
||||
"""Race window guard: an abort can fire between ``process_pending_chunks``
|
||||
and the scheduler's ``finally``-clause ``restore_queues`` call. The
|
||||
second purge inside ``restore_queues`` must drop the now-untracked
|
||||
request so it does not get re-injected onto ``running_queue``.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
|
||||
req = _req("late-abort", RequestStatus.WAITING_FOR_CHUNK)
|
||||
adapter.waiting_for_chunk_running_requests.append(req)
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
|
||||
# Tick N: process_pending_chunks runs while req is still tracked.
|
||||
scheduler_requests: dict = {req.request_id: req}
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert req in adapter.waiting_for_chunk_running_requests
|
||||
|
||||
# Mid-tick: abort fires and the scheduler's free path deletes the entry.
|
||||
del scheduler_requests[req.request_id]
|
||||
|
||||
# finally: restore_queues sees the now-untracked req and must drop it
|
||||
# instead of blindly extending it onto running_queue.
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests=scheduler_requests)
|
||||
assert running_queue == []
|
||||
assert req.request_id in adapter._cancelled_load_reqs
|
||||
|
||||
|
||||
def test_purge_is_noop_on_empty_deques(build_adapter):
|
||||
"""Empty deques short-circuit -- guards against any accidental
|
||||
``IndexError`` from ``popleft`` on an empty deque if a future caller
|
||||
shadows the empty check.
|
||||
"""
|
||||
adapter, _ = build_adapter(stage_id=1, max_num_seqs=8)
|
||||
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
|
||||
assert len(adapter.waiting_for_chunk_running_requests) == 0
|
||||
|
||||
waiting_queue = DummyWaitingQueue()
|
||||
running_queue: list = []
|
||||
|
||||
adapter.process_pending_chunks(waiting_queue, running_queue, scheduler_requests={})
|
||||
assert len(adapter.waiting_for_chunk_waiting_requests) == 0
|
||||
assert len(adapter.waiting_for_chunk_running_requests) == 0
|
||||
adapter.restore_queues(waiting_queue, running_queue, scheduler_requests={})
|
||||
assert running_queue == []
|
||||
assert waiting_queue == []
|
||||
@@ -13,7 +13,7 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
@pytest.fixture()
|
||||
def connector():
|
||||
c = SharedMemoryConnector({"shm_threshold_bytes": 64})
|
||||
c = SharedMemoryConnector({"shm_threshold_bytes": 0})
|
||||
yield c
|
||||
c.close()
|
||||
|
||||
|
||||
Loaded 100 of 271 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user