Compare commits
136
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
|
|
d98f94fa29 | ||
|
|
908e7d32d9 | ||
|
|
d6f1d3d61c | ||
|
|
c39ac9fffd | ||
|
|
580a0efe40 | ||
|
|
329c98ff9d | ||
|
|
36a5e9cfd2 | ||
|
|
0b04f2a0d6 | ||
|
|
5db89b41c4 | ||
|
|
e500c00545 | ||
|
|
8ddcb3621f | ||
|
|
2e6c83c89c | ||
|
|
d9063d7141 | ||
|
|
8c33793ffc | ||
|
|
4024a8ccd6 | ||
|
|
fb563fbd19 | ||
|
|
0fd2620d56 | ||
|
|
0d0f4e669e | ||
|
|
d94108693a | ||
|
|
d612233bef | ||
|
|
f6e8b77642 | ||
|
|
d8edf0c9d5 | ||
|
|
67bd02f767 | ||
|
|
0814d845dc | ||
|
|
bb973cb10c | ||
|
|
fb3abb38fb | ||
|
|
eb81914c64 | ||
|
|
1951dca786 | ||
|
|
d71cacc74c | ||
|
|
465f1464d2 | ||
|
|
f6894f4308 | ||
|
|
a2d1dafc65 | ||
|
|
6dc4638a7e | ||
|
|
d35cdd2dab | ||
|
|
0958e62ecc | ||
|
|
3901ad8165 | ||
|
|
b84c99c122 | ||
|
|
733f4f6a36 | ||
|
|
67d14b5ac4 | ||
|
|
64efa66869 | ||
|
|
6646a66d41 | ||
|
|
378359fc11 | ||
|
|
9c260b3774 | ||
|
|
9177b37e2f | ||
|
|
50bc51ffdf | ||
|
|
044240d8eb | ||
|
|
f4d922d8f6 | ||
|
|
eaa16e1941 | ||
|
|
8a6786731e | ||
|
|
fb0f17858f | ||
|
|
2447740308 | ||
|
|
a10eac0f8c | ||
|
+2 |
ec4ec0f604 | ||
|
|
ad07f98b73 | ||
|
|
a60242092a | ||
|
|
991b1ec752 | ||
|
|
6bdd85b384 | ||
|
|
36bd8029f3 | ||
|
|
f0d7fe18f2 | ||
|
|
ccb5c53ea5 | ||
|
|
8efd9d832d | ||
|
|
7ff11d8d0f | ||
|
|
fc32d9fe6f | ||
|
|
afbcc51672 | ||
|
|
1292c47804 | ||
|
|
e241005a71 | ||
|
|
950e66a44c | ||
|
|
a1e1ac0155 | ||
|
|
cead615288 | ||
|
|
17cf60a63d | ||
|
|
2d0b1d8a29 | ||
|
|
e70822ed54 | ||
|
|
6ff246bf3c | ||
|
|
951ae28eb0 | ||
|
|
7e152ca64b | ||
|
|
208932210f | ||
|
|
f05e03aea8 | ||
|
|
63ba11b8b1 | ||
|
|
37a442447e | ||
|
|
13e9094be0 | ||
|
|
1228a109f0 | ||
|
|
8f28b969e9 | ||
|
|
2d2258890b | ||
|
|
0383fde38a | ||
|
|
f75c437a60 | ||
|
|
dff8cda844 | ||
|
|
c782746133 | ||
|
|
db8bc2c71d | ||
|
|
14d70d9870 | ||
|
|
77d395150b | ||
|
|
f195e44bc1 | ||
|
|
2f64c712fd | ||
|
|
65bc737fd4 | ||
|
|
d744940b18 | ||
|
|
e957d16f8d | ||
|
|
8fa0394d29 | ||
|
|
ed37d8e7a2 | ||
|
|
d78a92ec54 | ||
|
|
679cbe226a | ||
|
|
dff757da9c | ||
|
|
890803ca06 | ||
|
|
64a5bb92de | ||
|
|
83e3b129f0 | ||
|
|
1761253005 | ||
|
|
d897852258 | ||
|
|
b46c5af06f | ||
|
|
f200e988a6 | ||
|
|
e3621ec633 | ||
|
|
e7f0db1067 | ||
|
|
864d22164c | ||
|
|
5f5b77aa9c | ||
|
|
3fa8817fc5 | ||
|
|
a62420450c | ||
|
|
5cce797c03 | ||
|
|
d8a72be206 | ||
|
|
ba4932c364 | ||
|
|
dc14380796 | ||
|
|
3bce5e2e3a | ||
|
|
46e0b71423 | ||
|
|
c58a02c873 | ||
|
|
a86c333c7e | ||
|
|
aff0af9754 | ||
|
|
b1a4b08a7e | ||
|
|
e2e79177db | ||
|
|
ae937320ff | ||
|
|
11829ccd75 | ||
|
|
a693ae67f6 | ||
|
|
3c62c53f59 | ||
|
|
f8340d078e | ||
|
|
c9d5545b52 | ||
|
|
10404e4c5b | ||
|
|
6c6fde3e6d | ||
|
|
7b837944a1 | ||
|
|
a9ee92ddec | ||
|
|
31886d6ece | ||
|
|
fa94f41807 |
No files matched your search
@@ -10,7 +10,7 @@ steps:
|
||||
DOCKER_BUILDKIT: "1"
|
||||
# Buildkite will automatically replace this with the actual commit hash
|
||||
VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}"
|
||||
VLLM_VERSION: "v0.22.0"
|
||||
VLLM_VERSION: "v0.23.0"
|
||||
priority: 100
|
||||
timeout_in_minutes: 60
|
||||
soft_fail: false
|
||||
+20
-9
@@ -1,11 +1,12 @@
|
||||
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only skip-ci
|
||||
# from git diff, then uploads document 2. When docs-only skip applies, image-build still runs for the nightly
|
||||
# exception in upload_pipeline.py (nightly-test label or main NIGHTLY=1); otherwise upload-nightly
|
||||
# can be skipped together with test-ready/test-merge. Weekly uploads also depend on image-build.
|
||||
# Document 1: Buildkite loads only this block on first parse. The next step resolves docs-only,
|
||||
# pytest skip-mark-only, or combined skip-ci from git diff, then uploads document 2. When skip-ci
|
||||
# applies, image-build still runs for the 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.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"
|
||||
- label: ":github: Resolve skip-ci (docs / skip marks) & upload pipeline"
|
||||
key: upload-ci-pipeline
|
||||
commands:
|
||||
- "python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/pipeline.yml"
|
||||
@@ -25,23 +26,33 @@ steps:
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L2 Test
|
||||
# L2 Test — PR ready label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
|
||||
- label: "Upload Ready Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
if: __UPLOAD_READY_IF__
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
- |
|
||||
EXTRA=""
|
||||
if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
|
||||
EXTRA="--e2e"
|
||||
fi
|
||||
python3 .buildkite/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
# L3 Test
|
||||
# L3 Test — PR merge-test label (diff-aware) or main+NIGHTLY=1 (--e2e, full E2E group)
|
||||
- label: "Upload Merge Pipeline"
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
if: __UPLOAD_MERGE_IF__
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
- |
|
||||
EXTRA=""
|
||||
if [ "${NIGHTLY:-}" = "1" ] && [ "${BUILDKITE_BRANCH}" = "main" ]; then
|
||||
EXTRA="--e2e"
|
||||
fi
|
||||
python3 .buildkite/scripts/upload_pipeline.py --upload $$EXTRA .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
|
||||
@@ -14,7 +14,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-ready-pipeline
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-ready.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-ready.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -23,7 +23,7 @@ steps:
|
||||
depends_on: image-build
|
||||
key: upload-merge-pipeline
|
||||
commands:
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload .buildkite/test-merge.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-merge.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
|
||||
@@ -33,6 +33,6 @@ steps:
|
||||
key: upload-nightly-pipeline
|
||||
if: build.env("NIGHTLY") == "1"
|
||||
commands:
|
||||
- buildkite-agent pipeline upload .buildkite/test-nightly.yml
|
||||
- python3 .buildkite/scripts/upload_pipeline.py --upload --all .buildkite/test-nightly.yml
|
||||
agents:
|
||||
queue: "cpu_queue_premerge"
|
||||
@@ -64,7 +64,8 @@ container_name="rocm_${BUILDKITE_COMMIT}_$(tr -dc A-Za-z0-9 < /dev/urandom | hea
|
||||
docker pull "${image_name}"
|
||||
|
||||
remove_docker_container() {
|
||||
docker rm -f "${container_name}" || docker image rm -f "${image_name}" || true
|
||||
docker ps -aq --filter "name=${container_name}" | xargs -r docker rm -f || true
|
||||
docker image rm -f "${image_name}" || true
|
||||
}
|
||||
trap remove_docker_container EXIT
|
||||
|
||||
@@ -107,7 +108,6 @@ if [[ $commands == *"--shard-id="* ]]; then
|
||||
--network=host \
|
||||
--shm-size=16gb \
|
||||
--group-add "$render_gid" \
|
||||
--rm \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
@@ -140,7 +140,6 @@ else
|
||||
--network=host \
|
||||
--shm-size=16gb \
|
||||
--group-add "$render_gid" \
|
||||
--rm \
|
||||
-e MIOPEN_DEBUG_CONV_DIRECT=0 \
|
||||
-e MIOPEN_DEBUG_CONV_GEMM=0 \
|
||||
-e VLLM_ROCM_USE_AITER=0 \
|
||||
|
||||
@@ -3,7 +3,8 @@
|
||||
"""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.
|
||||
- Detect docs-only, pytest skip-mark-only, or combined docs+skip-mark 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.
|
||||
@@ -12,19 +13,26 @@ Test pipeline mode (e.g. test-merge.yml):
|
||||
recognize this uploader-only key).
|
||||
|
||||
Usage:
|
||||
python3 upload_pipeline.py [--upload] <pipeline.yml>
|
||||
python3 upload_pipeline.py [--upload] [--all | --e2e] <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
|
||||
|
||||
# Full suite (rebase pipeline): keep all steps, still strip the uploader-only key:
|
||||
python3 upload_pipeline.py --upload --all .buildkite/test-merge.yml
|
||||
|
||||
# Nightly L2 E2E only: keep the E2E Test group, ignore source_file_dependencies:
|
||||
python3 upload_pipeline.py --upload --e2e .buildkite/test-ready.yml
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import re
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
@@ -42,6 +50,7 @@ LOG = "upload_pipeline"
|
||||
ROOT = Path(__file__).resolve().parent.parent.parent
|
||||
DOC_SEP = "\n---\n"
|
||||
BOOTSTRAP_MARKER = "__IMAGE_BUILD_IF__"
|
||||
E2E_GROUP_MARKER = "E2E Test"
|
||||
|
||||
|
||||
def _log(message: str) -> None:
|
||||
@@ -57,8 +66,8 @@ def _git(*args: str) -> subprocess.CompletedProcess[str]:
|
||||
)
|
||||
|
||||
|
||||
def resolve_changed_files() -> list[str] | None:
|
||||
"""Return changed file paths, or None when diff cannot be resolved."""
|
||||
def resolve_diff_range() -> str | None:
|
||||
"""Return a git diff range for PR or main builds, or None when unavailable."""
|
||||
is_pr = os.environ.get("BUILDKITE_PULL_REQUEST", "false") != "false" and os.environ.get(
|
||||
"BUILDKITE_PULL_REQUEST", ""
|
||||
)
|
||||
@@ -76,14 +85,22 @@ def resolve_changed_files() -> list[str] | None:
|
||||
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":
|
||||
return f"{base_ref}...{commit}"
|
||||
|
||||
if 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 f"{commit}^..{commit}"
|
||||
|
||||
_log("not PR/main build; using safe defaults")
|
||||
return None
|
||||
|
||||
|
||||
def resolve_changed_files() -> list[str] | None:
|
||||
"""Return changed file paths, or None when diff cannot be resolved."""
|
||||
diff_range = resolve_diff_range()
|
||||
if diff_range is None:
|
||||
return None
|
||||
|
||||
result = _git("diff", "--name-only", diff_range)
|
||||
@@ -96,30 +113,207 @@ def resolve_changed_files() -> list[str] | None:
|
||||
return files
|
||||
|
||||
|
||||
def _is_doc_file(file_path: str) -> bool:
|
||||
if not file_path:
|
||||
return False
|
||||
if file_path.startswith("docs/"):
|
||||
return True
|
||||
if file_path.endswith(".md"):
|
||||
return True
|
||||
return file_path == "mkdocs.yml"
|
||||
|
||||
|
||||
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/"):
|
||||
if not _is_doc_file(file_path):
|
||||
return False
|
||||
return has_any
|
||||
|
||||
|
||||
def _is_test_python_file(file_path: str) -> bool:
|
||||
path = Path(file_path)
|
||||
return path.suffix == ".py" and bool(path.parts) and path.parts[0] == "tests"
|
||||
|
||||
|
||||
_SKIP_MARK_RE = re.compile(r"pytest\.mark\.skip(?:if)?\b|pytest\.skip\s*\(")
|
||||
_PYTESTMARK_SKIP_RE = re.compile(r"pytest\.mark\.skip\b")
|
||||
_PYTEST_MARK_ONLY_RE = re.compile(r"pytest\.mark\.\w+")
|
||||
|
||||
|
||||
def _paren_balance(line: str) -> int:
|
||||
return line.count("(") - line.count(")")
|
||||
|
||||
|
||||
def _is_skip_mark_related_content(line: str) -> bool:
|
||||
stripped = line.strip()
|
||||
if not stripped:
|
||||
return True
|
||||
if stripped.startswith("#"):
|
||||
return True
|
||||
return _SKIP_MARK_RE.search(stripped) is not None
|
||||
|
||||
|
||||
def _is_pytestmark_adjacent_content(line: str) -> bool:
|
||||
"""Allow pytestmark list refactors that only add skip/skipif alongside existing marks."""
|
||||
stripped = line.strip().rstrip(",")
|
||||
if not stripped:
|
||||
return True
|
||||
if stripped in {"[", "]"}:
|
||||
return True
|
||||
if stripped.startswith("#"):
|
||||
return True
|
||||
if stripped.startswith("pytestmark"):
|
||||
return True
|
||||
return _PYTEST_MARK_ONLY_RE.search(stripped) is not None
|
||||
|
||||
|
||||
def diff_only_contains_skip_mark_changes(diff_text: str) -> bool:
|
||||
"""True when diff only edits skip marks or reformats pytestmark to add them."""
|
||||
pending_depth = 0
|
||||
saw_change = False
|
||||
has_skip_mark_edit = False
|
||||
for raw_line in diff_text.splitlines():
|
||||
if raw_line.startswith("@@"):
|
||||
pending_depth = 0
|
||||
continue
|
||||
if file_path.endswith(".md"):
|
||||
if not (raw_line.startswith("+") or raw_line.startswith("-")):
|
||||
continue
|
||||
if file_path == "mkdocs.yml":
|
||||
if raw_line.startswith("+++") or raw_line.startswith("---"):
|
||||
continue
|
||||
|
||||
saw_change = True
|
||||
content = raw_line[1:]
|
||||
|
||||
if pending_depth > 0:
|
||||
pending_depth += _paren_balance(content)
|
||||
if pending_depth < 0:
|
||||
pending_depth = 0
|
||||
continue
|
||||
|
||||
if _is_skip_mark_related_content(content):
|
||||
has_skip_mark_edit = True
|
||||
pending_depth = max(0, _paren_balance(content))
|
||||
continue
|
||||
|
||||
if _is_pytestmark_adjacent_content(content):
|
||||
continue
|
||||
|
||||
return False
|
||||
|
||||
return saw_change and has_skip_mark_edit
|
||||
|
||||
|
||||
def _git_file_diff(file_path: str, diff_range: str) -> str | None:
|
||||
result = _git("diff", diff_range, "--", file_path)
|
||||
if result.returncode != 0:
|
||||
_log(f"skip-mark-only: git diff failed for {file_path}")
|
||||
return None
|
||||
return result.stdout
|
||||
|
||||
|
||||
def _is_new_file_diff(diff_text: str) -> bool:
|
||||
return "new file mode" in diff_text or "--- /dev/null" in diff_text
|
||||
|
||||
|
||||
def _file_has_module_level_pytest_skip(content: str) -> bool:
|
||||
"""True when ``pytestmark`` applies unconditional ``pytest.mark.skip`` to the module."""
|
||||
match = re.search(r"^pytestmark\s*=\s*\[(.*?)^\s*\]", content, re.MULTILINE | re.DOTALL)
|
||||
if match is not None:
|
||||
return _PYTESTMARK_SKIP_RE.search(match.group(1)) is not None
|
||||
match = re.search(r"^pytestmark\s*=\s*pytest\.mark\.skip\b", content, re.MULTILINE)
|
||||
return match is not None
|
||||
|
||||
|
||||
def _file_content_at_commit(file_path: str, commit: str) -> str | None:
|
||||
result = _git("show", f"{commit}:{file_path}")
|
||||
if result.returncode != 0:
|
||||
_log(f"skip-mark-only: cannot read {file_path} at {commit}")
|
||||
return None
|
||||
return result.stdout
|
||||
|
||||
|
||||
def _qualifies_as_skip_mark_test_change(file_path: str, diff_range: str) -> bool:
|
||||
diff_text = _git_file_diff(file_path, diff_range)
|
||||
if diff_text is None:
|
||||
return False
|
||||
if not diff_text.strip():
|
||||
_log(f"skip-mark-only: empty diff for {file_path}")
|
||||
return False
|
||||
|
||||
if diff_only_contains_skip_mark_changes(diff_text):
|
||||
return True
|
||||
|
||||
if not _is_new_file_diff(diff_text):
|
||||
_log(f"skip-mark-only: non-skip changes in {file_path}")
|
||||
return False
|
||||
|
||||
commit = os.environ.get("BUILDKITE_COMMIT", "")
|
||||
if not commit:
|
||||
return False
|
||||
content = _file_content_at_commit(file_path, commit)
|
||||
if content is None:
|
||||
return False
|
||||
if _file_has_module_level_pytest_skip(content):
|
||||
_log(f"skip-mark-only: new test file with module-level pytest.mark.skip: {file_path}")
|
||||
return True
|
||||
|
||||
_log(f"skip-mark-only: new test file without module-level pytest.mark.skip: {file_path}")
|
||||
return False
|
||||
|
||||
|
||||
def is_skip_mark_only_change(changed_files: list[str], *, diff_range: str) -> bool:
|
||||
"""True when every changed file is a tests/ module with only skip-mark edits."""
|
||||
has_any = False
|
||||
for file_path in changed_files:
|
||||
if not file_path:
|
||||
continue
|
||||
has_any = True
|
||||
if not _is_test_python_file(file_path):
|
||||
return False
|
||||
if not _qualifies_as_skip_mark_test_change(file_path, diff_range):
|
||||
return False
|
||||
return has_any
|
||||
|
||||
|
||||
def is_docs_or_skip_mark_only_change(changed_files: list[str], *, diff_range: str) -> bool:
|
||||
"""True when each changed file is a doc path or a qualifying skip-mark test change."""
|
||||
has_any = False
|
||||
for file_path in changed_files:
|
||||
if not file_path:
|
||||
continue
|
||||
has_any = True
|
||||
if _is_doc_file(file_path):
|
||||
continue
|
||||
if _is_test_python_file(file_path) and _qualifies_as_skip_mark_test_change(file_path, diff_range):
|
||||
continue
|
||||
_log(f"docs/skip-mark-only: rejecting {file_path}")
|
||||
return False
|
||||
return has_any
|
||||
|
||||
|
||||
def resolve_skip_ci(changed_files: list[str] | None) -> bool:
|
||||
def resolve_skip_ci(changed_files: list[str] | None, *, diff_range: str | None = 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):
|
||||
|
||||
if diff_range is not None and is_docs_or_skip_mark_only_change(changed_files, diff_range=diff_range):
|
||||
if is_docs_only_change(changed_files):
|
||||
_log("docs-only change detected; skip-ci=1")
|
||||
elif is_skip_mark_only_change(changed_files, diff_range=diff_range):
|
||||
_log("pytest skip-mark-only change detected; skip-ci=1")
|
||||
else:
|
||||
_log("docs + pytest skip-mark-only change detected; skip-ci=1")
|
||||
return True
|
||||
|
||||
if diff_range is None and is_docs_only_change(changed_files):
|
||||
_log("docs-only change detected; skip-ci=1")
|
||||
return True
|
||||
_log("non-doc changes detected; skip-ci=0")
|
||||
|
||||
_log("non-doc/non-skip-mark changes detected; skip-ci=0")
|
||||
return False
|
||||
|
||||
|
||||
@@ -132,18 +326,18 @@ def render_bootstrap_pipeline(text: str, *, skip_ci: bool) -> str:
|
||||
nightly_only = (
|
||||
'(build.pull_request.labels includes "nightly-test") || (build.branch == "main" && build.env("NIGHTLY") == "1")'
|
||||
)
|
||||
nightly_main = 'build.branch == "main" && build.env("NIGHTLY") == "1"'
|
||||
ready_pr = 'build.branch != "main" && build.pull_request.labels includes "ready"'
|
||||
merge_main = 'build.branch == "main" && build.env("NIGHTLY") != "1" && build.env("WEEKLY") != "1"'
|
||||
merge_pr = 'build.branch != "main" && build.pull_request.labels includes "merge-test"'
|
||||
if skip_ci:
|
||||
image_if = f"'{nightly_only}'"
|
||||
ready_if = "'false'"
|
||||
merge_if = "'false'"
|
||||
ready_if = f"'({nightly_main})'"
|
||||
merge_if = f"'({nightly_main})'"
|
||||
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")\''
|
||||
)
|
||||
ready_if = f"'(({nightly_main}) || ({ready_pr}))'"
|
||||
merge_if = f"'(({nightly_main}) || (({merge_main}) || ({merge_pr})))'"
|
||||
|
||||
return (
|
||||
continuation.replace("__IMAGE_BUILD_IF__", image_if)
|
||||
@@ -206,6 +400,20 @@ def _filter_steps(steps: list[Any], changed_files: list[str]) -> list[Any]:
|
||||
return filtered
|
||||
|
||||
|
||||
def _is_e2e_group(step: dict[str, Any]) -> bool:
|
||||
group = step.get("group")
|
||||
return isinstance(group, str) and E2E_GROUP_MARKER in group
|
||||
|
||||
|
||||
def _select_e2e_group_steps(steps: list[Any]) -> list[Any]:
|
||||
selected = [step for step in steps if isinstance(step, dict) and _is_e2e_group(step)]
|
||||
if not selected:
|
||||
_log(f"no group matching {E2E_GROUP_MARKER!r} found")
|
||||
else:
|
||||
_log(f"keep {len(selected)} group(s) matching {E2E_GROUP_MARKER!r}")
|
||||
return selected
|
||||
|
||||
|
||||
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] = []
|
||||
@@ -232,11 +440,16 @@ def _strip_uploader_metadata_from_steps(steps: list[Any]) -> list[Any]:
|
||||
def render_test_pipeline(
|
||||
doc: dict[str, Any],
|
||||
changed_files: list[str] | None,
|
||||
*,
|
||||
e2e_only: bool = False,
|
||||
) -> dict[str, Any]:
|
||||
steps = doc.get("steps")
|
||||
if not isinstance(steps, list):
|
||||
return doc
|
||||
if changed_files is not None:
|
||||
if e2e_only:
|
||||
steps = _select_e2e_group_steps(steps)
|
||||
steps = _strip_uploader_metadata_from_steps(steps)
|
||||
elif changed_files is not None:
|
||||
steps = _filter_steps(steps, changed_files)
|
||||
else:
|
||||
steps = _strip_uploader_metadata_from_steps(steps)
|
||||
@@ -250,19 +463,34 @@ def resolve_pipeline_path(arg: str) -> Path:
|
||||
return ROOT / path
|
||||
|
||||
|
||||
def render_pipeline(path: Path) -> str:
|
||||
def render_pipeline(
|
||||
path: Path,
|
||||
*,
|
||||
force_all: bool = False,
|
||||
e2e_only: bool = False,
|
||||
) -> str:
|
||||
text = path.read_text(encoding="utf-8")
|
||||
diff_range = resolve_diff_range()
|
||||
changed_files = resolve_changed_files()
|
||||
|
||||
# ``--all`` forces the keep-all-steps path (no diff-aware skipping) while still
|
||||
# stripping ``source_file_dependencies``. Used by the rebase pipeline so main builds
|
||||
# run the full e2e suite (see .buildkite/rebase-pipeline.yaml).
|
||||
if force_all or e2e_only:
|
||||
changed_files = None
|
||||
|
||||
if BOOTSTRAP_MARKER in text:
|
||||
rendered = render_bootstrap_pipeline(text, skip_ci=resolve_skip_ci(changed_files))
|
||||
rendered = render_bootstrap_pipeline(
|
||||
text,
|
||||
skip_ci=resolve_skip_ci(changed_files, diff_range=diff_range),
|
||||
)
|
||||
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)
|
||||
doc = render_test_pipeline(doc, changed_files, e2e_only=e2e_only)
|
||||
|
||||
return yaml.safe_dump(doc, sort_keys=False)
|
||||
|
||||
@@ -289,6 +517,17 @@ def main() -> int:
|
||||
action="store_true",
|
||||
help="Pipe rendered YAML to buildkite-agent pipeline upload",
|
||||
)
|
||||
mode = parser.add_mutually_exclusive_group()
|
||||
mode.add_argument(
|
||||
"--all",
|
||||
action="store_true",
|
||||
help="Keep all steps (disable diff-aware skipping); still strips source_file_dependencies",
|
||||
)
|
||||
mode.add_argument(
|
||||
"--e2e",
|
||||
action="store_true",
|
||||
help="Keep only the E2E Test group; disable diff-aware skipping for those steps",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
path = resolve_pipeline_path(args.pipeline)
|
||||
@@ -296,7 +535,7 @@ def main() -> int:
|
||||
_log(f"missing pipeline file: {path}")
|
||||
return 1
|
||||
|
||||
rendered = render_pipeline(path)
|
||||
rendered = render_pipeline(path, force_all=args.all, e2e_only=args.e2e)
|
||||
if args.upload:
|
||||
upload_to_buildkite(rendered)
|
||||
else:
|
||||
|
||||
@@ -179,17 +179,30 @@ steps:
|
||||
pytest -s -v tests/e2e/online_serving/test_qwen3_tts_base.py tests/e2e/offline_inference/test_qwen3_tts_base.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
# It is the whisper ASR validation step that is slow, not the Voxtral TTS generation.
|
||||
# whisper ASR runs on CPU.
|
||||
- label: "Voxtral-TTS E2E Test (Online Serving)"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
timeout 40m bash -c '
|
||||
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"
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "Voxtral-TTS E2E Test (Offline Inference)"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c '
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
# TODO: Bagel test on ROCm is very unstable. @tjtanaa
|
||||
|
||||
@@ -185,20 +185,6 @@ steps:
|
||||
- export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1"
|
||||
- timeout 30m pytest -s -v tests/e2e/online_serving/test_qwen3_tts_customvoice.py -m "core_model" --run-level "core_model"
|
||||
|
||||
- label: "Voxtral-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
mirror_hardwares: [amdproduction]
|
||||
grade: Blocking
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c '
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/online_serving/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
pytest -s -v tests/e2e/offline_inference/test_voxtral_tts.py -m "advanced_model" --run-level "advanced_model"
|
||||
'
|
||||
|
||||
- label: "CosyVoice3-TTS E2E Test"
|
||||
agent_pool: mi325_1
|
||||
depends_on: amd-build
|
||||
|
||||
@@ -363,6 +363,42 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "TTS · IndexTTS2 Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_indextts2.py
|
||||
- tests/e2e/offline_inference/test_indextts2.py
|
||||
- tests/model_executor/models/indextts2/
|
||||
- tests/model_executor/stage_input_processors/test_indextts2.py
|
||||
- tests/assets/indextts2/
|
||||
- vllm_omni/model_executor/models/indextts2/
|
||||
- vllm_omni/model_executor/stage_input_processors/indextts2.py
|
||||
- vllm_omni/deploy/indextts2.yaml
|
||||
- vllm_omni/deploy/indextts2_low_latency.yaml
|
||||
- vllm_omni/entrypoints/openai/serving_speech.py
|
||||
- vllm_omni/entrypoints/openai/tts_adapters/indextts2.py
|
||||
- vllm_omni/config/pipeline_registry.py
|
||||
- vllm_omni/model_executor/models/registry.py
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
timeout 20m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
pytest -s -v tests/e2e/online_serving/test_indextts2.py tests/e2e/offline_inference/test_indextts2.py -m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
"
|
||||
agents:
|
||||
queue: "gpu_4_queue"
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
- "HF_TOKEN"
|
||||
gpus: all
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "Omni · Qwen3-Omni Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_qwen3_omni.py
|
||||
@@ -509,6 +545,34 @@ steps:
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
|
||||
- label: "TTS · SoulX-Singer Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/offline_inference/test_soulxsinger.py
|
||||
- vllm_omni/diffusion/models/soulx_singer/
|
||||
- vllm_omni/deploy/soulxsinger_svc.yaml
|
||||
- vllm_omni/deploy/soulxsinger_svs.yaml
|
||||
timeout_in_minutes: 30
|
||||
commands:
|
||||
- |
|
||||
timeout 30m bash -c "
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
export VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
pytest -s -v tests/e2e/offline_inference/test_soulxsinger.py -k 'svc or precomputed' -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: "Diffusion · Qwen Image Edit Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image_edit.py
|
||||
|
||||
@@ -1279,7 +1279,7 @@ steps:
|
||||
- label: ":full_moon: Diffusion X2V · Function Test"
|
||||
timeout_in_minutes: 90
|
||||
commands:
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
- pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py tests/e2e/offline_inference/test_wan22_autoround_w4a16_expansion.py tests/e2e/online_serving/test_ltx2_expansion.py -m "full_model and cuda" --run-level "full_model"
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
@@ -1322,6 +1322,7 @@ steps:
|
||||
- pytest -s -v tests/e2e/accuracy/hunyuanvideo15_i2v/test_hunyuanvideo15_i2v_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_ltx2_3_video_similarity.py -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_diffusers_backend_similarity.py -k '2v_matches_diffusers' -m full_model --run-level full_model
|
||||
- pytest -s -v tests/e2e/accuracy/test_video_streaming_output_similarity.py -m full_model --run-level full_model
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
|
||||
@@ -61,7 +61,7 @@ steps:
|
||||
steps:
|
||||
- label: "Diffusion · Offloader Test"
|
||||
commands:
|
||||
- timeout 10m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
|
||||
- timeout 40m pytest -s -v tests/diffusion/offloader -m "core_model and cuda" --run-level "core_model"
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
|
||||
@@ -85,6 +85,66 @@ steps:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Reliability Test - hunyuan-image"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
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: 4
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
|
||||
- label: "Reliability Test - voxcpm2"
|
||||
timeout_in_minutes: 180
|
||||
depends_on: upload-weekly-pipeline
|
||||
if: build.env("WEEKLY") == "1" || build.pull_request.labels includes "weekly-test"
|
||||
commands:
|
||||
- pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
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"
|
||||
|
||||
# Invalid HTTP / WS payloads vs live Omni servers (Issue #3649 tracks known skips).
|
||||
- group: "Reliability Test - Invalid parameters Test"
|
||||
depends_on: upload-weekly-pipeline
|
||||
|
||||
@@ -24,7 +24,7 @@ Check the model's HF repo for `model_index.json`. This determines your path:
|
||||
|
||||
| Scenario | How to identify | Migration path |
|
||||
|----------|----------------|----------------|
|
||||
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip to Step 5 (test) and Step 7 (docs) |
|
||||
| **Already supported** | `_class_name` in `model_index.json` matches a key in `_DIFFUSION_MODELS` in `registry.py` | Skip implementation, then validate model-specific examples, tests, and docs as needed |
|
||||
| **Diffusers-based** | Has standard `model_index.json` with `_diffusers_version`, subfolders for `transformer/`, `vae/`, etc. | Follow **Path A** below |
|
||||
| **Custom/private repo** | No diffusers `model_index.json`, weights in non-standard format, custom model code in a separate git repo | Follow **Path B** below |
|
||||
| **Hybrid** | Has some diffusers components (VAE) but custom transformer/fusion | Mix of Path A and Path B |
|
||||
@@ -52,7 +52,7 @@ vllm_omni/diffusion/models/your_model_name/
|
||||
2. Replace attention with `vllm_omni.diffusion.attention.layer.Attention` (QKV shape: `[B, seq, heads, head_dim]`).
|
||||
3. Add `od_config: OmniDiffusionConfig | None = None` to `__init__`.
|
||||
4. Add `load_weights()` method mapping diffusers weight names to vllm-omni names.
|
||||
5. Add class attributes: `_repeated_blocks`, `_layerwise_offload_blocks_attr`.
|
||||
5. Add class attributes for acceleration features such as `_repeated_blocks` and `_layerwise_offload_blocks_attrs` (see [references/transformer-adaptation.md](references/transformer-adaptation.md) for examples).
|
||||
|
||||
### A4. Adapt pipeline
|
||||
|
||||
@@ -78,13 +78,34 @@ class YourPipeline(nn.Module):
|
||||
|
||||
Add post/pre-process functions in the same pipeline file. Register them in `registry.py`.
|
||||
|
||||
### A4.1 Add progress bar support (recommended)
|
||||
|
||||
For pipelines with a standard denoising loop, prefer the existing progress bar pattern instead of hand-rolled logging.
|
||||
|
||||
```python
|
||||
from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin
|
||||
|
||||
class YourPipeline(nn.Module, ProgressBarMixin):
|
||||
def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput:
|
||||
# ... prepare timesteps / latents ...
|
||||
with self.progress_bar(total=len(timesteps)) as progress_bar:
|
||||
for i, t in enumerate(timesteps):
|
||||
# predict noise / scheduler step
|
||||
latents = ...
|
||||
progress_bar.update()
|
||||
|
||||
return DiffusionOutput(output=output)
|
||||
```
|
||||
|
||||
For custom loop structures, follow `vllm_omni/diffusion/models/progress_bar.py` and existing pipelines using `ProgressBarMixin`.
|
||||
|
||||
### A5. Register, test, docs → continue at Step 4 below.
|
||||
|
||||
---
|
||||
|
||||
## Path B: Custom/Private Repo Model
|
||||
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3.
|
||||
For models without a diffusers pipeline — weights in custom format, model code in a private repo. Real examples: DreamID-Omni, BAGEL, HunyuanImage3. See [references/custom-model-patterns.md](references/custom-model-patterns.md) for concrete custom-model integration patterns.
|
||||
|
||||
### B1. Understand the reference repo
|
||||
|
||||
@@ -134,7 +155,7 @@ Copy the essential model files into `vllm_omni/diffusion/models/<name>/` and ada
|
||||
|
||||
### B4. Handle custom weight loading
|
||||
|
||||
Custom models have two patterns for weight loading:
|
||||
Custom models have two common patterns for weight loading:
|
||||
|
||||
**Pattern 1: Bypass standard loader** (DreamID-Omni style)
|
||||
|
||||
@@ -149,8 +170,10 @@ class CustomPipeline(nn.Module):
|
||||
self.vae = custom_init_vae(model, device=self.device)
|
||||
self.text_encoder = custom_init_text_encoder(model, device=self.device)
|
||||
self.transformer = CustomFusionModel(CONFIG)
|
||||
load_custom_checkpoint(self.transformer,
|
||||
checkpoint_path=os.path.join(model, "model.safetensors"))
|
||||
load_custom_checkpoint(
|
||||
self.transformer,
|
||||
checkpoint_path=os.path.join(model, "model.safetensors"),
|
||||
)
|
||||
# NO weights_sources defined — bypasses standard loader
|
||||
|
||||
def load_weights(self, weights):
|
||||
@@ -162,6 +185,8 @@ class CustomPipeline(nn.Module):
|
||||
When weights are in safetensors format but need name remapping:
|
||||
|
||||
```python
|
||||
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
|
||||
|
||||
class CustomPipeline(nn.Module):
|
||||
def __init__(self, *, od_config, prefix=""):
|
||||
super().__init__()
|
||||
@@ -187,7 +212,9 @@ class CustomPipeline(nn.Module):
|
||||
# Remap original weight names to vllm-omni module names
|
||||
name = self._remap_weight_name(name)
|
||||
if name in params:
|
||||
default_weight_loader(params[name], tensor)
|
||||
param = params[name]
|
||||
weight_loader = getattr(param, "weight_loader", default_weight_loader)
|
||||
weight_loader(param, tensor)
|
||||
loaded.add(name)
|
||||
return loaded
|
||||
```
|
||||
@@ -281,16 +308,43 @@ For Omni or custom models, create:
|
||||
### Step 7: Update Documentation
|
||||
|
||||
Required updates:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — parallelism support table
|
||||
1. `docs/user_guide/diffusion/parallelism/overview.md` — parallelism support overview/table
|
||||
2. `docs/user_guide/diffusion/cpu_offload_diffusion.md` — if CPU offload supported (add to supported models table)
|
||||
3. `docs/user_guide/diffusion/teacache.md` — if TeaCache supported
|
||||
4. `docs/user_guide/diffusion/cache_dit_acceleration.md` — if Cache-DiT supported
|
||||
5. `examples/offline_inference/xxx/README.md` — offline example docs
|
||||
6. `examples/online_serve/xxx/README.md` — online serve docs
|
||||
3. `docs/user_guide/diffusion/cache_acceleration/teacache.md` — if TeaCache supported
|
||||
4. `docs/user_guide/diffusion/cache_acceleration/cache_dit.md` — if Cache-DiT supported
|
||||
5. Offline example docs under `examples/offline_inference/<name>/` (`README.md` or category-specific `.md`)
|
||||
6. `examples/online_serving/<name>/README.md` — online serving docs
|
||||
|
||||
### Step 8: Add E2E Tests (Recommended)
|
||||
### Step 8: Add E2E Tests
|
||||
|
||||
Create `tests/e2e/online_serving/test_your_model_expansion.py`.
|
||||
**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming, Buildkite wiring, and run commands. Also read [l4_functionality_tests.inc.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/test_examples/l4_functionality_tests.inc.md) and [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md).
|
||||
|
||||
Classify the model's **CI priority** first:
|
||||
|
||||
| Priority | Required test levels | Files & markers |
|
||||
|----------|---------------------|-----------------|
|
||||
| **High** (listed in [#1832](https://github.com/vllm-project/vllm-omni/issues/1832) or on the diffusion hot path) | **L1** · **L2** online · **L3** online + offline · **L4** feature + performance | See table below |
|
||||
| **Medium** (*normal priority* in L4 docs) | **L3** online + offline · **L4** feature only | Fewer L4 parametrized rows |
|
||||
| **Low** | **L4** feature only | One or two `*_expansion.py` cases |
|
||||
|
||||
**Per-level deliverables (diffusion / `pytest.mark.diffusion`):**
|
||||
|
||||
| Level | Location | Marker | CI pipeline | Notes |
|
||||
|-------|----------|--------|-------------|-------|
|
||||
| **L1** | `tests/diffusion/models/{slug}/`, `tests/diffusion/cache/`, transformer unit tests | `core_model` + `cpu` | `test-ready.yml` | Weight remap, `_sp_plan`, cache enabler registration, shape contracts |
|
||||
| **L2** | `tests/e2e/online_serving/test_{slug}.py` (and offline if the category is offline-first) | **`core_model` + `advanced_model`** (both on baseline smoke) + `diffusion` + `@hardware_test` / `hardware_marks` | `test-ready.yml` | Default deploy smoke — minimal `num_inference_steps`, single prompt |
|
||||
| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` when offline matters | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `diffusion`) | `test-merge.yml` **or** merged into nightly diffusion function job | Real weights, streaming/API paths, LoRA/offload smoke |
|
||||
| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py` (+ offline expansion if needed) | `full_model` + `diffusion` | `test-nightly.yml` (X2I / X2V / X2A function groups) | Feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); perf → `tests/dfx/perf/tests/test_{model}_vllm_omni.json` |
|
||||
|
||||
**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, minimal steps, single prompt) should carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** function so L2 (`test-ready.yml`) and L3 (`test-merge.yml`) share one smoke test. Heavier deploy variants or API paths in the same file use **`advanced_model` only**. When L3 moves to nightly, migrate those heavier cases into `test_{slug}_expansion.py` with `full_model` and remove the dedicated `test-merge.yml` job (see `test_longcat_image_expansion.py`, `test_qwen_image_expansion.py`).
|
||||
|
||||
**L4 design (high priority):** Combine multiple supported features (Cache-DiT, TP, USP, CFG, HSDP, CPU offload, quantization) into **few parametrized** `OmniServerParams` rows so each feature appears in at least one case without exploding GPU jobs. Shard single-GPU vs multi-GPU cases across the nightly X2I/X2V function steps (1-GPU vs `distributed_cuda`).
|
||||
|
||||
**L4 design (medium / low):** One or two parametrized rows covering the best quality/perf trade-off; skip perf JSON unless the model is high priority.
|
||||
|
||||
**Reference implementations:** `tests/e2e/online_serving/test_qwen_image_edit_expansion.py`, `tests/e2e/online_serving/test_longcat_image_expansion.py`, `tests/e2e/online_serving/test_hunyuan_video_15_expansion.py`.
|
||||
|
||||
**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:** deploy constants, prompts, sampling dicts, and inline `request_config` / `form_data` belong in each `test_{slug}.py` and `test_{slug}_expansion.py`. Do not add per-model files under `tests/helpers/`; reuse only repo-wide harness (`mark`, `media`, `runtime`, `stage_config`, `assertions`). **L2+ online/offline e2e:** reuse or add `send_*_request` in `tests/helpers/runtime.py` — tests call the handler, not raw `omni.generate` / HTTP. See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers**.
|
||||
|
||||
### Step 9: Add Cache-DiT Acceleration
|
||||
|
||||
@@ -510,7 +564,7 @@ class YourTransformer(nn.Module):
|
||||
#### 10e. Update parallelism documentation
|
||||
|
||||
After adding parallelism support, update:
|
||||
1. `docs/user_guide/diffusion/parallelism_acceleration.md` — add your model to the support table
|
||||
1. `docs/user_guide/diffusion/parallelism/overview.md` — add your model to the support overview/table
|
||||
2. Record which parallelism methods are supported (USP, Ring, CFG, TP, HSDP, VAE-Patch)
|
||||
|
||||
### Step 11: Add CPU Offload Support
|
||||
@@ -568,6 +622,7 @@ See the [Profiling Single-Stage Diffusion](../../../docs/contributing/profiling.
|
||||
|
||||
## Reference Files
|
||||
|
||||
- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands
|
||||
- [Transformer Adaptation](references/transformer-adaptation.md) — porting transformers from diffusers
|
||||
- [Custom Model Patterns](references/custom-model-patterns.md) — patterns for non-diffusers models
|
||||
- [Parallelism Patterns](references/parallelism-patterns.md) — TP, SP/USP, CFG parallel, HSDP implementation details
|
||||
|
||||
@@ -569,4 +569,4 @@ Check which parallelism methods a model supports:
|
||||
| **TP** | Uses `ColumnParallelLinear` / `QKVParallelLinear`. Search: `grep -r 'ParallelLinear\|QKVParallel' vllm_omni/diffusion/models/<model>/` |
|
||||
| **HSDP** | Transformer defines `_hsdp_shard_conditions`. Search: `grep -r '_hsdp_shard_conditions' vllm_omni/diffusion/models/` |
|
||||
|
||||
The canonical per-model support table is in `docs/user_guide/diffusion/parallelism_acceleration.md`.
|
||||
The canonical per-model support overview/table is in `docs/user_guide/diffusion/parallelism/overview.md`.
|
||||
@@ -219,13 +219,14 @@ def load_weights(self, weights):
|
||||
|
||||
### Class-Level Attributes for Features
|
||||
|
||||
|
||||
```python
|
||||
class MyTransformer(nn.Module):
|
||||
# torch.compile: list block class names that repeat and can be compiled
|
||||
_repeated_blocks = ["MyTransformerBlock"]
|
||||
|
||||
# CPU offload: attribute name of the nn.ModuleList containing blocks
|
||||
_layerwise_offload_blocks_attr = "blocks"
|
||||
_layerwise_offload_blocks_attrs = ["blocks"]
|
||||
|
||||
# LoRA: mapping of fused param names to original param names
|
||||
packed_modules_mapping = {"to_qkv": ["to_q", "to_k", "to_v"]}
|
||||
|
||||
@@ -305,13 +305,72 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
|
||||
]
|
||||
```
|
||||
|
||||
### Test Case Writing (CI Levels)
|
||||
|
||||
**Follow the [vllm-omni-test skill](../vllm-omni-test/SKILL.md)** for markers, file naming (`test_{slug}.py` / `test_{slug}_expansion.py`), Buildkite wiring, and copy-paste run commands. Also read [CI_5levels.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/CI_5levels.md) and [tests_style.md](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/ci/tests_style.md).
|
||||
|
||||
Classify the model's **CI priority** first (high / medium / low). High-priority TTS models are typically those on the integration hot path or listed in tracking issues such as [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); medium and low tiers cover the long tail. When unsure, ask the reviewer which tier applies.
|
||||
|
||||
| Priority | Required test levels | Files & markers |
|
||||
|----------|---------------------|-----------------|
|
||||
| **High** | **L1** unit/logic · **L2** online smoke · **L3** online + offline integration · **L4** feature + performance | See table below |
|
||||
| **Medium** | **L3** online + offline · **L4** feature only | Skip dedicated L1/L2 unless fixing a logic bug |
|
||||
| **Low** | **L4** feature only | One or two `*_expansion.py` parametrized cases |
|
||||
|
||||
**Per-level deliverables (TTS / `pytest.mark.tts`):**
|
||||
|
||||
| Level | Location | Marker | CI pipeline | Notes |
|
||||
|-------|----------|--------|-------------|-------|
|
||||
| **L1** | `tests/model_executor/…`, `tests/entrypoints/openai_api/…`, stage-processor tests | `core_model` + `cpu` | `test-ready.yml` | Prompt assembly, async_chunk helpers, `serving_speech` validation — no GPU |
|
||||
| **L2** | `tests/e2e/online_serving/test_{slug}.py` | **`core_model` + `advanced_model`** (both on baseline smoke) + `tts` + `@hardware_test(...)` | `test-ready.yml` (`ready` label) | Default deploy smoke: single `/v1/audio/speech` or offline `OmniRunner` path |
|
||||
| **L3** | `tests/e2e/online_serving/test_{slug}.py` **and** `tests/e2e/offline_inference/test_{slug}.py` | Baseline smoke: **`core_model` + `advanced_model`**; heavier cases: `advanced_model` only (+ `tts`) | `test-merge.yml` **or** merged into nightly TTS function job | Streaming, voice clone, batch/queue, async_chunk |
|
||||
| **L4** | `tests/e2e/online_serving/test_{slug}_expansion.py`, optional offline expansion | `full_model` + `tts` | `test-nightly.yml` (`:full_moon: TTS · Function Test with L4`) | Feature matrix; perf → `tests/dfx/perf/tests/test_tts.json` |
|
||||
|
||||
**L2 & L3 online — same file, dual marks on the baseline smoke:** The **first / simplest** case in `test_{slug}.py` (default deploy, single non-streaming `/v1/audio/speech` or equivalent offline path) should carry **both** `@pytest.mark.core_model` **and** `@pytest.mark.advanced_model` on the **same** function so it runs in L2 (`test-ready.yml`, `--run-level core_model`, basic validation) and L3 (`test-merge.yml`, `--run-level advanced_model`, deeper validation) without duplicating the test. In-tree examples: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`.
|
||||
|
||||
Heavier scenarios in the same file use **`advanced_model` only** (streaming, extra languages, concurrency, async_chunk, batch). Example: `test_voice_clone_en_streaming_001` → `advanced_model` only. When migrating L3 to nightly, move those heavier cases into `test_{slug}_expansion.py` with `full_model` and drop the dedicated merge job (see `test_ming_tts_expansion.py`, `test_glm_tts_expansion.py`).
|
||||
|
||||
```python
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(res={"cuda": "L4"}, num_cards=1)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_voice_clone_en_non_streaming_001(omni_server, openai_client) -> None:
|
||||
openai_client.send_audio_speech_request({...})
|
||||
```
|
||||
|
||||
**L4 consolidation:** Prefer parametrized `OmniServerParams` rows (`default`, `async_chunk`, feature flags) in one expansion module rather than many merge-only files ([#1832](https://github.com/vllm-project/vllm-omni/issues/1832)).
|
||||
|
||||
**L4 performance (high-priority models):** Add latency / throughput / stress rows in `tests/dfx/perf/tests/test_tts.json`, or a dedicated `tests/dfx/perf/tests/test_{slug}.json` when the model must not join the shared nightly server matrix before integration lands (see VoxCPM2 / Coqui XTTS pattern). Register the model in `benchmarks/tts/model_configs.yaml` for local `bench_tts.py`. Wire a **separate** `test-nightly.yml` Perf Test step when the JSON is not merged into `test_tts.json` yet.
|
||||
|
||||
**Keep model-specific code inside test modules — not `tests/helpers/{slug}.py`:**
|
||||
|
||||
- Put `MODEL`, deploy path, vendored `REF_AUDIO_URL`, `get_prompt()`, and inline `request_config` dicts **in each** `test_{slug}.py`, `test_{slug}_expansion.py`, offline `test_{slug}.py`, and L1 `test_{slug}_*.py` as needed.
|
||||
- **Do not** add `tests/helpers/{slug}.py` (or `tests/helpers/{model_name}.py`) to deduplicate constants or request builders across those files. A little duplication is intentional; follow in-tree references such as `tests/e2e/online_serving/test_glm_tts.py` and `tests/e2e/online_serving/test_cosyvoice3_tts_expansion.py`.
|
||||
- `tests/helpers/` is for **repo-wide** harness code only (`mark.py`, `media.py`, `runtime.py`, `stage_config.py`, `assertions.py`, `fixtures/`). Import those; do not extend the tree with per-model modules.
|
||||
|
||||
**Runtime send helpers (`tests/helpers/runtime.py`) — online and offline e2e:**
|
||||
|
||||
| Path | Fixture | Call |
|
||||
|------|---------|------|
|
||||
| Online `/v1/*` | `openai_client` | `openai_client.send_*_request(request_config)` |
|
||||
| Offline inference | `omni_runner_handler` | `omni_runner_handler.send_*_request(request_config)` |
|
||||
|
||||
1. **Grep `runtime.py` first** — reuse `send_omni_request`, `send_diffusion_request`, `send_audio_speech_request` (online + offline Qwen-style TTS), `send_single_stage_tts_request` (Coqui XTTS / MOSS-TTS-Nano offline), etc.
|
||||
2. **No matching helper** → add `send_<feature>_request` (or `send_<route>_http_request` for negative/dfx) in **`runtime.py`** with general `assert_*` bundled inside, **then** call it from the test.
|
||||
3. **Test file** holds `request_config` dicts only — not `omni.generate`, not `_collect_audio()`, not raw HTTP/SDK.
|
||||
4. Wire `tests/helpers/runtime.py` into Buildkite `source_file_dependencies` when you add helpers.
|
||||
|
||||
See [vllm-omni-test skill](../vllm-omni-test/SKILL.md) § **Runtime send helpers** for full tables and exceptions.
|
||||
|
||||
### Deliverables
|
||||
|
||||
- Updated `serving_speech.py` with all 5 integration points (single commit)
|
||||
- Client scripts and server launcher under `examples/online_serving/text_to_speech/<model>/`
|
||||
- Gradio demo with streaming and voice cloning UI in the same dir
|
||||
- E2E online serving test (`tests/e2e/online_serving/test_<model>.py`)
|
||||
- Buildkite CI entry in `.buildkite/test-merge.yml`
|
||||
- E2E tests per **Test Case Writing (CI Levels)** above (priority tier determines L1–L4 scope)
|
||||
- Buildkite wired per level: `test-ready.yml` (L1/L2), `test-merge.yml` or nightly function job (L3), `test-nightly.yml` (L4) — see [vllm-omni-test skill](../vllm-omni-test/SKILL.md)
|
||||
- New section in `examples/online_serving/text_to_speech/README.md` (table row + per-model section). Do **not** create a top-level `examples/online_serving/<model>/` dir or a per-model `README.md` inside `text_to_speech/<model>/`.
|
||||
|
||||
### E2E test pitfalls to avoid
|
||||
@@ -326,9 +385,9 @@ def build_voice_clone_prompt(ref_audio_path: str, text: str, codec) -> list:
|
||||
- **Use the harness readiness gate.** The fixture waits for HTTP 200 on
|
||||
`/health`; don't add `time.sleep` in tests. If warmup is incomplete, make
|
||||
`/health` return non-200 until you're actually ready.
|
||||
- **Mark with `@pytest.mark.core_model` + `hardware_test(res={"cuda": "H100"})`**
|
||||
so the test lands in `test-ready.yml` (triggered by the `ready` label) rather
|
||||
than only nightly.
|
||||
- **Mark tests per the CI Levels table** — baseline smoke: `core_model` + `advanced_model`; heavier cases: `advanced_model` only; L4 expansion: `full_model`
|
||||
- **No per-model helper modules** — do not create `tests/helpers/{slug}.py`; keep constants and `request_config` payloads in the test file
|
||||
- **Online and offline e2e go through `runtime.py`** — `openai_client.send_audio_speech_request` (online); `omni_runner_handler.send_audio_speech_request` (Qwen-style offline) or `send_single_stage_tts_request` (single-stage offline). Add a new `send_*_request` in `runtime.py` when none fits; do not embed `omni.generate` or HTTP in tests
|
||||
|
||||
## Phase 4: Async Chunk (Streaming)
|
||||
|
||||
@@ -461,8 +520,8 @@ Use this checklist when integrating a new TTS model:
|
||||
- [ ] Voice cloning works (if supported)
|
||||
- [ ] All response formats work (wav, mp3, flac, pcm)
|
||||
- [ ] Client scripts and server launcher created
|
||||
- [ ] E2E online serving test written (`tests/e2e/online_serving/test_<model>.py`)
|
||||
- [ ] Buildkite CI entry added to `.buildkite/test-merge.yml`
|
||||
- [ ] E2E tests added per model priority tier (see **Test Case Writing (CI Levels)**)
|
||||
- [ ] Buildkite entries match level: `test-ready.yml` / `test-merge.yml` or nightly TTS job / `test-nightly.yml`
|
||||
- [ ] Gradio demo working
|
||||
- [ ] Documentation added (offline + online docs, nav, supported models)
|
||||
|
||||
@@ -500,5 +559,6 @@ Project docs and adjacent skills:
|
||||
- [TTS audio skill](../vllm-omni-audio-tts/SKILL.md) — supported models and usage
|
||||
- [Fish Speech integration](../vllm-omni-audio-tts/references/fish-speech.md) — complete example of Phases 1–3
|
||||
- [Qwen3-TTS reference](../vllm-omni-audio-tts/references/qwen-tts.md) — complete example of all 5 phases
|
||||
- [vllm-omni-test skill](../vllm-omni-test/SKILL.md) — L1–L4 markers, naming, Buildkite wiring, run commands
|
||||
- [Adding a TTS model (developer guide)](https://github.com/vllm-project/vllm-omni/blob/main/docs/contributing/model/adding_tts_model.md)
|
||||
- `plan/voxcpm2_native_ar_design.md` — VoxCPM2's vLLM-native AR + side-computation pattern (distinct from the generator-based single-stage described above)
|
||||
@@ -63,6 +63,8 @@ If multiple rows apply (e.g., a diffusion model is also a new model), union the
|
||||
|
||||
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 ⚠.
|
||||
|
||||
**Also run the Code-Quality sweep on every PR**, regardless of type or mode: the five diff-scoped checks in [references/code-quality.md](references/code-quality.md) — kwargs fragility, broad-except swallow, `Any`/wrong type hints, hot-path `.clone()`/`deepcopy`, and event-loop blocking — plus the advisory conventions (log level, structured logging, synchronization, cleanup, dependencies, naming) in [checklists.md](references/checklists.md). These count only lines the PR *adds* — the pre-existing backlog across the repo is out of scope.
|
||||
|
||||
### Step 5: Print Report
|
||||
|
||||
```
|
||||
@@ -74,6 +76,7 @@ Pre-check report for <branch>
|
||||
Dimension Result
|
||||
───────────────── ──────
|
||||
PR title format ✓
|
||||
Code quality ⚠ 1 broad except, 2 Any hints
|
||||
PR desc integrity ✓
|
||||
Registry/config ✓
|
||||
Dead code ⚠ 2 warnings
|
||||
|
||||
@@ -35,6 +35,31 @@ PR title must follow the project convention documented in `docs/contributing/REA
|
||||
|
||||
---
|
||||
|
||||
## All PRs: Code Quality (run on every PR, quick & full)
|
||||
|
||||
Diff-scoped sweep for five fragility patterns that are pervasive in the existing codebase — **only lines the PR adds count**, not the pre-existing backlog. Detection commands, codebase examples, and full severity rules live in [code-quality.md](code-quality.md).
|
||||
|
||||
- [ ] **No new `**kwargs` string-lookup plumbing:** the diff adds no `**kwargs` + `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`. ⚠ any new instance; ✗ if a new string key is duplicated across ≥2 files without a shared constant, or unknown keys are silently dropped on a fail-fast path.
|
||||
- [ ] **No new broad exception swallow:** the diff adds no `except Exception:` / bare `except:`. ⚠ any new instance; ✗ for `except: pass` / `except Exception: return None` / `continue` on a fail-fast path (init, config validation, weight loading, request handling).
|
||||
- [ ] **No new `Any` / wrong type hints:** the diff adds no `: Any` / `-> Any` / untyped production signature, and no new `SimpleNamespace` in tests faking an object that already has a real typed stub. ⚠ any new instance; ✗ for a *wrong* (actively misleading) annotation or a `SimpleNamespace` test fake of a typed object.
|
||||
- [ ] **No new hot-path copy:** the diff adds no `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path. ⚠ any new instance; ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without a comment, or `deepcopy` of a scheduler / sampling-config on the request path.
|
||||
- [ ] **No new event-loop blocking:** the diff adds no `time.sleep` / blocking HTTP (`requests`, `urllib`) / `await` inside a held lock in `async def` code. ⚠ any new instance; ✗ for blocking sleep or HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock held across an `await`. Blocking inside a dedicated worker thread is fine.
|
||||
|
||||
Roll the ⚠/✗ counts into the report as a single **Code quality** dimension row.
|
||||
|
||||
### Conventions (eyeball — no grep)
|
||||
|
||||
Judgment calls that don't grep cleanly — apply them while reading the diff, not as a mechanical check:
|
||||
|
||||
- **Right log level** — `info` is for end-user-visible events; routine internal events are `debug`; hot paths are `trace` or removed. Logging is not free.
|
||||
- **Structured fields over f-strings** — pass dynamic values as structured fields (`logger.info("msg", extra={...})`), not interpolated strings, so they stay greppable and parseable.
|
||||
- **No reflexive synchronization** — don't add `threading.Lock` / `asyncio.Lock` / `deepcopy` / re-wrap shared state unless there is real concurrent mutation; the GIL covers most Python code. Owners decide their own synchronization — don't pre-lock in a constructor.
|
||||
- **Context managers for cleanup** — use `with` / `contextlib` for files, locks, CUDA streams, and profilers; not raw `open()` + manual `.close()` or ad-hoc `finally`.
|
||||
- **No new dependencies without justification** — prefer the stdlib and existing vllm-omni deps; a new top-level dep needs a reason (check `pyproject.toml` / `requirements*.txt`).
|
||||
- **Naming / interface discipline** — don't rename or change interfaces just for taste in this PR, but do fix names that actively mislead (e.g. `serve` implying a long-running server, `Instance` where multiple instance kinds coexist).
|
||||
|
||||
---
|
||||
|
||||
## Bug Fix PRs
|
||||
|
||||
### Quick
|
||||
@@ -49,7 +74,7 @@ PR title must follow the project convention documented in `docs/contributing/REA
|
||||
- [ ] **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
|
||||
- [ ] **No silent failure risk:** no bare `except: pass`, no `try: ... except: return None`, no empty fallback added. See the broad-except pattern in [code-quality.md](code-quality.md#2-broad-exception-swallow) for the full severity rule (⚠ any new broad catch; ✗ for a swallow on a fail-fast path).
|
||||
- [ ] **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
|
||||
|
||||
|
||||
@@ -0,0 +1,179 @@
|
||||
# Code-Quality Patterns
|
||||
|
||||
Five fragility patterns to sweep on **every** PR, in both quick and full mode. These are pervasive in the existing codebase (hundreds of pre-existing sites) — so the checks are **diff-scoped**: only lines the PR *adds* count. Whole-repo greps light up every PR with the backlog and are useless.
|
||||
|
||||
All commands assume:
|
||||
```bash
|
||||
BASE="$(git merge-base HEAD origin/main)"
|
||||
```
|
||||
|
||||
**Severity:** ⚠ for any *new* instance introduced by the diff; ✗ for the acute sub-cases defined per pattern. Acute ✗ findings belong in the report's blocking count.
|
||||
|
||||
---
|
||||
|
||||
## 1. `**kwargs` string-lookup passthrough
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
def forward(self, **kwargs):
|
||||
# BLOCKER-prone: raw dict + string keys duplicated across files
|
||||
if "runtime_additional_information" in kwargs and \
|
||||
"model_intermediate_buffer" not in kwargs:
|
||||
...
|
||||
```
|
||||
Real instances of this exact guard are copy-pasted across `cosyvoice3.py`, `indextts2_talker.py` (×2), `indextts2_s2mel_decoder.py`, `qwen3_tts_talker.py`, `qwen3_omni.py`. A typo in any string key silently no-ops; there is no type check, no IDE completion, and unknown keys are silently dropped.
|
||||
|
||||
### Why dangerous
|
||||
Raw `dict` + `**kwargs` + lists of param-name strings is **fragile by construction**: the compiler can't see a missing or misspelled key, so failures are silent. Each duplicated copy is a new place for the strings to drift.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `**kwargs` signature whose body does `kwargs["..."]` / `kwargs.get("...")` / `"..." in kwargs`.
|
||||
- ✗ when the diff introduces a **new** string key that is duplicated across ≥2 files without a shared module constant, or when `**kwargs` plumbing **silently drops unknown keys** on a fail-fast path (init, config validation, weight loading).
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E '\*\*kwargs|kwargs(\.get|\[)|"[a-z_]+"\s+in\s+kwargs'
|
||||
```
|
||||
|
||||
### Fix
|
||||
Prefer explicit typed params. The repo already has the right primitive: `msgspec.Struct` with `forbid_unknown_fields=True` rejects unknown keys instead of silently dropping them — see `_StructBase` in `vllm_omni/data_entry_keys.py:105` and `MoriPullRequest` in `vllm_omni/distributed/omni_connectors/connectors/mori_transfer_engine_connector.py:81`. A `@dataclass` or `TypedDict` works too.
|
||||
|
||||
If `**kwargs` is unavoidable (vLLM base-class compat), centralize the key strings as module-level constants behind **one** typed accessor, instead of repeating the literal across files.
|
||||
|
||||
---
|
||||
|
||||
## 2. Broad exception swallow
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# BLOCKER-prone: catches everything, returns None on any failure
|
||||
try:
|
||||
evt = build_event(...)
|
||||
return evt
|
||||
except Exception:
|
||||
return None
|
||||
```
|
||||
Real clusters: `vllm_omni/metrics/stats.py:240,274,367,450` and `vllm_omni/metrics/{modality,utils}.py` — all `except Exception: return None`.
|
||||
|
||||
### Why dangerous
|
||||
A generic `except Exception:` is **too flexible**. If we catch an exception type we don't expect, it is almost always an unhandled edge case that needs to be *fixed* — not swallowed into `None`/`pass`/`continue`. Broad catches turn fixable bugs into silent wrong behavior and make debugging impossible.
|
||||
|
||||
The repo backs this with ruff `BLE001` (broad-exception catch); it is currently suppressed via `# noqa: BLE001` in `vllm_omni/patch.py:413`. Prefer fixing over suppressing.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `except Exception:` / `except BaseException:` / bare `except:`.
|
||||
- ✗ for `except:` / `except Exception: pass|return None|continue` on a **fail-fast path** (init, config validation, weight loading, request handling, connector setup). This extends the Bug-Fix "No silent failure risk" ✗ in [checklists.md](checklists.md).
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
# Any new broad catch (⚠ baseline):
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E 'except\s*(Exception|BaseException)?\s*:'
|
||||
# Note: this matches only the immediate-colon swallow forms (`except :`,
|
||||
# `except Exception:`) — it deliberately excludes specific catches like
|
||||
# `except ValueError:` and the `except Exception as e:` variant. Eyeball
|
||||
# the latter; `as e` usually means it is logged, but still verify.
|
||||
```
|
||||
|
||||
### Fix
|
||||
Catch the **specific** types you actually expect (`ValueError`, `KeyError`, `AttributeError`, `OSError`, `torch.cuda.OutOfMemoryError`, `msgspec.ValidationError`, …). At a genuine top-level / best-effort boundary (metrics, signal handlers), at minimum **log** the exception; never swallow into a bare `pass`/`return None` on a path that should fail loudly.
|
||||
|
||||
---
|
||||
|
||||
## 3. `Any` / wrong type annotations
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# ⚠ Any param + Any return — no contract, no narrowing
|
||||
def _extract_mm_output(engine_outputs: Any) -> dict[str, Any]:
|
||||
...
|
||||
```
|
||||
|
||||
### Why dangerous (root cause)
|
||||
`Any` disables the type checker for that position, so typos and wrong shapes pass silently. The repo has 636 `: Any` params, 244 `-> Any` returns, and **1,173 `SimpleNamespace` in `tests/`**. The `Any` leak is largely test-driven: vibe-coded unit tests fake objects with `SimpleNamespace`, which forces the production code consuming them to be typed `Any`, and the `Any` then propagates outward into public signatures.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `: Any` / `-> Any` / fully-untyped signature in production code.
|
||||
- ✗ for a **wrong** annotation (actively misleading — e.g. `-> bool` that returns `Optional[bool]`), or a *new* `SimpleNamespace` in a test that mimics an object which already has a real typed stub / `@dataclass` / `TypedDict` / `Protocol`.
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
# Any leaking into new signatures:
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E ':\s*Any\b|->\s*Any\b'
|
||||
# SimpleNamespace newly added in tests (the Any-leak source):
|
||||
git diff ${BASE}...HEAD -- "tests/" | grep "^+[^+]" | grep "SimpleNamespace"
|
||||
```
|
||||
|
||||
### Fix
|
||||
Replace `Any` with the concrete type, a `Protocol`, or a `Union`. If a type is genuinely dynamic, prefer `object` + `isinstance` narrowing over `Any` — `object` keeps the checker engaged. In tests, replace `SimpleNamespace` with the real class, a small `@dataclass`/`TypedDict`, or a `Protocol` + fake, so the system under test stays fully typed.
|
||||
|
||||
---
|
||||
|
||||
## 4. Unnecessary copy — `.clone()` / `copy.deepcopy` in hot paths
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# Per-step in a diffusion/AR loop: clones the full latent every iteration
|
||||
for t in timesteps:
|
||||
trajectory_latents.append(x_t.clone()) # O(N) GPU alloc + memcpy each step
|
||||
```
|
||||
Real clusters: `vllm_omni/diffusion/models/bagel/bagel_transformer.py` calls `x_t.clone()` ~9× across a trajectory loop (lines 1820–2186); the cache hooks `vllm_omni/diffusion/cache/teacache/hook.py:155` and `magcache/hook.py:182` clone `hidden_states` per step. Separately, `copy.deepcopy(self.scheduler)` / `copy.deepcopy(req.sampling_params)` is copy-pasted across diffusion pipelines — `pipeline_qwen_image.py:786`, `pipeline_hunyuan_image3.py:1904`, `pipeline_ltx2*.py` (×3), `pipeline_dreamzero.py:1085`, `diffusion_model_runner.py:398` — a full object-graph deepcopy on the request path. (Contrast `dreamzero/causal_wan_model.py:203`, which carries a `# No .clone() needed.` comment — the aware pattern.)
|
||||
|
||||
### Why dangerous
|
||||
A `.clone()` of a GPU tensor or a `copy.deepcopy` of a config object inside a per-step / per-request path is a hidden allocator: it reads as one Python line but triggers a GPU memcpy or a deep object-graph walk every iteration. Across 50 diffusion steps or hundreds of AR steps it dominates latency and VRAM bandwidth. `deepcopy` additionally ignores any custom device / `__copy__` logic and can drag tensors through CPU.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `.clone()` / `.copy_()` / `copy.deepcopy(...)` inside a per-step loop or per-request path.
|
||||
- ✗ for `.clone()` of a full latent/activations tensor inside an AR or diffusion step loop without an explanatory comment, or `copy.deepcopy` of a scheduler / sampling-config on the request hot path.
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E '\.clone\(\)|\.copy_\(|copy\.deepcopy|deepcopy\('
|
||||
```
|
||||
|
||||
### Fix
|
||||
Tensors: prefer a view (`x[..., :k]`, `torch.narrow`, `as_strided`) or write into a pre-allocated buffer with `dst.copy_(src)` instead of `.clone()`; if the caller no longer needs the original, `move`/`del` it rather than copy. Schedulers / configs: build a fresh lightweight instance per request, or split out the mutable per-request bits so the shared immutable part isn't deepcopied. If a copy is genuinely required (e.g. the source is mutated in place later), keep it and add a one-line comment saying why.
|
||||
|
||||
---
|
||||
|
||||
## 5. Blocking the asyncio event loop / lock held across `await`
|
||||
|
||||
### Anti-pattern
|
||||
```python
|
||||
# BLOCKER-prone: blocking sleep / blocking HTTP / lock held across an await,
|
||||
# all inside an async (event-loop) function
|
||||
async def handle(self, req):
|
||||
with self._lock:
|
||||
resp = await self._client.call(req) # lock held across await -> serializes all requests
|
||||
time.sleep(0.1) # stalls the single loop for every concurrent request
|
||||
```
|
||||
|
||||
### Why dangerous
|
||||
vLLM-Omni serves over a single asyncio event loop. Any synchronous blocking work on an `async def` path — `time.sleep`, blocking HTTP (`requests`/`urllib`), heavy CPU work, or a lock held across an `await` — stalls the loop for *every* concurrent request, not just the blocking one. A lock held across `await` silently serializes the whole pipeline; blocking I/O tanks throughput and inflates tail latency.
|
||||
|
||||
### Severity
|
||||
- ⚠ any *new* `time.sleep(`, blocking HTTP call (`requests.get/post`, `urllib.request`, `urlopen`), or `await` inside a `with Lock` / `async with` critical section, in `async def` code.
|
||||
- ✗ for `time.sleep` or blocking HTTP on the serving path (`engine/`, `entrypoints/`, async `connectors/`), or a lock acquired and held across an `await` on the request path.
|
||||
|
||||
### Detect (diff-scoped)
|
||||
```bash
|
||||
# New blocking primitives introduced by the diff — then eyeball whether each
|
||||
# sits in an async function on the loop vs a dedicated worker thread:
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" \
|
||||
| grep -E 'time\.sleep|requests\.(get|post|put|delete)|urllib\.request|urlopen'
|
||||
# Lock acquired with an await potentially inside the critical section (manual eyeball):
|
||||
git diff ${BASE}...HEAD -- "*.py" | grep "^+[^+]" | grep -E 'async with|\.acquire\(\)'
|
||||
```
|
||||
Note: blocking inside a **dedicated worker thread is correct** — the repo already does this (`diffusion/model_loader/hub_prefetch.py`, `distributed/omni_connectors/kv_transfer_manager.py`, the `mooncake`/`mori` connectors poll with `time.sleep` on their own threads). The check targets the *event-loop* path only.
|
||||
|
||||
### Fix
|
||||
Move blocking work off the loop with `asyncio.to_thread(...)` / `run_in_executor`, or use a native async client (`httpx.AsyncClient` / `aiohttp`). Never `await` while holding a lock — drop the lock before the `await` and re-acquire after, or restructure so the critical section contains no `await`.
|
||||
|
||||
---
|
||||
|
||||
## Running the sweep
|
||||
|
||||
Run all five detections against `${BASE}...HEAD`, then for each hit decide ⚠ vs ✗ using the severity rules above. Roll the results into the report as a single **Code quality** dimension row (count of ⚠ and ✗), alongside the type-specific checklist from [checklists.md](checklists.md).
|
||||
@@ -30,6 +30,12 @@ include:
|
||||
- `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
|
||||
- `vllm-omni-test`: guides generation and execution of CI-aligned tests (L1–L4),
|
||||
pytest marker selection (`core_model` / `advanced_model` / `full_model`,
|
||||
`omni` / `tts` / `diffusion`), Buildkite wiring (`test-ready.yml`,
|
||||
`test-merge.yml`, `test-nightly.yml`, `test-weekly.yml`), and copy-paste
|
||||
local plus CI-like `pytest` commands; see `references/test-routing.md` for
|
||||
level-to-command mapping
|
||||
- `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
|
||||
|
||||
@@ -0,0 +1,890 @@
|
||||
---
|
||||
name: vllm-omni-test
|
||||
description: Generate and run tests for vllm-project/vllm-omni with CI-aligned levels and markers; wire new tests into Buildkite (test-ready.yml for L1/L2, test-merge.yml for L3, test-nightly.yml for L4). On completion, always provide copy-paste local and CI-like pytest commands plus prerequisites. Use when creating regression tests, adding L1-L4 coverage, selecting pytest markers, or validating fixes from issues/PRs.
|
||||
---
|
||||
|
||||
# vLLM-Omni Test Generator & Runner
|
||||
|
||||
## Purpose
|
||||
|
||||
Use this skill to generate minimal, stable test cases and run them with the correct marker/level strategy for [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni).
|
||||
|
||||
**Link convention:** Paths such as `.buildkite/` and `docs/contributing/` live at the **vllm-omni repo root**. Markdown links use repo-relative paths from this skill file (e.g. `../../../.buildkite/test-ready.yml`, `../../../docs/contributing/ci/CI_5levels.md`).
|
||||
|
||||
Default priorities:
|
||||
|
||||
1. Reproducible regression coverage for bug fixes
|
||||
2. Correct test level and marker selection
|
||||
3. Low flake, low dependency tests first
|
||||
4. CI-compatible run commands
|
||||
5. **Actionable run commands for the human**: whenever you add or change tests, always finish with **copy-paste-ready** `pytest` lines (local: single file and/or single test; CI-like: markers + `--run-level`), plus short **prerequisites** (GPU tier, HF cache, optional `model_prefix`). Do not assume the reader will infer commands from `test-routing.md` alone.
|
||||
|
||||
## Inputs
|
||||
|
||||
- Issue/PR link and summary
|
||||
- Changed files or suspected code path
|
||||
- Whether the user wants local quick validation or CI-equivalent validation
|
||||
- Hardware constraints (CPU only / CUDA / ROCm / NPU)
|
||||
|
||||
## Workflow
|
||||
|
||||
### Step 1: Classify Test Goal
|
||||
|
||||
- **Bugfix regression**: start from a minimal failing scenario and add assertions that prevent recurrence. Before writing tests, output **`required` / `recommended` / `not_needed`**: **`required`** — stable logic/contract bug that should have been caught; **`recommended`** — environment-sensitive but a small regression still helps; **`not_needed`** — one-off external/config failure or existing tests already cover the path. Prefer the narrowest stable L1 (CPU) case; escalate to L2/L3 only when the bug needs real weights or serving.
|
||||
- **Feature coverage**: verify new behavior and one negative/boundary case.
|
||||
- **Perf/benchmark claim**: require benchmark-oriented tests and explicit metrics.
|
||||
|
||||
### Step 2: Select Test Level
|
||||
|
||||
- **L1**: unit/logic, deterministic, CPU-friendly, fastest feedback.
|
||||
- **L2**: basic e2e/integration and platform-dependent checks.
|
||||
- **L3/L4**: advanced model/integration/perf validation.
|
||||
|
||||
Use [references/test-routing.md](references/test-routing.md) for level-to-marker and command mapping.
|
||||
|
||||
### Step 3: Pick Markers
|
||||
|
||||
Always attach markers deliberately:
|
||||
|
||||
- **Level**: `core_model` (L1/L2) and/or `advanced_model` (L3) and/or `full_model` (L4 nightly)
|
||||
- **Model type** (required on model-centric e2e — pick exactly one):
|
||||
- `omni` — end-to-end multimodal LLM pipelines (thinker/talker/stages; Qwen-Omni family)
|
||||
- `tts` — speech synthesis / TTS-only models (`/v1/audio/speech`, voice clone, etc.)
|
||||
- `diffusion` — generative diffusion models (image / audio / text / video from noise)
|
||||
- **Cross-cutting area** (when relevant): `parallel`, `cache`, `example`, `benchmark`
|
||||
- **Hardware**: `cpu`, `gpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
|
||||
- Optional: `slow`, distributed markers when multi-card is required
|
||||
|
||||
**Baseline smoke (L2 + L3):** The simplest e2e case per model — default deploy, minimal request — should usually carry **both** `@pytest.mark.core_model` and `@pytest.mark.advanced_model` on the **same** test function so `test-ready.yml` and `test-merge.yml` share one test. `send_*_request` picks validation depth from `--run-level`. References: `test_voxcpm2_tts.py::test_text_to_audio_001`, `test_qwen3_tts_customvoice.py::test_text_to_audio_001`. Heavier scenarios use **`advanced_model` only**; L4 expansion uses **`full_model`**.
|
||||
|
||||
For hardware-aware tests, prefer `@hardware_test(...)` or `hardware_marks(...)` in `tests/helpers/mark.py`.
|
||||
|
||||
**Diffusion nightly split** (`test-nightly.yml`): all diffusion tests use `pytest.mark.diffusion`, but CI groups them by **output modality**, not by separate pytest markers:
|
||||
|
||||
| Nightly group | Scope | Typical models / paths |
|
||||
|---------------|--------|-------------------------|
|
||||
| **Diffusion X2I(&A&T) Model Test** | x2**i** (image), x2**a** (audio), x2**t** (text) and other **non-video** diffusion | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero, … — `test_*_expansion.py` under X2I shards; L4 sweep uses `-k "not test_wan and not test_bagel_expansion and not hunyuan"` for L4 |
|
||||
| **Diffusion X2V Model Test** | x2**v** (video) only | Wan2.2, HunyuanVideo 1.5, Wan VACE, … — e.g. `test_wan22_expansion.py`, `test_hunyuan_video_15_expansion.py` |
|
||||
|
||||
Wire new **video** diffusion expansion tests into the **X2V** group; wire **image/audio/text** diffusion expansion tests into **X2I(&A&T)**. Do not place x2v modules in X2I shards (see comment in `test-nightly.yml` above the X2I group).
|
||||
|
||||
### Naming: generated test module files (L2–L4, model-centric e2e)
|
||||
|
||||
When adding **new** pytest modules whose primary scope is a **specific model** (typical under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`), use this filename pattern:
|
||||
|
||||
| Level | Filename pattern | Example (model `Qwen/Qwen2.5-Omni-7B`) |
|
||||
|-------|------------------|----------------------------------------|
|
||||
| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` |
|
||||
| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` |
|
||||
|
||||
**Slug rules for `{lowercase_model_slug}`:**
|
||||
|
||||
1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`.
|
||||
2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree).
|
||||
3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) only then.
|
||||
4. **L1** unit tests are **not** bound to this pattern; use `tests/<area>/test_<feature>.py` as today.
|
||||
|
||||
Routing tables and commands: [references/test-routing.md](references/test-routing.md) § *Model-centric e2e filename convention*.
|
||||
|
||||
Existing references: `tests/e2e/offline_inference/test_qwen2_5_omni.py` (L2-style omni), `tests/e2e/offline_inference/test_qwen3_5_9b.py` (L2-style omni, single-stage VL), `tests/e2e/online_serving/test_qwen3_omni_expansion.py` (L4-style omni), `tests/e2e/online_serving/test_qwen_image_edit_expansion.py` / `test_qwen_image_expansion.py` (L4-style diffusion).
|
||||
|
||||
### Step 4: Generate Test Case Skeleton
|
||||
|
||||
**1. Pick the functional scenario** (then choose directory, fixtures, and markers):
|
||||
|
||||
| Scenario | Typical location | Fixtures / runner pattern | Baseline markers & level |
|
||||
|----------|------------------|---------------------------|---------------------------|
|
||||
| **Offline inference e2e** | `tests/e2e/offline_inference/` | **Module (default):** `omni_runner` + `omni_runner_handler`. **Function (isolation only):** `omni_runner_function` + `omni_runner_handler_function`. Diffusion/TTS may use `Omni(...).generate` directly | L2: `core_model` + **one of** `omni` / `tts` / `diffusion`; `@hardware_test(...)` when GPU/NPU is required |
|
||||
| **Online serving e2e** | `tests/e2e/online_serving/` | **Module (default):** `omni_server` + `openai_client`. **Function (isolation only):** `omni_server_function` + `openai_client_function`. Clients: `send_omni_request` (omni), `send_audio_speech_request` (tts), `send_diffusion_request` / `send_video_diffusion_request` / `send_images_generations_request` (diffusion) | Baseline smoke: **`core_model` + `advanced_model`**; heavier paths: `advanced_model` only; L4 expansion: `full_model` |
|
||||
| **Documentation / runnable examples** | `tests/examples/offline_inference/`, `tests/examples/online_serving/` | **Offline docs (preferred):** extract Python/Bash blocks from the doc README (e.g. `ReadmeSnippet.extract_readme_snippets`), `pytest.mark.parametrize` each snippet, run via `example_runner.run` with a stable `output_subfolder`. **Online docs:** copy client/request scripts into dedicated tests and keep them in sync with the doc page. | Usually **L4**: `advanced_model`, often `example` plus hardware marks matching the nightly docs-example job (see `.buildkite/test-nightly.yml`). Full conventions: [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) (introduced in [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910): naming, output directory layout, skip rules, avoid trimming `num_inference_steps` without a strong CI reason). |
|
||||
| **Performance / benchmark** | `tests/dfx/perf/tests/*.json` + `run_*_benchmark.py` | JSON or script-driven server + load config; assert explicit metrics / baselines | L4 Perf: `full_model` + `benchmark`; wire `test-nightly.yml` Perf steps |
|
||||
| **Invalid parameter / negative HTTP validation** | `tests/dfx/reliability/invalid_param_test/` | Live `omni_server` + low-level `send_*_http_request` with `err_code` / `err_message` | `pytest.mark.slow` + `omni` / `tts` / `diffusion` + `@hardware_marks` (`H100` or `L4`); CI in **`test-weekly.yml`** (not ready/merge/nightly) |
|
||||
|
||||
**If the user’s test plan includes invalid parameter validation / invalid params / negative HTTP / 400 validation:** do **not** add those `test_*` functions to `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move or author them** under `tests/dfx/reliability/invalid_param_test/` in the **endpoint-matching script** (see **Invalid parameter validation** below). Success-path e2e and invalid-param dfx tests must stay in separate modules.
|
||||
|
||||
**1b. Model type — `omni` vs `tts` vs `diffusion`**
|
||||
|
||||
After choosing offline/online/docs/perf, classify the **product under test** and attach **exactly one** model-type marker. All three can live under the same `tests/e2e/` trees; conventions diverge:
|
||||
|
||||
| Dimension | **Omni** (`pytest.mark.omni`) | **TTS** (`pytest.mark.tts`) | **Diffusion** (`pytest.mark.diffusion`) |
|
||||
|-----------|------------------------------|----------------------------|----------------------------------------|
|
||||
| **What it is** | Multimodal LLM pipeline (thinker/talker/stages; text + vision + audio I/O) | Speech synthesis / voice models | Generative diffusion (noise → image, audio, text, or video) |
|
||||
| **Examples** | Qwen2.5-Omni, Qwen3-Omni | Qwen3-TTS, VoxCPM2, Higgs-Audio, Step-Audio2 | Qwen-Image, BAGEL, Wan2.2, HunyuanVideo |
|
||||
| **Offline runner** | `omni_runner` + `omni_runner_handler`, `generate_multimodal` | `Omni(...).generate` or stage YAML + TTS params | `Omni(...).generate` + `OmniDiffusionSamplingParams` |
|
||||
| **Online client** | `openai_client.send_omni_request` (chat completions, modalities) | `openai_client.send_audio_speech_request` (`/v1/audio/speech`) | `send_diffusion_request` (chat/T2I), `send_video_diffusion_request` (`/v1/videos`, X2V), or `send_images_generations_http_request` / `send_images_edits_http_request` (DALL-E routes) |
|
||||
| **Typical assertions** | Stage outputs, text/audio keywords via `OmniRunnerHandler` / response handler | WAV bytes, stream chunks, speech endpoint contract | Image/video dimensions, `final_output_type`, `assert_diffusion_response` |
|
||||
| **Stage / deploy YAML** | Per-model omni stage configs (`ci/qwen3_omni_moe.yaml`, …) | `qwen3_tts.yaml`, `voxcpm2.yaml`, … | Often default serve; parallel/offload YAML for heavy DiT |
|
||||
| **Nightly group** (`test-nightly.yml`) | **Omni Model Test** — `-m "full_model and omni"` | **TTS Model Test** — `-m "full_model and tts"` | **Diffusion X2I(&A&T)** *or* **Diffusion X2V** (see Step 3 table; same `diffusion` marker, different YAML group / file shard) |
|
||||
| **L4 pressure** | Expansion per modality/model as needed | Expansion + accuracy/perf in TTS group | X2I: merge feature combos per [#1832](https://github.com/vllm-project/vllm-omni/issues/1832); X2V: separate nightly group |
|
||||
|
||||
Do **not** mix fixtures across types (e.g. do not use `omni_runner` layout for a pure diffusion or TTS model without mirroring an in-tree test in that family).
|
||||
|
||||
**Diffusion only — X2I(&A&T) vs X2V (nightly routing, not extra markers):**
|
||||
|
||||
- **X2I(&A&T)**: image / audio / text generation — Qwen-Image*, FLUX, SD3, Z-Image, BAGEL (expansion), LongCat, audio diffusion, etc.
|
||||
- **X2V**: video generation only — Wan2.2, HunyuanVideo 1.5, Wan VACE, LTX video similarity paths.
|
||||
|
||||
When adding a new `test_*_expansion.py`, place it in the matching **nightly group** step (explicit file list in `test-nightly.yml`), not only by marker expression.
|
||||
|
||||
**2. Use the narrowest deterministic skeleton for the scenario**
|
||||
|
||||
*L1 unit / logic (CPU-first):*
|
||||
|
||||
**Mocking rule (L1 only):** use **pytest** integration — **`mocker`** (`pytest-mock`) or **`monkeypatch`** (built-in). **Do not** import or call **`unittest.mock`** (`patch`, `MagicMock`, `@patch`, `with patch(...)`, etc.) in L1 tests; patches must auto-revert with the test lifecycle.
|
||||
|
||||
```python
|
||||
import pytest
|
||||
|
||||
pytestmark = [pytest.mark.core_model, pytest.mark.cpu]
|
||||
|
||||
|
||||
def test_<scenario_name>(mocker):
|
||||
# Prefer mocker.patch / mocker.spy / mocker.Mock — not unittest.mock.patch
|
||||
fake_fn = mocker.patch("vllm_omni.some.module.expensive_call", return_value=...)
|
||||
# Act
|
||||
# Assert
|
||||
fake_fn.assert_called_once()
|
||||
|
||||
|
||||
def test_<env_or_attr>(monkeypatch):
|
||||
# For simple env / attribute substitution without a Mock object
|
||||
monkeypatch.setenv("SOME_FLAG", "1")
|
||||
monkeypatch.setattr("vllm_omni.some.module.CONST", 42)
|
||||
...
|
||||
```
|
||||
|
||||
**Avoid in L1:**
|
||||
|
||||
```python
|
||||
# BAD — do not use in L1 unit tests
|
||||
from unittest.mock import patch, MagicMock
|
||||
|
||||
@patch("vllm_omni.some.module.fn")
|
||||
def test_bad(mock_fn): ...
|
||||
|
||||
def test_also_bad():
|
||||
with patch("...") as m: ...
|
||||
```
|
||||
|
||||
See **L1 unit test constraints (mocking)** below for the full do/don't list.
|
||||
|
||||
*Offline multimodal e2e — **Omni** (representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(...)
|
||||
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
|
||||
def test_<scenario>(omni_runner, omni_runner_handler) -> None:
|
||||
request_config = {"prompts": ..., "modalities": [...]} # optional: images, videos, audios
|
||||
omni_runner_handler.send_omni_request(request_config)
|
||||
```
|
||||
|
||||
*Offline generative e2e — **Diffusion** (representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@hardware_test(...)
|
||||
@pytest.mark.parametrize("omni_runner", test_params, indirect=True)
|
||||
def test_text_to_image_001(omni_runner_handler) -> None:
|
||||
omni_runner_handler.send_diffusion_request({"prompt": "...", "extra_body": {"num_inference_steps": 4, ...}})
|
||||
```
|
||||
|
||||
*Offline TTS e2e — **Qwen3-TTS** (two-stage; representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(...)
|
||||
@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
|
||||
def test_text_to_audio_001(omni_runner, omni_runner_handler) -> None:
|
||||
omni_runner_handler.send_audio_speech_request({
|
||||
"input": "...",
|
||||
"task_type": "Base",
|
||||
"ref_audio": REF_AUDIO_URL,
|
||||
"ref_text": REF_TEXT,
|
||||
})
|
||||
```
|
||||
|
||||
*Offline TTS e2e — **single-stage** (Coqui XTTS, MOSS-TTS-Nano; representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(...)
|
||||
@pytest.mark.parametrize("omni_runner", tts_server_params, indirect=True)
|
||||
def test_voice_clone_001(omni_runner_handler) -> None:
|
||||
omni_runner_handler.send_single_stage_tts_request({
|
||||
"input": "...",
|
||||
"language": "en",
|
||||
"prompt_audio_path": REF_AUDIO_PATH,
|
||||
"response_format": "wav",
|
||||
"run_level": "advanced_model",
|
||||
})
|
||||
```
|
||||
|
||||
*Online serving e2e — **Omni** (representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.omni
|
||||
@hardware_test(...)
|
||||
@pytest.mark.parametrize("omni_server", test_params, indirect=True)
|
||||
def test_text_to_text_001(omni_server, openai_client) -> None:
|
||||
request_config = {"model": omni_server.model, "messages": ..., "modalities": ["text"]}
|
||||
openai_client.send_omni_request(request_config)
|
||||
```
|
||||
|
||||
*Online serving e2e — **TTS** (representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.advanced_model
|
||||
@pytest.mark.tts
|
||||
@hardware_test(...)
|
||||
@pytest.mark.parametrize("omni_server", tts_server_params, indirect=True)
|
||||
def test_text_to_audio_001(omni_server, openai_client) -> None:
|
||||
request_config = {"model": omni_server.model, "input": "...", "response_format": "wav", ...}
|
||||
openai_client.send_audio_speech_request(request_config)
|
||||
```
|
||||
|
||||
*Online serving e2e — **Diffusion** X2I (representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
@pytest.mark.parametrize("omni_server", _get_default_case(MODEL), indirect=True)
|
||||
def test_text_to_image_001(omni_server, openai_client) -> None:
|
||||
openai_client.send_diffusion_request({...}) # chat completions + extra_body
|
||||
```
|
||||
|
||||
*Online serving e2e — **Diffusion** X2V (representative):*
|
||||
|
||||
```python
|
||||
@pytest.mark.core_model
|
||||
@pytest.mark.diffusion
|
||||
def test_text_to_video_001(omni_server, openai_client) -> None:
|
||||
openai_client.send_video_diffusion_request({"model": ..., "form_data": {...}}) # /v1/videos
|
||||
```
|
||||
|
||||
*Documentation example tests:* follow the **Preferred Test Strategy** in [l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md): dynamic extraction for offline READMEs; explicit copied client code for online pages until extraction is justified; use the documented **naming**, **output directory** (page folder + case id), and **skipping** rules (e.g. Gradio-only scripts).
|
||||
|
||||
*Performance tests:* add or extend entries under `tests/perf/` (and JSON configs where the project uses them), with **explicit baselines** and the same marker/run-level pairing as the CI step that will execute them.
|
||||
|
||||
**3. Cross-cutting rules**
|
||||
|
||||
- Reuse existing fixtures for the chosen scenario; do not mix “online client” assumptions into offline `OmniRunner` tests without a clear reason.
|
||||
- Avoid external network dependency in assertions unless the scenario is explicitly “online serving” or doc examples that require a model hub (then align with CI secrets/cache).
|
||||
- Keep **one test function = one intent** (one modality combo, one endpoint contract, or one acceleration combo).
|
||||
- **E2E test function layout**: **one case → one `test_<scenario>` function** with a name that states what is validated (endpoint, size/`n`, server flag, or route). **Do not** merge multiple cases into a single test that branches on `request.node.callspec.id`, `param.id`, or `if case_id == ...`. Use `@pytest.mark.parametrize("omni_server", [...], indirect=True)` per function (usually **one** `OmniServerParams` per test). A loop inside one test is OK only when it serves **that function’s single intent** (e.g. three standard sizes in `test_*_sizes_256_512_1024`).
|
||||
- **Runtime fixture scope** (`tests/helpers/fixtures/runtime.py`): default **`omni_server` / `omni_runner`** (module) + matching client/handler; use **`omni_server_function` / `omni_runner_function`** only when each `test_*` must start a fresh instance (see below).
|
||||
- **L1 mocks**: never `unittest.mock`; use `mocker` or `monkeypatch` only (see below).
|
||||
- **API calls (L2+ e2e, online and offline)**: **reuse** `send_*_request` in `tests/helpers/runtime.py` when it exists; **otherwise add** the helper in `runtime.py` first, then call it from the test. General `assert_*` inside `send_*_request`; special `assert_*` only in the test. See **Runtime send helpers** below — do not call `omni.generate`, raw HTTP, or SDK clients from `test_*.py`.
|
||||
- **Response assertions**: reusable checks on API bodies / decoded media belong in **`tests/helpers/assertions.py`** — not as private `_assert_*` helpers inside `test_*.py` (see below).
|
||||
- **Model-specific payloads stay in test modules** — per-model `MODEL`, deploy path, `REF_AUDIO_URL`, `get_prompt()`, `_build_request()`, and inline `request_config` dicts live in `test_{slug}.py` / `test_{slug}_expansion.py` (and offline/L1 siblings). **Do not** create `tests/helpers/{slug}.py` to deduplicate them; a little copy across files is preferred (see `test_glm_tts.py`, `test_cosyvoice3_tts_expansion.py`). `tests/helpers/` is repo-wide harness only (`mark`, `media`, `runtime`, `stage_config`, `assertions`, `fixtures/`).
|
||||
|
||||
#### L1 unit test constraints (mocking)
|
||||
|
||||
L1 tests (`core_model and cpu`, under `tests/diffusion/`, `tests/engine/`, `tests/model_executor/`, etc.) must follow the repo’s **pytest-mock** convention (`pytest-mock>=3.10.0` in `pyproject.toml` `[project.optional-dependencies] dev`).
|
||||
|
||||
| Do | Don't |
|
||||
|----|-------|
|
||||
| `def test_foo(mocker):` + `mocker.patch(...)`, `mocker.spy(...)`, `mocker.Mock()`, `mocker.MagicMock()`, `mocker.AsyncMock()` | `from unittest.mock import patch, MagicMock, Mock, AsyncMock` |
|
||||
| `def test_bar(monkeypatch):` + `monkeypatch.setattr` / `setenv` / `delenv` / `setitem` | `@patch(...)` decorator |
|
||||
| Let `mocker` auto-stop patches after the test | `with patch(...):` / `patch.object(...)` context managers |
|
||||
| Mirror neighboring L1 tests in the same directory | `unittest.mock.create_autospec` unless an existing file already documents an exception |
|
||||
|
||||
**Rationale:** `mocker` ties patch lifecycle to pytest fixtures (no leaked patches across tests). `unittest.mock` decorators/context managers are easy to compose incorrectly with parametrized or async tests and are inconsistent with in-tree L1 style.
|
||||
|
||||
**Minimal patterns:**
|
||||
|
||||
```python
|
||||
def test_returns_cached_config(mocker):
|
||||
loader = mocker.patch(
|
||||
"vllm_omni.foo.load_yaml",
|
||||
return_value={"stages": []},
|
||||
)
|
||||
result = get_deploy_config("ci/foo.yaml")
|
||||
assert result["stages"] == []
|
||||
loader.assert_called_once_with("ci/foo.yaml")
|
||||
|
||||
|
||||
def test_skips_when_env_unset(monkeypatch):
|
||||
monkeypatch.delenv("VLLM_OMNI_FEATURE", raising=False)
|
||||
assert should_enable_feature() is False
|
||||
```
|
||||
|
||||
E2E levels (L2+) generally **avoid mocks**; if a rare L2 stub is unavoidable, still prefer `mocker` over `unittest.mock` for consistency.
|
||||
|
||||
#### Runtime fixtures — scope (`tests/helpers/fixtures/runtime.py`)
|
||||
|
||||
vLLM-Omni e2e tests start a real **OmniServer** (online) or **OmniRunner** (offline). Pick scope by **how often the process must be recreated**, not by test level alone.
|
||||
|
||||
| Scope | Online fixtures | Offline fixtures | When to use |
|
||||
|-------|-----------------|------------------|-------------|
|
||||
| **Module** (default) | `omni_server` → `openai_client` | `omni_runner` → `omni_runner_handler` | **Default for L2/L3/L4 expansion** — amortize model/server init across `test_*` in the same module. Same `OmniServerParams` / runner config can be reused by multiple tests. |
|
||||
| **Function** | `omni_server_function` → `openai_client_function` | `omni_runner_function` → `omni_runner_handler_function` | **Only when required** — each `test_*` must get a **clean** server/runner (no shared engine/GPU state). Typical: `tests/dfx/reliability/`, sleep/wakeup, crash/restart, tests that mutate global server state. |
|
||||
|
||||
**Rules:**
|
||||
|
||||
1. **Default to module scope** (`omni_server` / `omni_runner`) unless the scenario **explicitly** needs a fresh instance per test function.
|
||||
2. **Indirect parametrize name must match the fixture name**: `@pytest.mark.parametrize("omni_server", ...)` with `omni_server` + `openai_client`; `@pytest.mark.parametrize("omni_server_function", ...)` with `omni_server_function` + `openai_client_function`. Do not mix module fixture with function client (or vice versa).
|
||||
3. Different `OmniServerParams` per `test_*` (e.g. default vs `--enable-cpu-offload`) is still OK with **module** `omni_server` — pytest parametrizes per test node; only switch to `_function` when isolation between tests matters, not merely because `server_args` differ.
|
||||
4. One `test_*` with **many** server configs (expansion matrix) → single function + `@pytest.mark.parametrize("omni_server", [...], indirect=True)` + module `omni_server` (see in-tree `test_qwen_image_expansion.py`).
|
||||
|
||||
```python
|
||||
# Default — module-scoped server (L4 expansion)
|
||||
@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=H100)], indirect=True)
|
||||
def test_foo_images_generations_default_1024(omni_server, openai_client) -> None:
|
||||
openai_client.send_images_generations_request({...})
|
||||
|
||||
# Function-scoped — reliability / per-test clean state only
|
||||
@pytest.mark.parametrize(
|
||||
"omni_server_function",
|
||||
[pytest.param(OmniServerParams(model=MODEL), marks=H100)],
|
||||
indirect=True,
|
||||
)
|
||||
def test_foo_sleep_wakeup_cycle(omni_server_function, openai_client_function) -> None:
|
||||
openai_client_function.send_omni_sleep_http_request({...})
|
||||
openai_client_function.send_omni_wakeup_http_request({...})
|
||||
```
|
||||
|
||||
#### Runtime send helpers — online **and** offline (`tests/helpers/runtime.py`)
|
||||
|
||||
**L2+ e2e (online serving and offline inference) must call APIs through `tests/helpers/runtime.py`.** Fixtures live in `tests/helpers/fixtures/runtime.py`; **send/assert implementation** lives in `tests/helpers/runtime.py` + `tests/helpers/assertions.py`.
|
||||
|
||||
| Principle | Action |
|
||||
|-----------|--------|
|
||||
| **Reuse first** | Grep `runtime.py` for an existing **`send_*_request`** on `OpenAIClientHandler` (online) or `OmniRunnerHandler` (offline) that matches the endpoint / pipeline shape. |
|
||||
| **Extend when close** | If an existing `send_*_request` almost fits (e.g. missing one optional field), extend it in `runtime.py` — do not fork logic in the test file. |
|
||||
| **Add when missing** | No suitable helper → add **`send_<feature>_request`** (high-level: call + general `assert_*`) or **`send_<route>_<verb>_http_request`** (low-level HTTP for negative/dfx) in **`runtime.py` first**, then call it from tests. |
|
||||
| **Test module owns payload only** | In `test_*.py`: `MODEL`, deploy path, vendored media, `get_prompt()`, and inline **`request_config` dicts** only. **No** `omni.generate(...)`, raw `requests.post`, OpenAI SDK calls, or `_collect_audio()` / `_process_output()` in e2e tests. |
|
||||
|
||||
**Online** (`openai_client` from `omni_server`): `OpenAIClientHandler.send_*_request`.
|
||||
|
||||
**Offline** (`omni_runner_handler` from `omni_runner`): `OmniRunnerHandler.send_*_request`.
|
||||
|
||||
| Do | Don't |
|
||||
|----|-------|
|
||||
| Online: `openai_client.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_video_diffusion_request`, `send_images_generations_request`, … | `requests.post(f"{base_url}/v1/...", json=…)` or `client.chat.completions.create(...)` inside a test |
|
||||
| Offline: `omni_runner_handler.send_omni_request`, `send_diffusion_request`, `send_audio_speech_request`, `send_single_stage_tts_request`, `send_single_stage_tts_batch_request`, … | `omni_runner.omni.generate(...)` + hand-rolled tensor/WAV extraction in `test_*.py` |
|
||||
| Add missing `send_*` to **`runtime.py`** first; bundle general `assert_*` inside the send helper | A one-off `def _post_*` / `def _collect_audio` at the bottom of a test module |
|
||||
| Mirror naming/style of neighboring `send_*` (docstring, `request_config` dict, optional `run_level`, `err_code` / `err_message` for negative cases) | Different parameter shapes per test file for the same endpoint |
|
||||
|
||||
**Workflow when generating online/offline e2e tests:**
|
||||
|
||||
1. Decide whether the needed check is **general** (every success call of this `send_*`) or **special** (one case / one parameter combo only). See **General vs special assert placement** below.
|
||||
2. **Search `runtime.py`** for a matching **`send_*_request`**; if missing, add it (with general `assert_*` inside) before writing the test body.
|
||||
3. In the test module: build `request_config` → call **`send_*_request` only** for the general contract; call **`assert_*` in the test file only** for special, case-specific checks (import from `assertions.py`).
|
||||
4. Reserve **low-level** `send_*_http_request` for negative/dfx tests (`err_code` / `err_message`) — not for ordinary L2+ success-path e2e.
|
||||
5. When wiring Buildkite `source_file_dependencies`, include **`tests/helpers/runtime.py`** and/or **`tests/helpers/assertions.py`** when new helpers were added.
|
||||
|
||||
**Exceptions (document in the test docstring why):** models whose offline prompt path cannot go through existing handlers yet (e.g. `test_higgs_audio_v2.py`, `test_voxtral_tts.py` with custom tokenizer compose) may call `omni.generate` directly until a `send_*_request` is added to `runtime.py` — treat as **debt**, not the default for new TTS/diffusion/omni e2e.
|
||||
|
||||
**General vs special assert placement:**
|
||||
|
||||
| Kind | Definition | Where it lives | Test file calls |
|
||||
|------|------------|----------------|-----------------|
|
||||
| **General** | Default success contract for this endpoint/client method — e.g. HTTP 200, `data[]` shape, `n` count, `size` dimensions, decodable image/audio, bundled omni/diffusion fields | Implement in `assertions.py`, invoke **inside** `send_*_request` in `runtime.py` (after low-level send / SDK call, when `err_code` is not set) | **`send_*_request` only** — do **not** repeat the same `assert_*` |
|
||||
| **Special** | Extra check tied to **this test case only** — e.g. seed byte-identical replay, accuracy/CLIP threshold, perf ceiling, model-specific optional field | Add or reuse `assert_*` in `assertions.py` (never inline in test) | `send_*_request` **then** `assert_<special>(...)` **once** for that case |
|
||||
|
||||
Changing `request_config` fields (`size`, `n`, `seed`, server flags) is **not** special validation — the general `assert_*` should read those from `request_config` inside `send_*`.
|
||||
|
||||
**`send_*` ↔ `assert_*` pairing (online `OpenAIClientHandler`):**
|
||||
|
||||
| High-level `send_*` (prefer in L2+ e2e) | Assert already invoked inside `runtime.py` | Low-level HTTP-only sibling |
|
||||
|----------------------------------------|--------------------------------------------|-----------------------------|
|
||||
| `send_omni_request` | `assert_omni_response` | `send_chat_completions_http_request` → `assert_http_error` only |
|
||||
| `send_diffusion_request` | `assert_diffusion_response` | — |
|
||||
| `send_audio_speech_request` | `assert_audio_speech_response` | `send_audio_speech_http_request` → `assert_http_error` only |
|
||||
| `send_video_diffusion_request` | `assert_diffusion_response` | `send_videos_*_http_request` → `assert_http_error` only |
|
||||
| `send_images_generations_request` *(add when needed)* | **`assert_images_generations_response`** *(general — inside send)* | `send_images_generations_http_request` → `assert_http_error` only |
|
||||
| `send_images_edits_request` *(add when needed)* | **`assert_images_edits_response`** *(general — inside send)* | `send_images_edits_http_request` → `assert_http_error` only |
|
||||
|
||||
**Rule:** Tests call high-level `send_*_request` for the general contract. **Never** call the same bundled `assert_*` again. Call an extra `assert_*` in the test **only** for special, case-specific validation.
|
||||
|
||||
**Common `OpenAIClientHandler` entry points** (non-exhaustive — grep `runtime.py` before adding):
|
||||
|
||||
| Area | High-level (SDK + assert) | Low-level HTTP (`*_http_request`) |
|
||||
|------|---------------------------|-----------------------------------|
|
||||
| Omni chat | `send_omni_request` | `send_chat_completions_http_request` |
|
||||
| Diffusion T2I (chat route) | `send_diffusion_request` | — |
|
||||
| Diffusion X2V | `send_video_diffusion_request` | `send_videos_create_http_request`, `send_video_content_http_request`, … |
|
||||
| DALL-E T2I / edit | `send_images_generations_request`, `send_images_edits_request` | `send_images_generations_http_request`, `send_images_edits_http_request` |
|
||||
| TTS (online) | `send_audio_speech_request` | `send_audio_speech_http_request`, `send_audio_generate_http_request`, … |
|
||||
| Ops / meta | — | `send_health_http_request`, `send_models_http_request`, `send_omni_sleep_http_request`, … |
|
||||
|
||||
**Common `OmniRunnerHandler` entry points** (offline — grep `runtime.py` before adding):
|
||||
|
||||
| Area | High-level (`send_*_request` + assert) | Notes |
|
||||
|------|----------------------------------------|-------|
|
||||
| Omni multimodal | `send_omni_request` | `generate_multimodal` path |
|
||||
| Diffusion offline | `send_diffusion_request` | chat-route / `OmniTextPrompt` offline |
|
||||
| Qwen-style TTS | `send_audio_speech_request` | two-stage, `generate_multimodal` + `mm_processor_kwargs` |
|
||||
| Single-stage TTS (Coqui XTTS, MOSS-TTS-Nano, …) | `send_single_stage_tts_request`, `send_single_stage_tts_batch_request` | `prompt` + `additional_information` + `omni.generate` — **do not** duplicate in tests |
|
||||
| New model family | **Add** `send_<family>_request` here first | Then call from `tests/e2e/offline_inference/test_*.py` |
|
||||
|
||||
L1 tests under `tests/entrypoints/` may use **FastAPI `TestClient`** or direct handler calls with mocks; they do **not** need `OpenAIClientHandler` / `OmniRunnerHandler`. **L2+ online and offline e2e** must use `runtime.py` `send_*` helpers.
|
||||
|
||||
#### Invalid parameter validation (`tests/dfx/reliability/invalid_param_test/`)
|
||||
|
||||
When the user asks for **invalid parameter validation**, **invalid request bodies**, **HTTP 4xx contract tests**, or any case that sends **malformed / out-of-range / mismatched** API payloads against a **live** server:
|
||||
|
||||
1. **Do not** place these in `tests/e2e/online_serving/` or `*_expansion.py`. If already drafted there, **move** the `test_*` into the correct `invalid_param_test` script and delete the duplicate from e2e.
|
||||
2. **Pick the script by HTTP route** (extend an existing file; add a new `test_invalid_<area>.py` only when no in-tree script covers that route family):
|
||||
|
||||
| API / area | Script |
|
||||
|------------|--------|
|
||||
| Omni chat completions, WebSocket video/realtime paths | `test_invalid_omni_chat.py` |
|
||||
| `POST /v1/audio/speech`, stream, batch, voices | `test_invalid_audio_speech.py` |
|
||||
| Audio diffusion endpoints | `test_invalid_audio_diffusion.py` |
|
||||
| `POST /v1/images/generations` | `test_invalid_image_generation.py` |
|
||||
| `POST /v1/images/edits` | `test_invalid_image_editing.py` |
|
||||
| `POST/GET/DELETE /v1/videos*` | `test_invalid_video_generation.py` |
|
||||
| Sleep / wakeup / server control | `test_invalid_server_control.py` |
|
||||
|
||||
3. **Match in-tree style** in the chosen script:
|
||||
|
||||
| Element | Convention |
|
||||
|---------|------------|
|
||||
| **Module markers** | `pytestmark = [pytest.mark.slow, pytest.mark.<omni or tts or diffusion>]` |
|
||||
| **Server fixture** | `_PARAMS` / `_QWEN3_TTS_SPEECH`-style list of `pytest.param(OmniServerParams(...), id="...", marks=hardware_marks(...))`; `@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)` |
|
||||
| **Hardware** | `hardware_marks(res={"cuda": "H100"})` for heavy diffusion/omni/video; `hardware_marks(res={"cuda": "L4"})` for smaller TTS models (must match weekly `-m "slow and L4"` step) |
|
||||
| **HTTP client** | **Low-level** `openai_client.send_*_http_request({..., "err_code": 400, "err_message": (...)} )` — **not** `send_*_request` (success path) |
|
||||
| **Case shape** | Prefer **one `test_*` per route family** + `@pytest.mark.parametrize("body_spec, err_message", [...])` with stable `id=` per case; or dedicated `test_<route>_malformed_json` when not parametrized |
|
||||
| **Body helpers** | `_minimal_<endpoint>_json(omni_server)` / `_minimal_*_form_data()`; merge overrides with `body.update(body_spec)` |
|
||||
| **Known gaps** | `pytest.mark.skip(reason="…#3649")` as `_SKIP_ISSUE_3649` when server validation is not yet strict (mirror neighboring cases) |
|
||||
| **Sections** | Route banner comments (`# ─── POST /v1/images/generations ───`) like existing files |
|
||||
| **Shared fixtures** | Reuse `tests/dfx/reliability/invalid_param_test/conftest.py` (`tiny_png_bytes`, env defaults) |
|
||||
|
||||
4. **Adding a new model** to invalid-param coverage: append a `pytest.param(OmniServerParams(model="...", stage_config_path=..., server_args=...), ...)` entry to the script’s `_PARAMS` list — do **not** create `test_invalid_<model>.py` unless the route is new.
|
||||
|
||||
5. **CI — `.buildkite/test-weekly.yml` only** (not `test-ready.yml` / `test-merge.yml` / `test-nightly.yml`):
|
||||
|
||||
| Weekly step | Command | When your cases run |
|
||||
|-------------|---------|---------------------|
|
||||
| **Invalid parameters Test · H100** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and H100"` | Diffusion / omni / video invalid-param tests with `H100` hardware mark |
|
||||
| **Invalid parameters Test · L4** | `pytest -s -v tests/dfx/reliability/invalid_param_test/ -m "slow and L4"` | TTS / lighter models marked `L4` |
|
||||
|
||||
- **Trigger:** `build.env("WEEKLY") == "1"` or PR label `weekly-test`.
|
||||
- **Default:** extending an existing `invalid_param_test` script needs **no YAML edit** — the weekly steps already sweep the whole directory.
|
||||
- **Edit YAML only** when adding a **new hardware queue**, a **new top-level script** that must run in isolation, or a **model-specific weekly shard** (mirror neighboring reliability steps).
|
||||
- Weekly steps **do not** use `source_file_dependencies`.
|
||||
|
||||
**Example — append to `test_invalid_image_generation.py`:**
|
||||
|
||||
```python
|
||||
@pytest.mark.parametrize(
|
||||
"body_spec, err_message",
|
||||
[
|
||||
pytest.param({"seed": -1}, ("seed", "greater_than_equal", "0"), id="seed_negative"),
|
||||
],
|
||||
)
|
||||
@pytest.mark.parametrize("omni_server", _PARAMS, indirect=True)
|
||||
def test_images_generations_invalid_requests(
|
||||
omni_server: OmniServer,
|
||||
openai_client: OpenAIClientHandler,
|
||||
body_spec: dict[str, object],
|
||||
err_message: str | tuple[str, ...],
|
||||
) -> None:
|
||||
body = _minimal_images_gen_json(omni_server)
|
||||
body.update(body_spec)
|
||||
openai_client.send_images_generations_http_request(
|
||||
{"json": body, "timeout": 300, "err_code": 400, "err_message": err_message}
|
||||
)
|
||||
```
|
||||
|
||||
See [references/test-routing.md](references/test-routing.md) **Invalid parameter / weekly CI**.
|
||||
|
||||
#### Assertion helpers (`tests/helpers/assertions.py`)
|
||||
|
||||
**Do not** add module-local helpers such as `_assert_images_generations_payload` or `_send_and_assert_*` in e2e test files. Response/media validation belongs in **`tests/helpers/assertions.py`**, grouped by **category**:
|
||||
|
||||
| Category | Existing anchors | When to extend vs add |
|
||||
|----------|------------------|------------------------|
|
||||
| **Image** (chat diffusion + DALL-E JSON) | `assert_image_diffusion_response`, `assert_image_valid` | DALL-E `/v1/images/generations` JSON → add **`assert_images_generations_response`** beside image helpers (reuses `assert_image_valid`); do not fork decode logic into tests |
|
||||
| **Video** | `assert_video_diffusion_response`, `assert_video_valid` | Extend these for new video contracts |
|
||||
| **Audio** | `assert_audio_diffusion_response`, `assert_audio_speech_response`, `assert_audio_valid` | Extend for new TTS/audio endpoints |
|
||||
| **Omni multimodal** | `assert_omni_response` | Extend for new modality combos |
|
||||
| **HTTP errors** | `assert_http_error`, `assert_err_message_in_text` | Used inside low-level `send_*_http_request` and negative tests |
|
||||
|
||||
| Do | Don't |
|
||||
|----|-------|
|
||||
| Put **general** `assert_*` inside `send_*_request` in `runtime.py`; tests call `send_*_request` only | Call `assert_diffusion_response` after `send_diffusion_request` (general assert already bundled) |
|
||||
| Put **special** `assert_*` in `assertions.py` and call it **in the test** after `send_*_request` when that case needs extra checks | Put general decode/count/size logic in the test because “this case uses `n=4`” (that belongs in general assert reading `request_config`) |
|
||||
| Extend the matching category function, or add `assert_<endpoint>_response` next to its category | `_send_and_assert_*` or PIL/base64 loops in `test_*.py` |
|
||||
| Use low-level `send_*_http_request` in dfx/negative tests (`err_code`) | Use low-level HTTP send + manual general assert in every L2+ expansion file |
|
||||
|
||||
**Workflow when generating tests:**
|
||||
|
||||
1. Implement or reuse **general** `assert_*` in `assertions.py` (by category).
|
||||
2. Wire it into **`send_*_request`** in `runtime.py` if not already bundled.
|
||||
3. In the test: `request_config` → `send_*_request`. Add a separate `assert_*` import/call **only** for special case validation.
|
||||
4. Update `source_file_dependencies` when `runtime.py` or `assertions.py` changes.
|
||||
|
||||
**Example — DALL-E `/v1/images/generations`:**
|
||||
|
||||
```python
|
||||
# tests/helpers/assertions.py — general contract (image category)
|
||||
def assert_images_generations_response(resp_body: dict, request_config: dict, *, run_level: str | None = None) -> None:
|
||||
... # data[], n, size → assert_image_valid
|
||||
|
||||
# tests/helpers/runtime.py — general assert INSIDE send
|
||||
def send_images_generations_request(self, request_config: dict[str, Any], ...) -> list[HttpResponse]:
|
||||
responses = self.send_images_generations_http_request(request_config, ...)
|
||||
if request_config.get("err_code") is None:
|
||||
body = responses[0].json_body
|
||||
assert isinstance(body, dict)
|
||||
assert_images_generations_response(body, request_config, run_level=self.run_level)
|
||||
return responses
|
||||
|
||||
# tests/e2e/.../test_foo_expansion.py — one case per test_* (no case_id branching)
|
||||
@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
|
||||
def test_foo_images_generations_default_1024(omni_server, openai_client) -> None:
|
||||
openai_client.send_images_generations_request({"json": body, "timeout": 300})
|
||||
|
||||
@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
|
||||
def test_foo_images_generations_sizes_256_512_1024(omni_server, openai_client) -> None:
|
||||
for size in ("256x256", "512x512", "1024x1024"):
|
||||
openai_client.send_images_generations_request({"json": {**body, "size": size}, "timeout": 300})
|
||||
|
||||
# chat diffusion — separate test; send_diffusion_request bundles assert_diffusion_response
|
||||
@pytest.mark.parametrize("omni_server", [pytest.param(OmniServerParams(model=MODEL), marks=SINGLE_L4)], indirect=True)
|
||||
def test_foo_chat_completions_t2i_fallback(omni_server, openai_client) -> None:
|
||||
openai_client.send_diffusion_request({...})
|
||||
```
|
||||
|
||||
L1 tests under `tests/entrypoints/` may keep **minimal** asserts next to the handler under test when validating a single branch; still prefer `assertions.py` when the same JSON/media check appears in more than one file or level (e.g. L1 protocol test + L4 e2e).
|
||||
|
||||
#### Omni Test Writing Guidance (L1-L4 Layering)
|
||||
|
||||
When the goal is a general Omni/multimodal test case, prioritize mapping the test to the correct purpose, directory, and resource assumptions aligned with the layers (see [CI_5levels.md](../../../docs/contributing/ci/CI_5levels.md)):
|
||||
|
||||
- **L1**: Unit/logic validation on CPU (`core_model and cpu`). Cover input validation, branches, and exception paths (`tests/<component>/test_*.py`). **Mock with `mocker` / `monkeypatch` only — not `unittest.mock`.**
|
||||
- **L2**: Basic e2e (online/offline basic scenarios). Prefer dummy/lightweight models to validate the end-to-end request-to-output-structure/streaming chain (typically `tests/e2e/online_serving/` and `tests/e2e/offline_inference/`).
|
||||
- **L3/L4**: Important integration, performance, and accuracy validation. L4 emphasizes “full functional scenarios + performance/stress + runnable doc examples” (typically `*_expansion.py` plus related expansion cases).
|
||||
|
||||
Also keep markers consistent with the run level: use `core_model` for L1/L2 and `advanced_model` for L3/L4, and pair with `--run-level` to select the intended CI strategy.
|
||||
|
||||
#### Diffusion Test Writing Guidance (L4 Coverage Combinations)
|
||||
|
||||
When the task involves diffusion models/features, organize L4 test cases following [`#1832`](https://github.com/vllm-project/vllm-omni/issues/1832): combine multiple diffusion features into as few test cases as possible to fit limited CI GPU resources.
|
||||
|
||||
Implementation strategy:
|
||||
|
||||
1. If “full L4 is too heavy”: first provide a **reduced local validation case in L1/L2** (so the key assertion and the contract fix point are covered deterministically).
|
||||
2. Then provide **CI/nightly-ready L3/L4 high-marked cases** (e.g. `advanced_model`) to broaden coverage under resource constraints.
|
||||
|
||||
#### L4 nightly sub-pillars — Function vs Accuracy vs Perf
|
||||
|
||||
**L4 is not only `test_*_expansion.py`.** In `test-nightly.yml`, each model-type group typically runs **separate jobs**:
|
||||
|
||||
| Sub-pillar | Nightly step label pattern | What you add | Typical paths |
|
||||
|------------|----------------------------|--------------|---------------|
|
||||
| **Function** | `· Function Test with …` | E2e expansion / feature matrix | `tests/e2e/online_serving/test_<model>_expansion.py` (or offline) |
|
||||
| **Accuracy** | `· Accuracy Test` | Quality / similarity vs baseline | `tests/e2e/accuracy/test_<model>*.py` |
|
||||
| **Perf** | `· Perf Test · <Model>` | Throughput / latency / memory benchmarks | `tests/dfx/perf/tests/test_<model>_vllm_omni.json` + runner script |
|
||||
| **Doc** (optional) | `· Doc Test` | Runnable doc examples | `tests/examples/*/test_text_to_image.py`, … |
|
||||
|
||||
**Default when the user asks for “L4 functional cases”:** deliver **Function** pillar only (`*_expansion.py` + Function Test shard in `test-nightly.yml`). **Do not** silently add Perf or Accuracy unless the user also asks for **performance** / **benchmark** / **accuracy** / **full L4** / **full L4 coverage**.
|
||||
|
||||
**When the user explicitly asks for L4 perf (or full L4 including perf):**
|
||||
|
||||
1. **Do not** put throughput/latency assertions inside `test_*_expansion.py` — perf uses the **dfx benchmark harness**, not `omni_server` e2e fixtures.
|
||||
2. Add or extend a **JSON config** under `tests/dfx/perf/tests/` (mirror in-tree names: `test_qwen_image_vllm_omni.json`, `test_bagel_vllm_omni.json`).
|
||||
3. Each JSON entry: `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]` with explicit **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`, …).
|
||||
4. Run locally via the matching script (model type → runner):
|
||||
|
||||
| Model type | Runner script | Example config |
|
||||
|------------|---------------|----------------|
|
||||
| **Diffusion X2I/X2V** | `tests/dfx/perf/scripts/run_diffusion_benchmark.py` | `tests/dfx/perf/tests/test_qwen_image_vllm_omni.json` |
|
||||
| **TTS** | `tests/dfx/perf/scripts/run_benchmark.py` | `tests/dfx/perf/tests/test_tts.json` (shared) or `test_{slug}.json` (dedicated; use when the model must not join the shared nightly matrix before integration — e.g. `test_voxcpm2.json`, `test_coqui_tts.json`) |
|
||||
| **Omni** | (see existing Omni Perf steps — `run_benchmark.py` / dedicated omni configs in-tree) | `tests/dfx/perf/tests/test_omni*.json` |
|
||||
|
||||
5. Wire **`test-nightly.yml` · Perf Test** step for that model (separate from Function Test): export `DIFFUSION_BENCHMARK_DIR` / `BENCHMARK_DIR`, run pytest on the script with `--test-config-file`, **upload artifacts** (`buildkite-agent artifact upload`), often **multi-GPU H100** for diffusion perf.
|
||||
6. **One benchmark scenario → one `test_name` block** in JSON (same spirit as one `test_*` per function case). Combine server `serve_args` + workload in one entry; use multiple `benchmark_params` rows for size/step sweeps under the same server config.
|
||||
|
||||
**Perf local command (diffusion example):**
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
|
||||
--test-config-file dfx/perf/tests/test_<model>_vllm_omni.json
|
||||
```
|
||||
|
||||
**Clarify in the test plan** which L4 pillars you deliver: `Function only` | `Function + Perf` | `Function + Accuracy + Perf`.
|
||||
|
||||
See [references/test-routing.md](references/test-routing.md) **L4 nightly pillars** for Buildkite step patterns.
|
||||
|
||||
### Step 5: Wire Buildkite (when CI must run the new test)
|
||||
|
||||
If the test is not already collected by an existing pipeline command (for example, L1 tests marked `core_model and cpu` are already covered by the **Simple Unit Test** step in ready/merge), update the appropriate pipeline under `.buildkite/`:
|
||||
|
||||
| Test level | Edit this file | Typical trigger / intent |
|
||||
|------------|----------------|---------------------------|
|
||||
| **L1** and **L2** | [`.buildkite/test-ready.yml`](../../../.buildkite/test-ready.yml) | PR **ready** label; L1 CPU + L2 GPU/basic e2e (steps labeled **Omni ·**, **TTS ·**, **Diffusion ·**) |
|
||||
| **L3** | [`.buildkite/test-merge.yml`](../../../.buildkite/test-merge.yml) | Post-merge; `advanced_model` integration per model type |
|
||||
| **L4** | [`.buildkite/test-nightly.yml`](../../../.buildkite/test-nightly.yml) | Nightly; grouped by model type (see below) |
|
||||
| **Invalid param / reliability (weekly)** | [`.buildkite/test-weekly.yml`](../../../.buildkite/test-weekly.yml) | Weekly; `tests/dfx/reliability/invalid_param_test/` — **not** L1–L4 e2e pipelines |
|
||||
|
||||
**Level-specific Buildkite delivery (follow the user’s requested level):**
|
||||
|
||||
| User asks for | Document / edit **only** | Do **not** add by default |
|
||||
|---------------|---------------------------|---------------------------|
|
||||
| **L4** (`*_expansion.py`, `full_model`) | `test-nightly.yml` — append file to the matching nightly shard (X2I / X2V / Omni / TTS) **or** note it is already collected by an existing `-m` / `-k` sweep | `test-merge.yml` (L3 / `advanced_model`) or `test-ready.yml` (L2) |
|
||||
| **L3** | `test-merge.yml` with `advanced_model` + `source_file_dependencies` | `test-nightly.yml` unless user also wants nightly |
|
||||
| **L2** | `test-ready.yml` with `core_model` + `source_file_dependencies` | merge / nightly pipelines |
|
||||
| **Invalid param** | `test-weekly.yml` — **Invalid parameters Test** group (`-m "slow and H100"` / `-m "slow and L4"`). Usually **no YAML edit** when appending cases to existing `invalid_param_test` scripts | `test-ready.yml`, `test-merge.yml`, `test-nightly.yml` |
|
||||
|
||||
`test-nightly.yml` shards use **explicit pytest file paths** in `commands` (and PR labels like `diffusion-x2iat-test`); they generally **do not** use `source_file_dependencies`. Only suggest merge/ready YAML when the requested level is L3/L2 (or the user explicitly asks for multi-level CI).
|
||||
|
||||
**`test-nightly.yml` top-level groups** (each uses `full_model` + type marker in pytest `-m`):
|
||||
|
||||
| Group | Type marker | Notes |
|
||||
|-------|-------------|--------|
|
||||
| **Omni Model Test** | `omni` | Function / doc / accuracy / **perf** for Qwen-Omni family |
|
||||
| **TTS Model Test** | `tts` | `-m "full_model and L4 and tts"` (and H100 variants when enabled); **Perf Test** uses `run_benchmark.py` + `test_tts.json` |
|
||||
| **Diffusion X2I(&A&T) Model Test** | `diffusion` | Image/audio/text diffusion; explicit file shards for **Function**; **Perf Test · <Model>** uses `run_diffusion_benchmark.py` + JSON config |
|
||||
| **Diffusion X2V Model Test** | `diffusion` | Video-only (Wan, HunyuanVideo, …); PR label `diffusion-x2v-test`; separate **Perf** steps when in-tree |
|
||||
|
||||
Guidelines when editing YAML:
|
||||
|
||||
- **Match markers and `--run-level`** to the test: L1/L2 in `test-ready.yml` use `core_model` + `omni` / `tts` / `diffusion` as appropriate; L3 merge uses `advanced_model`; L4 nightly uses `full_model` with the same type marker.
|
||||
- **Diffusion L4**: add x2i/x2a/x2t expansion files to an **X2I(&A&T)** step; add x2v expansion files to **X2V** — never rely on a broad `tests/e2e/` sweep alone (shards exist to save GPU and avoid wrong queue).
|
||||
- **`source_file_dependencies` (required for new E2E jobs in `test-ready.yml` and `test-merge.yml`)**: every step under the **`:card_index_dividers: E2E Test`** group **must** declare `source_file_dependencies` so Buildkite only schedules the job when relevant paths change. Omitting it breaks path-based CI skipping for GPU e2e jobs.
|
||||
- **Reuse or extend an existing step** when the new test shares the same marker expression, queue, and timeout; otherwise add a new `steps:` entry with the correct `agents.queue`, `timeout_in_minutes`, and docker/kubernetes plugin blocks consistent with neighboring jobs.
|
||||
- **Job timeout — use `timeout_in_minutes` on the step, not `timeout` before pytest**: when adding or documenting a new E2E job in `test-ready.yml` / `test-merge.yml`, set `timeout_in_minutes` on the step and run `pytest` directly in `commands`. **Do not** wrap pytest in `timeout 40m bash -c "..."` / `timeout 20m bash -c "..."`; Buildkite already enforces the step deadline via `timeout_in_minutes`.
|
||||
- **Never omit `plugins` in skill examples or PR YAML snippets**: H100 E2E steps on `mithril-h100-pool` use the full `kubernetes` block (`resources`, `volumeMounts`, `env` with `HF_TOKEN` secret, `nodeSelector: gpu-h100-sxm`, `devshm` + `hf-cache` volumes). L4/docker steps on `gpu_1_queue` use the full `docker#v5.2.0` block (`shm-size`, `HF_HOME`, volumes). Do not write `# ... plugins unchanged` — paste the complete block from a sibling step.
|
||||
- **Platform forks** (e.g. AMD-ready / AMD-merge) live alongside these files; apply the same level → file mapping for those pipelines when the test is platform-specific.
|
||||
|
||||
#### `source_file_dependencies` for E2E Test jobs (`test-ready.yml`, `test-merge.yml`)
|
||||
|
||||
When adding or splitting a step inside the **E2E Test** group, **always** include `source_file_dependencies` on that step. List paths the job actually depends on — at minimum:
|
||||
|
||||
1. **Test module(s)** — every `tests/e2e/...` file the `pytest` command runs (online and offline if both are in one step).
|
||||
2. **Model implementation** — `vllm_omni/model_executor/models/<family>/` and/or `vllm_omni/diffusion/models/<family>/` as applicable.
|
||||
3. **Stage / input plumbing** (omni & multi-stage TTS) — `vllm_omni/model_executor/stage_input_processors/<family>.py` when present.
|
||||
4. **Deploy / stage YAML** — `vllm_omni/deploy/<config>.yaml` or `vllm_omni/deploy/ci/<config>.yaml` referenced by the test’s `OmniServerParams` / `get_deploy_config_path(...)`.
|
||||
5. **Test helpers** (when changed in the same PR) — `tests/helpers/runtime.py` for new `send_*` client methods; `tests/helpers/assertions.py` for new shared `assert_*` helpers.
|
||||
|
||||
Use **repo-relative paths** (no leading `./`). Prefer **directory** entries (`.../models/bagel/`) when the whole package matters; use **file** entries for single test modules and YAML configs. Mirror a neighboring step for the same model family (e.g. **Diffusion · Bagel Test** in `test-ready.yml`).
|
||||
|
||||
**Example** (L2 ready — diffusion online smoke, H100 kubernetes):
|
||||
|
||||
```yaml
|
||||
- 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
|
||||
timeout_in_minutes: 40
|
||||
commands:
|
||||
- |
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v tests/e2e/online_serving/test_bagel.py -m 'core_model' --run-level 'core_model'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 1
|
||||
volumeMounts:
|
||||
- name: devshm
|
||||
mountPath: /dev/shm
|
||||
- name: hf-cache
|
||||
mountPath: /root/.cache/huggingface
|
||||
env:
|
||||
- name: HF_HOME
|
||||
value: /root/.cache/huggingface
|
||||
- name: HF_TOKEN
|
||||
valueFrom:
|
||||
secretKeyRef:
|
||||
name: hf-token-secret
|
||||
key: token
|
||||
nodeSelector:
|
||||
node.kubernetes.io/instance-type: gpu-h100-sxm
|
||||
volumes:
|
||||
- name: devshm
|
||||
emptyDir:
|
||||
medium: Memory
|
||||
- name: hf-cache
|
||||
hostPath:
|
||||
path: /mnt/hf-cache
|
||||
type: DirectoryOrCreate
|
||||
```
|
||||
|
||||
**Example** (L3 merge — TTS online + offline, L4 docker):
|
||||
|
||||
```yaml
|
||||
- label: "TTS · Qwen3-TTS Base Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen3_tts_base.py
|
||||
- tests/e2e/offline_inference/test_qwen3_tts_base.py
|
||||
- vllm_omni/model_executor/models/qwen3_tts/
|
||||
- vllm_omni/model_executor/stage_input_processors/qwen3_tts.py
|
||||
- vllm_omni/deploy/qwen3_tts.yaml
|
||||
timeout_in_minutes: 20
|
||||
commands:
|
||||
- |
|
||||
export VLLM_LOGGING_LEVEL=DEBUG
|
||||
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 and cuda' --run-level 'advanced_model'
|
||||
agents:
|
||||
queue: "gpu_1_queue" # g6.4xlarge instance on AWS, has 1 L4 GPU
|
||||
plugins:
|
||||
- docker#v5.2.0:
|
||||
image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
always-pull: true
|
||||
propagate-environment: true
|
||||
shm-size: "8gb"
|
||||
environment:
|
||||
- "HF_HOME=/fsx/hf_cache"
|
||||
volumes:
|
||||
- "/fsx/hf_cache:/fsx/hf_cache"
|
||||
```
|
||||
|
||||
**Example** (L3 merge — diffusion online + offline, H100 kubernetes):
|
||||
|
||||
```yaml
|
||||
- label: "Diffusion · Qwen Image Test"
|
||||
source_file_dependencies:
|
||||
- tests/e2e/online_serving/test_qwen_image.py
|
||||
- tests/e2e/offline_inference/test_qwen_image.py
|
||||
- vllm_omni/diffusion/models/qwen_image/
|
||||
timeout_in_minutes: 40
|
||||
commands:
|
||||
- |
|
||||
export VLLM_IMAGE_FETCH_TIMEOUT=60
|
||||
pytest -s -v \
|
||||
tests/e2e/online_serving/test_qwen_image.py \
|
||||
tests/e2e/offline_inference/test_qwen_image.py \
|
||||
-m 'advanced_model and cuda' --run-level 'advanced_model'
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
- kubernetes:
|
||||
podSpec:
|
||||
containers:
|
||||
- image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT
|
||||
resources:
|
||||
limits:
|
||||
nvidia.com/gpu: 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
|
||||
```
|
||||
|
||||
When extending an **existing** E2E step to run an additional test file, **append** that file (and any new model/deploy paths) to `source_file_dependencies` on the same step. When authoring a PR, mention the dependency list in the **Buildkite change** section of your test plan.
|
||||
|
||||
#### L3 / L4: validating Buildkite from a feature branch (root `pipeline.yml`)
|
||||
|
||||
Production behavior is driven by [`.buildkite/pipeline.yml`](../../../.buildkite/pipeline.yml): it builds the CI image, then uploads **L2** (`test-ready.yml` on non-`main`), **L3** (`test-merge.yml` when `build.branch == "main"`), and **L4** (`test-nightly.yml` when `build.env("NIGHTLY") == "1"`).
|
||||
|
||||
When you add or debug **L3** or **L4** tests and need those child pipelines to run **off `main`** or **without** `NIGHTLY=1`, apply **temporary** edits (revert before merge):
|
||||
|
||||
1. **Point the upload at the right file (optional shortcut)**
|
||||
In the **Upload Ready Pipeline** step, the command is `buildkite-agent pipeline upload .buildkite/test-ready.yml`. For a one-off run you may change that path to `.buildkite/test-merge.yml` (L3) or `.buildkite/test-nightly.yml` (L4), and adjust the step’s `if:` so it does not conflict with the other upload steps (avoid double-uploading unless intentional).
|
||||
|
||||
2. **L3 — merge pipeline**
|
||||
The gate `if: build.branch == "main"` lives on the **Upload Merge Pipeline** step in `pipeline.yml` (not inside `test-merge.yml`). **Comment out that `if` line** so `test-merge.yml` is uploaded on your feature branch. Alternatively rely on step (1) instead of the dedicated merge upload step.
|
||||
|
||||
3. **L4 — nightly pipeline**
|
||||
**Comment out** `if: build.env("NIGHTLY") == "1"` on the **Upload Nightly Pipeline** step in `pipeline.yml` so the nightly definition is uploaded without setting the env var.
|
||||
Child steps in [`test-nightly.yml`](../../../.buildkite/test-nightly.yml) often repeat `if: build.env("NIGHTLY") == "1"`; **comment out those lines on the steps you need to run**, otherwise they will still be skipped after upload.
|
||||
|
||||
### Step 6: Run Tests
|
||||
|
||||
Pick one command path:
|
||||
|
||||
- **Quick local regression (preferred first)**: single file or `file.py::test_name`
|
||||
- **CI-like level run**: marker expression + `--run-level`
|
||||
|
||||
Full templates live in [references/test-routing.md](references/test-routing.md). **After authoring tests, always emit concrete commands** (see **Output Format**).
|
||||
|
||||
**Typical copy-paste examples** (run from repo `tests/` directory; requires matching vLLM/vllm-omni install and hardware):
|
||||
|
||||
| Area | Example |
|
||||
|------|---------|
|
||||
| **Omni offline L2** | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` |
|
||||
| **Omni online L2** | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` |
|
||||
| **TTS online L2** | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` |
|
||||
| **Diffusion X2I L2 online** | `pytest -s -v e2e/online_serving/test_qwen_image.py -m "core_model and diffusion" --run-level=core_model` |
|
||||
| **Diffusion X2V L2 online** | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` |
|
||||
| **Diffusion X2I L4 nightly** | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` |
|
||||
| **Diffusion X2I L4 perf (nightly)** | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` |
|
||||
| **Diffusion X2V L4 nightly** | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` |
|
||||
| **Invalid param (weekly H100)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` |
|
||||
| **Invalid param (weekly L4)** | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` |
|
||||
| **L1 CPU** | `pytest -s -v -m "core_model and cpu"` |
|
||||
|
||||
**Prerequisites to mention when relevant**: GPU model (e.g. L4 vs H100), `HF_HOME` / token for hub weights, module-level `skipif` (NPU/XPU-only gaps), and whether CI already collects the path (e.g. `test_*_expansion.py` glob in `test-nightly.yml`).
|
||||
|
||||
### Step 7: Validate Result Quality
|
||||
|
||||
Before finishing:
|
||||
|
||||
- Is the new assertion directly tied to the bug/feature contract?
|
||||
- Are API calls made via **`runtime.py` `send_*_request`** (not inline HTTP/SDK in the test file)?
|
||||
- Is the **general** `assert_*` bundled inside `send_*_request` (not duplicated in the test)?
|
||||
- Is each scenario a **separate `test_*` function** (no `case_id` / `if param.id` branching across cases)?
|
||||
- **Fixture scope**: default **`omni_server` / `omni_runner`** (module); **`_function`** variants only when each test must restart the instance?
|
||||
- Is the test deterministic (no fragile timing/network coupling)?
|
||||
- Is runtime appropriate for the selected level?
|
||||
- Are markers and `--run-level` consistent?
|
||||
- **Invalid-param cases** in `tests/dfx/reliability/invalid_param_test/` (not e2e), using `send_*_http_request` + `err_code`, with `pytest.mark.slow` and correct `H100` / `L4` hardware mark?
|
||||
|
||||
## Output Format
|
||||
|
||||
When completing a request, return:
|
||||
|
||||
1. **Test plan** (level, markers, file target, **L4 pillar(s)** if applicable: Function / Accuracy / Perf / Doc; **or** **Invalid param** pillar with target `invalid_param_test/test_invalid_<area>.py`; module basename for e2e: `test_{slug}.py` / `test_{slug}_expansion.py`)
|
||||
2. **Generated/updated test file(s)** — model-specific constants and request payloads stay **inside** those test modules (**never** `tests/helpers/{slug}.py`). Extend **`tests/helpers/runtime.py`** / **`tests/helpers/assertions.py`** only when new **repo-wide** `send_*` or shared assert helpers are required (do not leave ad-hoc HTTP or `_assert_*` in test modules). **If L4 Perf was requested**, also list **`tests/dfx/perf/tests/*.json`**. **If invalid-param cases were requested**, list the **`invalid_param_test/` script** and parametrized `test_*` / new `id=` rows — **not** e2e paths.
|
||||
3. **Buildkite change** — **match the requested level and pillar**:
|
||||
- **L4 Function**: `test-nightly.yml` **Function Test** shard (file list / note existing sweep). **No** `test-merge.yml` unless user also asked for L3.
|
||||
- **L4 Perf**: `test-nightly.yml` **Perf Test · <Model>** step (new step or extend commands); include env exports + artifact upload pattern from a sibling perf job.
|
||||
- **L4 Accuracy**: `test-nightly.yml` **Accuracy Test** step under the same model-type group.
|
||||
- **L3**: `test-merge.yml` + full `source_file_dependencies` + `agents` + `plugins`.
|
||||
- **L2**: `test-ready.yml` + same E2E block requirements.
|
||||
- **Invalid param**: `test-weekly.yml` — note **Invalid parameters Test · H100/L4** group; usually **no YAML change** when only extending existing scripts.
|
||||
If the user asked only for **L4 functional cases**, state explicitly that **Perf / Accuracy / Invalid param** were not included (offer to add if needed). If the plan mixes **success e2e** and **invalid param**, split deliverables across e2e vs `invalid_param_test/` and call out both CI files.
|
||||
4. **Run commands (required)** — always include, in fenced `bash` blocks:
|
||||
- **Local — whole file**: `cd tests` then `pytest -s -v <path> …`
|
||||
- **Local — single test** (optional but preferred when the change is one function): `pytest -s -v path::test_func …`
|
||||
- **CI-like** (when not L1 CPU): the same **marker + `--run-level`** pairing the level uses (see [references/test-routing.md](references/test-routing.md) and **Step 6** table above)
|
||||
- **Prerequisites** (one line): e.g. “Requires CUDA L4 + weights cached locally”; “Requires H100 + `HF_TOKEN`”; “Syntax-only validation when vllm is not installed locally”
|
||||
5. **Result summary** (pass/fail if you executed tests; if not executed, state that explicitly and what the user should run)
|
||||
|
||||
## Additional Resources
|
||||
|
||||
- Marker and command routing: [references/test-routing.md](references/test-routing.md)
|
||||
- CI pipelines (vllm-omni): [test-ready.yml](../../../.buildkite/test-ready.yml) (L1/L2), [test-merge.yml](../../../.buildkite/test-merge.yml) (L3), [test-nightly.yml](../../../.buildkite/test-nightly.yml) (L4), [test-weekly.yml](../../../.buildkite/test-weekly.yml) (invalid param / reliability)
|
||||
- L4 documentation example tests (naming, extraction vs copied scripts, output dirs, skips): [docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md](../../../docs/contributing/ci/test_examples/l4_doc_example_tests.inc.md) — see also [PR #1910](https://github.com/vllm-project/vllm-omni/pull/1910)
|
||||
@@ -0,0 +1,372 @@
|
||||
# Test Routing Reference
|
||||
|
||||
Use this reference to map testing goals to levels, markers, and runnable commands.
|
||||
|
||||
**Repo paths** (`.buildkite/`, `docs/contributing/…`): link with repo-relative paths from this file (e.g. `../../../../.buildkite/test-ready.yml`, `../../../../docs/contributing/ci/CI_5levels.md`).
|
||||
|
||||
## Model-centric e2e filename convention (L2–L4)
|
||||
|
||||
When generating a **new** test module tied to a **specific model** under `tests/e2e/offline_inference/` or `tests/e2e/online_serving/`:
|
||||
|
||||
| Level | Pattern | Example (`Qwen/Qwen2.5-Omni-7B`) |
|
||||
|-------|---------|----------------------------------|
|
||||
| **L2**, **L3** | `test_{lowercase_model_slug}.py` | `test_qwen2_5_omni.py` |
|
||||
| **L4** | `test_{lowercase_model_slug}_expansion.py` | `test_qwen2_5_omni_expansion.py` |
|
||||
|
||||
**Slug rules for `{lowercase_model_slug}`** (canonical: [SKILL.md](../SKILL.md) § *Naming: generated test module files*):
|
||||
|
||||
1. Start from the HuggingFace-style id (e.g. `Qwen/Qwen2.5-Omni-7B`), but **do not** put the org into the filename: use the **repo segment only** (`Qwen2.5-Omni-7B`), not `Qwen_...` / `qwen_qwen2_5_...`.
|
||||
2. **Lowercase**; replace `.`, `-`, and whitespace with a single `_` (e.g. `Qwen2.5-Omni-7B` → `qwen2_5_omni`). **Omit** trailing size tokens such as `7b` / `30b` in the basename when a single file covers that model line in the directory (matches `test_qwen2_5_omni.py` in-tree).
|
||||
3. If two checkpoints in the same folder need separate modules, add a **minimal** disambiguator (e.g. `_7b` vs `_3b`) **only then**.
|
||||
4. **L1** unit tests are **not** bound to this pattern; use `tests/<area>/test_<feature>.py` as today.
|
||||
|
||||
## Model type markers (`omni` / `tts` / `diffusion`)
|
||||
|
||||
Every model-centric e2e test must declare **exactly one** type marker:
|
||||
|
||||
| Marker | Model family | Typical APIs |
|
||||
|--------|--------------|--------------|
|
||||
| `pytest.mark.omni` | Multimodal LLM (Qwen-Omni, …) | `send_omni_request`, `omni_runner`, stage YAML |
|
||||
| `pytest.mark.tts` | TTS / speech synthesis | `send_audio_speech_request`, `/v1/audio/speech` |
|
||||
| `pytest.mark.diffusion` | Diffusion generative models | `send_diffusion_request`, `send_video_diffusion_request`, `send_images_generations_http_request`, `OmniDiffusionSamplingParams` |
|
||||
|
||||
**Diffusion** uses one marker for all output modalities. Nightly CI splits diffusion by **group**, not by extra markers:
|
||||
|
||||
| Nightly group (`test-nightly.yml`) | Diffusion scope | Examples |
|
||||
|-----------------------------------|-----------------|----------|
|
||||
| **Diffusion X2I(&A&T) Model Test** | x2**i** / x2**a** / x2**t** — image, audio, text (non-video) | Qwen-Image*, BAGEL, FLUX, SD3, Z-Image, LongCat, DreamZero |
|
||||
| **Diffusion X2V Model Test** | x2**v** — video only | Wan2.2, HunyuanVideo 1.5, Wan VACE |
|
||||
|
||||
PR labels for selective nightly runs: `diffusion-x2iat-test`, `diffusion-x2v-test` (plus `nightly-test` / `NIGHTLY=1`).
|
||||
|
||||
## Level and Marker Mapping
|
||||
|
||||
| Goal | Suggested Level | Marker baseline | Typical location |
|
||||
|------|------------------|-----------------|------------------|
|
||||
| Unit logic, regression on pure code path | L1 | `core_model and cpu` | `tests/<component>/test_*.py` |
|
||||
|
||||
### L1 unit tests — mocking (`mocker`, not `unittest.mock`)
|
||||
|
||||
L1 modules run in **Simple Test** (`test-ready.yml` / `test-merge.yml`: `-m 'core_model and cpu'`). When isolation requires doubles:
|
||||
|
||||
- **Use** the **`mocker`** fixture (`pytest-mock`): `mocker.patch`, `mocker.spy`, `mocker.Mock`, `mocker.MagicMock`, `mocker.AsyncMock`.
|
||||
- **Use** **`monkeypatch`** for env vars and simple `setattr` / `delattr` without mock objects.
|
||||
- **Do not use** `unittest.mock` — no `from unittest.mock import ...`, no `@patch`, no `with patch(...)`.
|
||||
|
||||
```python
|
||||
# Good (L1)
|
||||
def test_handler(mocker):
|
||||
mocker.patch("vllm_omni.pkg.fn", return_value=0)
|
||||
|
||||
# Bad (L1)
|
||||
from unittest.mock import patch
|
||||
|
||||
@patch("vllm_omni.pkg.fn")
|
||||
def test_handler(mock_fn): ...
|
||||
```
|
||||
|
||||
E2E (L2+) should not rely on mocks unless documenting a rare exception; prefer real (or lightweight) execution paths.
|
||||
| Basic integration/e2e | L2 | `core_model` + **one of** `omni` / `tts` / `diffusion` | `tests/e2e/...` |
|
||||
| Advanced integration | L3 | `advanced_model` + type marker | `tests/e2e/...` |
|
||||
| Full function / nightly | L4 | `full_model` (nightly) + type marker | **Function:** `tests/e2e/*_expansion.py`; **Perf:** `tests/dfx/perf/tests/*.json`; **Accuracy:** `tests/e2e/accuracy/` |
|
||||
| Invalid HTTP / param validation | Weekly (dfx) | `pytest.mark.slow` + type marker + `H100` or `L4` | `tests/dfx/reliability/invalid_param_test/test_invalid_*.py` |
|
||||
|
||||
## Marker Selection Rules
|
||||
|
||||
1. **Level** (pick one):
|
||||
- `core_model` — L1/L2 (`test-ready.yml`)
|
||||
- `advanced_model` — L3 (`test-merge.yml`)
|
||||
- `full_model` — L4 nightly (`test-nightly.yml`); some expansion tests still carry both `advanced_model` and `full_model` during migration
|
||||
2. **Model type** (pick one): `omni`, `tts`, or `diffusion`
|
||||
3. **Cross-cutting** when relevant: `parallel`, `cache`, `example`, `benchmark`
|
||||
4. **Hardware**: `cpu`, `cuda`, `rocm`, `npu`, `L4`, `H100`, `distributed_cuda`, …
|
||||
5. Multi-card: `@hardware_test(...)` in `tests/helpers/mark.py`
|
||||
|
||||
## Command Templates
|
||||
|
||||
### Quick local checks
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v test_xxxx.py
|
||||
```
|
||||
|
||||
### L1
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "core_model and cpu"
|
||||
```
|
||||
|
||||
### L2
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "core_model and not cpu" --run-level=core_model
|
||||
```
|
||||
|
||||
### L3 (merge)
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "advanced_model" --run-level=advanced_model
|
||||
```
|
||||
|
||||
### L4 (nightly)
|
||||
|
||||
**Function** (e2e expansion — default when user asks for “L4 functional cases”):
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model
|
||||
```
|
||||
|
||||
**Perf** (dfx benchmark harness — separate pillar; not `omni_server` fixtures):
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py \
|
||||
--test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json
|
||||
```
|
||||
|
||||
Broad marker sweep (when no explicit file shard):
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "full_model" --run-level=full_model
|
||||
```
|
||||
|
||||
## L4 nightly pillars (Function / Accuracy / Perf)
|
||||
|
||||
Nightly **L4** for a model is often **multiple jobs** in `test-nightly.yml`, not a single pytest module:
|
||||
|
||||
| Pillar | User says | Deliver | Nightly step |
|
||||
|--------|-----------|---------|--------------|
|
||||
| **Function** | “L4 functional cases” / “functional” / `*_expansion.py` | `tests/e2e/.../test_<model>_expansion.py` | `· Function Test with H100/L4` |
|
||||
| **Perf** | “performance” / “perf” / “benchmark” / “full L4” | `tests/dfx/perf/tests/test_<model>_vllm_omni.json` + runner | `· Perf Test · <Model>` |
|
||||
| **Accuracy** | “accuracy” / “similarity” | `tests/e2e/accuracy/test_<model>*.py` | `· Accuracy Test` |
|
||||
|
||||
**Rule:** “L4 functional cases” → **Function only** by default. State that Perf/Accuracy are separate pillars; add them only when requested.
|
||||
|
||||
**Diffusion perf JSON** (`tests/dfx/perf/tests/test_<slug>_vllm_omni.json`): array of objects with `test_name`, `server_params` (`model`, `serve_args`), `benchmark_params[]`, and per-workload **`baseline`** (`throughput_qps`, `latency_mean`, `peak_memory_mb_mean`). Copy structure from `test_qwen_image_vllm_omni.json` in-tree.
|
||||
|
||||
**Diffusion perf nightly step** (under **Diffusion X2I(&A&T)**; copy full `kubernetes` plugins from `· Perf Test · Qwen-Image`):
|
||||
|
||||
```yaml
|
||||
- label: ":full_moon: Diffusion X2I(&A&T) · Perf Test · <Model-Display-Name>"
|
||||
key: nightly-diffusion-x2iat-performance-<slug>
|
||||
timeout_in_minutes: 180
|
||||
commands:
|
||||
- export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results
|
||||
- export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN
|
||||
- export CACHE_DIT_VERSION=1.3.0
|
||||
- |
|
||||
set +e
|
||||
pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_<slug>_vllm_omni.json
|
||||
EXIT=$$?
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json"
|
||||
buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log"
|
||||
exit $$EXIT
|
||||
agents:
|
||||
queue: "mithril-h100-pool"
|
||||
plugins:
|
||||
# … full kubernetes block (often multi-GPU) — copy from Perf Test · Qwen-Image
|
||||
```
|
||||
|
||||
| Model type | Perf runner | Config example |
|
||||
|------------|-------------|----------------|
|
||||
| Diffusion X2I/X2V | `dfx/perf/scripts/run_diffusion_benchmark.py` | `test_qwen_image_vllm_omni.json` |
|
||||
| TTS | `dfx/perf/scripts/run_benchmark.py` | `test_tts.json` (`BENCHMARK_DIR`) |
|
||||
| Omni | `run_benchmark.py` / omni configs in-tree | `test_omni*.json` |
|
||||
|
||||
Do **not** put throughput/latency baselines inside `test_*_expansion.py` — that belongs in the dfx perf JSON + nightly Perf job.
|
||||
|
||||
## Invalid parameter validation (weekly / dfx)
|
||||
|
||||
When the user’s test plan includes **invalid parameter validation**, **invalid request bodies**, or **HTTP 4xx** validation against a live server:
|
||||
|
||||
1. **Do not** author these in `tests/e2e/online_serving/test_*.py` or `*_expansion.py`. **Move** any drafted `test_*` into `tests/dfx/reliability/invalid_param_test/`.
|
||||
2. **Pick script by route** (extend in-tree file):
|
||||
|
||||
| Route family | Script |
|
||||
|--------------|--------|
|
||||
| Omni chat / WS video·realtime | `test_invalid_omni_chat.py` |
|
||||
| `/v1/audio/speech` (+ stream / batch / voices) | `test_invalid_audio_speech.py` |
|
||||
| Audio diffusion | `test_invalid_audio_diffusion.py` |
|
||||
| `/v1/images/generations` | `test_invalid_image_generation.py` |
|
||||
| `/v1/images/edits` | `test_invalid_image_editing.py` |
|
||||
| `/v1/videos*` | `test_invalid_video_generation.py` |
|
||||
| Sleep / wakeup / server control | `test_invalid_server_control.py` |
|
||||
|
||||
3. **Style:** `pytestmark = [pytest.mark.slow, pytest.mark.<type>]`; `_PARAMS` + `hardware_marks`; `send_*_http_request` with `err_code` + `err_message`; parametrized `body_spec` rows with `id=`; `_minimal_*_json()` helpers; `_SKIP_ISSUE_3649` when tracked in [#3649](https://github.com/vllm-project/vllm-omni/issues/3649).
|
||||
4. **CI:** [`.buildkite/test-weekly.yml`](../../../../.buildkite/test-weekly.yml) group **Reliability Test - Invalid parameters Test** — **not** ready/merge/nightly.
|
||||
|
||||
```bash
|
||||
# Weekly H100 (diffusion / omni / video invalid-param)
|
||||
cd tests
|
||||
pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"
|
||||
|
||||
# Weekly L4 (e.g. Qwen3-TTS speech invalid-param)
|
||||
pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"
|
||||
|
||||
# Single script / case
|
||||
pytest -s -v dfx/reliability/invalid_param_test/test_invalid_image_generation.py::test_images_generations_invalid_requests -m "slow and H100"
|
||||
```
|
||||
|
||||
**Trigger:** `WEEKLY=1` or PR label `weekly-test`. Appending cases to existing scripts usually needs **no YAML edit** (directory sweep). No `source_file_dependencies` on weekly steps.
|
||||
|
||||
### Platform-targeted examples
|
||||
|
||||
```bash
|
||||
cd tests
|
||||
pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model
|
||||
```
|
||||
|
||||
### Concrete e2e paths (common in-tree)
|
||||
|
||||
Paths are relative to `tests/` after `cd tests`.
|
||||
|
||||
| Scenario | Example command |
|
||||
|----------|------------------|
|
||||
| **Omni** offline L2 | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py -m "core_model and omni and not cpu" --run-level=core_model` |
|
||||
| **Omni** offline L2 — one test | `pytest -s -v e2e/offline_inference/test_qwen2_5_omni.py::test_text_to_text -m "core_model and omni and not cpu" --run-level=core_model` |
|
||||
| **Omni** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_omni.py -m "core_model and omni" --run-level=core_model` |
|
||||
| **Omni** offline L2 (Qwen3.5-9B VL) | `pytest -s -v e2e/offline_inference/test_qwen3_5_9b.py -m "core_model and omni and not cpu" --run-level=core_model` |
|
||||
| **TTS** online L2 | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py -m "core_model and tts" --run-level=core_model` |
|
||||
| **TTS** online L2 — one test | `pytest -s -v e2e/online_serving/test_qwen3_tts_base.py::test_text_to_audio_001 -m "core_model and tts" --run-level=core_model` |
|
||||
| **Diffusion X2I** offline L2 | `pytest -s -v e2e/offline_inference/test_t2i_model.py -m "core_model and diffusion and not cpu" --run-level=core_model` |
|
||||
| **Diffusion X2I** L2 online smoke | `pytest -s -v e2e/online_serving/test_qwen_image.py e2e/online_serving/test_bagel.py -m "core_model and diffusion" --run-level=core_model` |
|
||||
| **Diffusion X2V** L2 online smoke | `pytest -s -v e2e/online_serving/test_wan22_t2v.py -m "core_model and diffusion" --run-level=core_model` |
|
||||
| **Diffusion X2I** L4 online expansion (H100) | `pytest -s -v e2e/online_serving/test_qwen_image_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` |
|
||||
| **Diffusion X2I** L4 perf (nightly) | `pytest -s -v dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file dfx/perf/tests/test_qwen_image_vllm_omni.json` |
|
||||
| **Diffusion X2I** L4 online expansion (Edit) | `pytest -s -v e2e/online_serving/test_qwen_image_edit_expansion.py -m "full_model and diffusion and H100" --run-level=full_model` |
|
||||
| **Diffusion X2V** L4 nightly | `pytest -s -v e2e/online_serving/test_wan22_expansion.py -m "full_model and cuda" --run-level=full_model` |
|
||||
| **Invalid param** weekly H100 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and H100"` |
|
||||
| **Invalid param** weekly L4 | `pytest -s -v dfx/reliability/invalid_param_test/ -m "slow and L4"` |
|
||||
|
||||
### Agent / author completion checklist
|
||||
|
||||
When adding or modifying tests, do not stop at “where the file lives” — also deliver:
|
||||
|
||||
1. **Local**: `cd tests` + `pytest -s -v <path>` (and `path::test_func` when a single case is enough).
|
||||
2. **CI-like**: marker string + `--run-level` matching the pipeline (`core_model` / `advanced_model` / `full_model`).
|
||||
3. **L1 mocks**: `mocker` / `monkeypatch` only; no `unittest.mock`.
|
||||
4. **API client + assert placement**: **General** validation → implement in `assertions.py`, call **inside** `send_*_request` in `runtime.py`; tests call **`send_*_request` only**. **Special** case validation → `assert_*` in `assertions.py`, called **in the test** after `send_*_request`. Low-level `send_*_http_request` is for negative/dfx tests (`err_code`), not ordinary L2+ success e2e.
|
||||
5. **Shared assertions**: logic in **`tests/helpers/assertions.py`**; general checks inside `send_*_request`; special checks only in the test. No `_assert_*` in `test_*.py`.
|
||||
6. **One case → one `test_*`**: function name reflects what is validated; no `if case_id == ...` mega-test merging multiple scenarios.
|
||||
7. **Fixture scope**: default **`omni_server` + `openai_client`** / **`omni_runner` + `omni_runner_handler`** (module). Use **`omni_server_function` + `openai_client_function`** / **`omni_runner_function` + `omni_runner_handler_function`** only when each `test_*` must spawn a fresh instance. Parametrize name must match fixture (`omni_server` vs `omni_server_function`).
|
||||
8. **Type marker**: `omni`, `tts`, or `diffusion` on every model e2e module.
|
||||
9. **Diffusion L4 Function**: wire `*_expansion.py` into **X2I(&A&T)** or **X2V** **Function Test** in **`test-nightly.yml` only** — do not add `test-merge.yml` unless the user also requested L3.
|
||||
10. **Diffusion L4 Perf** (only when requested): add `tests/dfx/perf/tests/test_<slug>_vllm_omni.json` + **Perf Test · <Model>** step (artifact upload); not part of “L4 functional cases” by default.
|
||||
11. **E2E Buildkite (L2/L3 only)**: `test-ready.yml` / `test-merge.yml` steps need **`source_file_dependencies`** + full **`agents` + `plugins`**. **L4 nightly** uses explicit file lists or perf scripts in `test-nightly.yml` (no merge job).
|
||||
12. **Invalid param**: cases in **`tests/dfx/reliability/invalid_param_test/`** (route-matching script), `send_*_http_request` + `err_code`, `pytest.mark.slow` + `H100`/`L4`; CI = **`test-weekly.yml`** only — do not put in e2e or nightly.
|
||||
13. **Prerequisites**: GPU tier, HF cache/token, and any module `skipif` / platform-only YAML.
|
||||
|
||||
## Buildkite pipeline mapping
|
||||
|
||||
| Level | Repo file | Model-type grouping |
|
||||
|-------|-----------|---------------------|
|
||||
| L1, L2 | [`.buildkite/test-ready.yml`](../../../../.buildkite/test-ready.yml) | Steps prefixed **Omni ·**, **TTS ·**, **Diffusion ·** under **E2E Test** — **`source_file_dependencies` required** |
|
||||
| L3 | [`.buildkite/test-merge.yml`](../../../../.buildkite/test-merge.yml) | Per-model E2E steps; `-m "advanced_model and …"` — **`source_file_dependencies` required** |
|
||||
| L4 | [`.buildkite/test-nightly.yml`](../../../../.buildkite/test-nightly.yml) | **Omni / TTS / Diffusion X2I(&A&T) / X2V** — each group may have **Function**, **Accuracy**, **Perf**, **Doc** steps |
|
||||
| Invalid param (weekly) | [`.buildkite/test-weekly.yml`](../../../../.buildkite/test-weekly.yml) | **Invalid parameters Test · H100** / **· L4** — sweeps `tests/dfx/reliability/invalid_param_test/` |
|
||||
|
||||
### `source_file_dependencies` (E2E Test only — ready & merge)
|
||||
|
||||
Required on each step inside the **E2E Test** group in `test-ready.yml` and `test-merge.yml`. Typical entries:
|
||||
|
||||
| Category | Path pattern |
|
||||
|----------|----------------|
|
||||
| Test module | `tests/e2e/online_serving/test_<slug>.py`, `tests/e2e/offline_inference/test_<slug>.py` |
|
||||
| API client helpers | `tests/helpers/runtime.py` (when the e2e job uses newly added `send_*` methods) |
|
||||
| Shared assert helpers | `tests/helpers/assertions.py` (when the e2e job uses newly added `assert_*` helpers) |
|
||||
| AR / omni model | `vllm_omni/model_executor/models/<family>/` |
|
||||
| Diffusion model | `vllm_omni/diffusion/models/<family>/` |
|
||||
| Stage processor | `vllm_omni/model_executor/stage_input_processors/<family>.py` |
|
||||
| Deploy config | `vllm_omni/deploy/<name>.yaml` or `vllm_omni/deploy/ci/<name>.yaml` |
|
||||
|
||||
Copy the dependency block from the closest in-tree E2E step for the same model family. Update the list whenever the pytest command or server YAML changes.
|
||||
|
||||
### E2E `agents` + `plugins` blocks (required — do not truncate)
|
||||
|
||||
When generating or documenting E2E steps, **always include the complete `plugins` section**. Set **`timeout_in_minutes`** on the step for the job deadline; run **`pytest` directly** in `commands` (no `timeout 40m bash -c` wrapper around pytest).
|
||||
|
||||
Common patterns in `test-ready.yml` / `test-merge.yml`:
|
||||
|
||||
| Queue | Plugin | Required fields |
|
||||
|-------|--------|-----------------|
|
||||
| `mithril-h100-pool` | `kubernetes` | `resources.limits.nvidia.com/gpu`, `volumeMounts` (`devshm`, `hf-cache`), `env` (`HF_HOME`, `HF_TOKEN` via `secretKeyRef`), `nodeSelector: gpu-h100-sxm`, `volumes` (`devshm` emptyDir, `hf-cache` hostPath `/mnt/hf-cache`) |
|
||||
| `gpu_1_queue` / `gpu_4_queue` | `docker#v5.2.0` | `image`, `always-pull`, `propagate-environment`, `shm-size: "8gb"`, `environment: HF_HOME`, `volumes: /fsx/hf_cache` (add `HF_TOKEN` when the step needs hub access) |
|
||||
|
||||
**H100 kubernetes template** (copy from **Diffusion · Bagel Test** / **Diffusion · Qwen Image Test** in-tree):
|
||||
|
||||
```yaml
|
||||
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
|
||||
```
|
||||
|
||||
**L4 docker template** (copy from **TTS · Qwen3-TTS Base Test** in `test-merge.yml`):
|
||||
|
||||
```yaml
|
||||
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"
|
||||
```
|
||||
|
||||
When extending an existing step (e.g. add `tests/e2e/offline_inference/test_qwen_image.py` to **Diffusion · Qwen Image Test**), update `source_file_dependencies` and `commands` only; **keep the existing `agents` + `plugins` block unchanged** unless the hardware tier changes.
|
||||
|
||||
The root [`.buildkite/pipeline.yml`](../../../../.buildkite/pipeline.yml) decides **which** child file is uploaded. To run **L3** on a feature branch, comment out `if: build.branch == "main"` on the merge upload step. To run **L4** without `NIGHTLY=1`, comment out the nightly upload step’s `if: build.env("NIGHTLY") == "1"` and the same `if` on relevant steps inside `test-nightly.yml`. Revert such edits before merging.
|
||||
|
||||
Platform-specific pipelines (e.g. AMD) follow the same level → file pairing under `.buildkite/`.
|
||||
|
||||
## Diffusion RFC (#1832) Alignment Tips
|
||||
|
||||
For **X2I(&A&T)** coverage planning:
|
||||
|
||||
- Prioritize high-value feature combinations with minimal case count.
|
||||
- Split into lightweight L1/L2 validation plus L4 `*_expansion.py` under the **X2I** nightly shards.
|
||||
- **X2V** models (Wan, HunyuanVideo) stay in the **X2V** group — do not merge into X2I sweeps.
|
||||
|
||||
If hardware is insufficient, provide an executable reduced case plus a deferred full CI/nightly plan.
|
||||
@@ -19,6 +19,7 @@ body:
|
||||
|
||||
Consider redacting or replacing sensitive values with placeholders like `<YOUR_TOKEN_HERE>` when sharing configuration or code examples.
|
||||
- type: textarea
|
||||
id: current-environment
|
||||
attributes:
|
||||
label: Your current environment
|
||||
description: |
|
||||
@@ -42,6 +43,7 @@ body:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: code-version
|
||||
attributes:
|
||||
label: Your code version
|
||||
description: |
|
||||
@@ -66,6 +68,7 @@ body:
|
||||
required: true
|
||||
|
||||
- type: textarea
|
||||
id: bug-description
|
||||
attributes:
|
||||
label: 🐛 Describe the bug
|
||||
description: |
|
||||
|
||||
@@ -241,6 +241,7 @@ datasets/
|
||||
*.json
|
||||
!tests/dfx/perf/tests/*.json
|
||||
!apps/ComfyUI-vLLM-Omni/example_workflows/*.json
|
||||
!tests/assets/soulxsinger/*.json
|
||||
*.jsonl
|
||||
*.parquet
|
||||
|
||||
|
||||
@@ -68,6 +68,7 @@ Visit our [documentation](https://vllm-omni.readthedocs.io/en/latest/) to learn
|
||||
- [Installation](https://vllm-omni.readthedocs.io/en/latest/getting_started/installation/)
|
||||
- [Quickstart](https://vllm-omni.readthedocs.io/en/latest/getting_started/quickstart/)
|
||||
- [List of Supported Models](https://vllm-omni.readthedocs.io/en/latest/models/supported_models/)
|
||||
- [Deployment Recipes](https://recipes.vllm.ai) for vLLM-Omni model serving
|
||||
|
||||
## Contributing
|
||||
|
||||
|
||||
@@ -0,0 +1,87 @@
|
||||
# LVSA Long-Video Showcase
|
||||
|
||||
**Training-free block-sparse attention for long-video diffusion**, running on
|
||||
vLLM-Omni via the public plugin entry points (no core changes). LVSA speeds up
|
||||
extended-horizon generation and prevents the freeze/loop failure mode dense
|
||||
attention exhibits beyond a model's training length — with no fine-tuning and no
|
||||
weight changes.
|
||||
|
||||
Upstream: **[LongVideoSparseAttention](https://github.com/JiusiServe/LongVideoSparseAttention)**
|
||||
(algorithm, all model adapters, full docs).
|
||||
|
||||
## Results
|
||||
|
||||
Wan 2.1 1.3B, single A100 80 GB, mean over 5 prompts (Dense vs LVSA-FlashInfer):
|
||||
|
||||
| Horizon | 1× | 2× | 3× | 4× | 5× | 6× |
|
||||
|---|---|---|---|---|---|---|
|
||||
| **Speedup** | 0.9× | 1.4× | 1.9× | 2.5× | 3.0× | **3.5×** |
|
||||
| **VQeval composite** (Dense→LVSA) | 63→63 | 59→63 | 61→63 | 61→64 | 59→64 | 59→63 |
|
||||
|
||||
The speedup **grows with length** (LVSA is the dense regime at 1× by design, so
|
||||
the win starts at 2×), and quality *improves* at extension — LVSA's rotating
|
||||
keyframes prevent the looping/static output dense produces beyond its horizon
|
||||
(VQeval loop-quality +27 to +34 at ≥3×). At long horizons dense **OOMs on 80 GB**
|
||||
(HunyuanVideo ≥2×, Wan2.1-14B and Cosmos at 6×) where LVSA still generates.
|
||||
|
||||
→ Full sweep (5 models × 6 horizons, VQeval + VBench-Long) and the
|
||||
quality-metric methodology:
|
||||
**[benchmarks/](https://github.com/JiusiServe/LongVideoSparseAttention/tree/main/benchmarks)**.
|
||||
|
||||
## Quickstart
|
||||
|
||||
```bash
|
||||
# 1. vLLM-Omni 0.22.0 (stable) + matching vLLM
|
||||
pip install "vllm==0.22.0"
|
||||
pip install --no-build-isolation \
|
||||
"vllm-omni @ git+https://github.com/vllm-project/vllm-omni.git@v0.22.0"
|
||||
|
||||
# 2. LVSA core + plugin
|
||||
git clone https://github.com/JiusiServe/LongVideoSparseAttention.git
|
||||
( cd LongVideoSparseAttention && pip install -e . && pip install -e lvsa-vllm-omni/ )
|
||||
|
||||
# 3. this showcase's client dep
|
||||
pip install requests
|
||||
```
|
||||
|
||||
FlashInfer (used by the default `flashinfer` backend) **ships with vLLM** — step 1
|
||||
already pulls the pinned `flashinfer-python` + `flashinfer-cubin`, so the fused
|
||||
backend works with no extra install. Two optional extras:
|
||||
|
||||
```bash
|
||||
# (optional) faster startup — prebuilt JIT cache, matched to your CUDA + flashinfer
|
||||
CU=cu128 # match your CUDA: cu121 / cu124 / cu128 / …
|
||||
FIVER=$(python -c "import importlib.metadata as m; print(m.version('flashinfer-python'))")
|
||||
pip install --extra-index-url "https://flashinfer.ai/whl/${CU}" \
|
||||
"flashinfer-jit-cache==${FIVER}+${CU}"
|
||||
|
||||
# (fallback) no FlashInfer in your environment? run the SDPA backend instead:
|
||||
export LVSA_BACKEND=sdpa
|
||||
```
|
||||
|
||||
```bash
|
||||
# Serve an LVSA-enabled endpoint (Wan 2.1 1.3B at 2× horizon):
|
||||
MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan FRAMES=161 \
|
||||
bash serve_lvsa.sh
|
||||
|
||||
# In another shell — generate:
|
||||
python generate.py --frames 161 --height 480 --width 832 \
|
||||
--prompt "A dog running through a sunlit forest." --out wan_lvsa.mp4
|
||||
```
|
||||
|
||||
`serve_lvsa.sh` sets the LVSA env vars for the chosen `MODEL_FAMILY`
|
||||
(`wan` / `hunyuan`). Wan and HunyuanVideo run through the LVSA **attention
|
||||
backend** by default — it stays sparse even under sequence-parallel, unlike the
|
||||
monkey-patch hooks.
|
||||
|
||||
## Configuration
|
||||
|
||||
The one knob you must get right is the model's training horizon
|
||||
(`LVSA_REFERENCE_LATENT_FRAMES`: Wan = 21, HunyuanVideo = 33); `serve_lvsa.sh`
|
||||
sets it per family. For the full environment-variable reference, the multi-GPU
|
||||
support matrix, and tuning (`sparsity_scale`, rotating keyframes, non-standard
|
||||
resolutions), see the upstream docs:
|
||||
|
||||
- **[lvsa-vllm-omni plugin reference](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/lvsa-vllm-omni/README.md)** — every `LVSA_*` env var
|
||||
- **[docs/parallelism.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/parallelism.md)** — TP / Ulysses / HSDP support per model
|
||||
- **[docs/tuning.md](https://github.com/JiusiServe/LongVideoSparseAttention/blob/main/docs/tuning.md)** — picking knobs for your model and horizon
|
||||
@@ -0,0 +1,93 @@
|
||||
#!/usr/bin/env python
|
||||
"""Generate a long video through an LVSA-enabled vLLM-Omni server.
|
||||
|
||||
A thin client over the Omni video API — start the server first with
|
||||
``serve_lvsa.sh``. The sparse attention runs server-side; this just submits a
|
||||
job and saves the result.
|
||||
|
||||
Uses the **asynchronous** endpoint (``POST /v1/videos`` → poll
|
||||
``GET /v1/videos/{id}`` → download ``GET /v1/videos/{id}/content``) rather than
|
||||
``POST /v1/videos/sync``. The sync endpoint is capped by a fixed server-side
|
||||
timeout (``VIDEO_SYNC_TIMEOUT_S``, 600 s) and is documented as
|
||||
"for benchmark and testing scenarios" — long videos (LVSA's use case) routinely
|
||||
exceed it. The async job has no such cap.
|
||||
|
||||
python generate.py --frames 481 --height 480 --width 832 \
|
||||
--prompt "A golden retriever running through a sunlit forest." \
|
||||
--out dog_forest_lvsa.mp4
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
import time
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def main():
|
||||
ap = argparse.ArgumentParser()
|
||||
ap.add_argument("--server", default=os.environ.get("OMNI_SERVER", "http://localhost:8000"))
|
||||
ap.add_argument("--prompt", default="A dog running through a sunlit forest.")
|
||||
ap.add_argument("--frames", type=int, default=161) # Wan 2x horizon
|
||||
ap.add_argument("--height", type=int, default=480)
|
||||
ap.add_argument("--width", type=int, default=832)
|
||||
ap.add_argument("--fps", type=int, default=16)
|
||||
ap.add_argument("--steps", type=int, default=50)
|
||||
ap.add_argument("--guidance", type=float, default=5.0)
|
||||
ap.add_argument("--out", default="lvsa_long_video.mp4")
|
||||
ap.add_argument("--poll-interval", type=float, default=5.0, help="Seconds between status polls.")
|
||||
ap.add_argument(
|
||||
"--max-wait", type=float, default=7200.0, help="Give up if the job hasn't finished after this many seconds."
|
||||
)
|
||||
args = ap.parse_args()
|
||||
|
||||
base = args.server.rstrip("/")
|
||||
form = {
|
||||
"prompt": args.prompt,
|
||||
"size": f"{args.width}x{args.height}",
|
||||
"num_frames": str(args.frames),
|
||||
"fps": str(args.fps),
|
||||
"num_inference_steps": str(args.steps),
|
||||
"guidance_scale": str(args.guidance),
|
||||
}
|
||||
|
||||
# ── 1. submit the async job ──────────────────────────────────────────────
|
||||
print(f"[lvsa-app] POST {base}/v1/videos {args.width}x{args.height} {args.frames}f steps={args.steps}")
|
||||
resp = requests.post(f"{base}/v1/videos", data=form, timeout=120)
|
||||
resp.raise_for_status()
|
||||
job = resp.json()
|
||||
vid = job["id"]
|
||||
print(f"[lvsa-app] job {vid} ({job.get('status', '?')})")
|
||||
|
||||
# ── 2. poll until completed / failed ─────────────────────────────────────
|
||||
deadline = time.time() + args.max_wait
|
||||
last = None
|
||||
while True:
|
||||
if time.time() > deadline:
|
||||
raise SystemExit(f"[lvsa-app] timed out after {args.max_wait:.0f}s (job {vid} still running)")
|
||||
time.sleep(args.poll_interval)
|
||||
r = requests.get(f"{base}/v1/videos/{vid}", timeout=60)
|
||||
if r.status_code != 200:
|
||||
# retrieve_video returns a non-200 JSON when the job has FAILED
|
||||
raise SystemExit(f"[lvsa-app] job failed: {r.text}")
|
||||
st = r.json()
|
||||
status = st.get("status")
|
||||
progress = st.get("progress", 0)
|
||||
if (status, progress) != last:
|
||||
print(f"[lvsa-app] {status} {progress}%")
|
||||
last = (status, progress)
|
||||
if status == "completed":
|
||||
break
|
||||
if status == "failed":
|
||||
raise SystemExit(f"[lvsa-app] job failed: {st.get('error')}")
|
||||
|
||||
# ── 3. download the rendered mp4 ─────────────────────────────────────────
|
||||
c = requests.get(f"{base}/v1/videos/{vid}/content", timeout=600)
|
||||
c.raise_for_status()
|
||||
with open(args.out, "wb") as f:
|
||||
f.write(c.content)
|
||||
print(f"[lvsa-app] saved {args.out}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,45 @@
|
||||
#!/usr/bin/env bash
|
||||
# Start an LVSA-enabled vLLM-Omni server for Wan or HunyuanVideo.
|
||||
#
|
||||
# LVSA is selected as the attention backend via --diffusion-attention-config (see
|
||||
# below); the `LVSA_*` env vars only tune it. The `lvsa-vllm-omni` plugin registers
|
||||
# itself through vLLM-Omni's `general_plugins` entry points — nothing in vLLM-Omni
|
||||
# core is modified. Watch the log for "[LVSA] Geometry detected" to confirm sparse
|
||||
# attention is active for your request geometry.
|
||||
#
|
||||
# MODEL=/path/to/Wan2.1-T2V-1.3B-Diffusers MODEL_FAMILY=wan FRAMES=161 bash serve_lvsa.sh
|
||||
# MODEL=/path/to/HunyuanVideo-1.5-... MODEL_FAMILY=hunyuan FRAMES=193 bash serve_lvsa.sh
|
||||
set -euo pipefail
|
||||
|
||||
MODEL=${MODEL:?set MODEL=/path/to/Wan-or-HunyuanVideo-Diffusers}
|
||||
MODEL_FAMILY=${MODEL_FAMILY:-wan} # wan | hunyuan
|
||||
PORT=${PORT:-8000}
|
||||
FRAMES=${FRAMES:-161} # target horizon in frames
|
||||
VAE_T=${VAE_T:-4} # VAE temporal compression (Wan/Hunyuan = 4)
|
||||
BACKEND=${LVSA_BACKEND:-flashinfer} # flashinfer (fused, fastest) | sdpa
|
||||
|
||||
T_LAT=$(( (FRAMES - 1) / VAE_T + 1 ))
|
||||
|
||||
# Per-family training reference horizon (Wan train ref = 21, Hunyuan = 33).
|
||||
case "$MODEL_FAMILY" in
|
||||
wan) REFERENCE=${REFERENCE:-21} ;;
|
||||
hunyuan) REFERENCE=${REFERENCE:-33} ;;
|
||||
*) echo "MODEL_FAMILY must be 'wan' or 'hunyuan'"; exit 1 ;;
|
||||
esac
|
||||
export LVSA_BACKEND="$BACKEND"
|
||||
export LVSA_TOTAL_LATENT_FRAMES="$T_LAT"
|
||||
export LVSA_REFERENCE_LATENT_FRAMES="$REFERENCE"
|
||||
export LVSA_ROTATE_KEYFRAMES=${LVSA_ROTATE_KEYFRAMES:-1}
|
||||
|
||||
echo "[lvsa-app] serving $MODEL ($MODEL_FAMILY) on :$PORT"
|
||||
echo "[lvsa-app] LVSA backend=$BACKEND T_lat=$T_LAT (frames=$FRAMES) reference=$REFERENCE"
|
||||
echo "[lvsa-app] confirm activation by grep '\\[LVSA\\] Geometry detected' in this log"
|
||||
|
||||
# vLLM-Omni 0.22 selects the LVSA attention backend per-role via
|
||||
# --diffusion-attention-config. The backend is what drives sparse attention
|
||||
# here; it engages under tensor-parallel too, unlike the LVSA_*_HOOK monkey-patch
|
||||
# path (which falls back to dense under sequence-parallel), so we deliberately do
|
||||
# NOT set a hook.
|
||||
exec vllm serve "$MODEL" --omni --host 0.0.0.0 --port "$PORT" \
|
||||
--diffusion-attention-config '{"per_role": {"self": {"backend": "LVSA"}}}' \
|
||||
--init-timeout 1800 ${EXTRA_ARGS:-}
|
||||
@@ -79,6 +79,7 @@ Usage:
|
||||
import argparse
|
||||
import ast
|
||||
import asyncio
|
||||
import base64
|
||||
import glob
|
||||
import json
|
||||
import logging
|
||||
@@ -106,6 +107,9 @@ from tqdm.asyncio import tqdm
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
_STAGE_METRICS_ENDPOINTS = {"/v1/chat/completions"}
|
||||
_RETURN_STAGE_METRICS_FIELD = "return_stage_metrics"
|
||||
|
||||
|
||||
class BaseDataset(ABC):
|
||||
def __init__(self, args, api_url: str, model: str):
|
||||
@@ -1013,6 +1017,81 @@ def calculate_metrics(
|
||||
return metrics
|
||||
|
||||
|
||||
def _save_generated_outputs(
|
||||
outputs: list[RequestFuncOutput],
|
||||
requests_list: list[RequestFuncInput],
|
||||
save_dir: str,
|
||||
) -> None:
|
||||
"""Decode and save base64 images/videos from successful responses."""
|
||||
os.makedirs(save_dir, exist_ok=True)
|
||||
saved = 0
|
||||
failed = 0
|
||||
|
||||
for idx, (req, out) in enumerate(zip(requests_list, outputs)):
|
||||
if not out.success or not out.response_body:
|
||||
continue
|
||||
|
||||
if isinstance(out.response_body, bytes):
|
||||
fname = f"req_{idx:04d}.mp4"
|
||||
fpath = os.path.join(save_dir, fname)
|
||||
with open(fpath, "wb") as f:
|
||||
f.write(out.response_body)
|
||||
saved += 1
|
||||
continue
|
||||
|
||||
media_urls: list[str] = []
|
||||
|
||||
# Chat-completions style: choices[*].message.content[*].image_url.url
|
||||
choices = out.response_body.get("choices", [])
|
||||
if isinstance(choices, list):
|
||||
for choice in choices:
|
||||
content = (choice or {}).get("message", {}).get("content")
|
||||
if not isinstance(content, list):
|
||||
continue
|
||||
for item in content:
|
||||
if not isinstance(item, dict) or item.get("type") != "image_url":
|
||||
continue
|
||||
url = (item.get("image_url") or {}).get("url", "")
|
||||
if isinstance(url, str) and url.startswith("data:"):
|
||||
media_urls.append(url)
|
||||
|
||||
# Images endpoint style: data[*].b64_json
|
||||
data_items = out.response_body.get("data", [])
|
||||
if isinstance(data_items, list):
|
||||
for data_item in data_items:
|
||||
if not isinstance(data_item, dict):
|
||||
continue
|
||||
b64_json = data_item.get("b64_json", "")
|
||||
if isinstance(b64_json, str) and b64_json:
|
||||
media_urls.append(f"data:image/png;base64,{b64_json}")
|
||||
|
||||
for img_idx, url in enumerate(media_urls):
|
||||
if "," not in url:
|
||||
continue
|
||||
|
||||
try:
|
||||
header, b64_data = url.split(",", 1)
|
||||
ext = "png"
|
||||
if "image/jpeg" in header:
|
||||
ext = "jpg"
|
||||
elif "image/webp" in header:
|
||||
ext = "webp"
|
||||
elif "video/mp4" in header:
|
||||
ext = "mp4"
|
||||
|
||||
img_bytes = base64.b64decode(b64_data)
|
||||
fname = f"req_{idx:04d}_{img_idx}.{ext}"
|
||||
fpath = os.path.join(save_dir, fname)
|
||||
with open(fpath, "wb") as f:
|
||||
f.write(img_bytes)
|
||||
saved += 1
|
||||
except Exception as e:
|
||||
failed += 1
|
||||
logger.warning(f"Failed to save image for request {idx}: {e}", exc_info=True)
|
||||
|
||||
logger.info(f"Saved {saved} generated image(s) to {save_dir}. Failed to save {failed} image(s).")
|
||||
|
||||
|
||||
def wait_for_service(base_url: str, timeout: int = 120) -> None:
|
||||
print(f"Waiting for service at {base_url}...")
|
||||
start_time = time.time()
|
||||
@@ -1095,6 +1174,10 @@ async def benchmark(args):
|
||||
requests_list = dataset.get_requests()
|
||||
print(f"Prepared {len(requests_list)} requests from {args.dataset} dataset.")
|
||||
|
||||
if args.return_stage_metrics and args.endpoint in _STAGE_METRICS_ENDPOINTS:
|
||||
for req in requests_list:
|
||||
req.extra_body.setdefault(_RETURN_STAGE_METRICS_FIELD, True)
|
||||
|
||||
if args.endpoint == "/v1/images/edits":
|
||||
for req in requests_list:
|
||||
req.default_bot_task = args.bot_task
|
||||
@@ -1202,6 +1285,9 @@ async def benchmark(args):
|
||||
|
||||
print("\n" + "=" * 60)
|
||||
|
||||
if args.save_dir:
|
||||
_save_generated_outputs(outputs, requests_list, args.save_dir)
|
||||
|
||||
if args.output_file:
|
||||
with open(args.output_file, "w") as f:
|
||||
json.dump(metrics, f, indent=2)
|
||||
@@ -1278,11 +1364,18 @@ if __name__ == "__main__":
|
||||
default=1,
|
||||
help="Number of warmup requests to run before measurement.",
|
||||
)
|
||||
# NOTE Changed default from 1 to 2 because some models (e.g., Bagel) run
|
||||
# `num_timesteps - 1` denoising iterations. A default of 1 results in 0 steps,
|
||||
# which causes errors.
|
||||
# TODO If this slightly longer warmup causes regression issues for other
|
||||
# diffusion pipelines in the future, consider implementing model-specific
|
||||
# overrides instead of a global default.
|
||||
parser.add_argument(
|
||||
"--warmup-num-inference-steps",
|
||||
type=int,
|
||||
default=1,
|
||||
help="num_inference_steps used for warmup requests.",
|
||||
default=2,
|
||||
help="Number of inference steps used for warmup requests. "
|
||||
"Default is 2 to ensure at least one denoising step is executed.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--warmup-concurrency",
|
||||
@@ -1323,6 +1416,13 @@ if __name__ == "__main__":
|
||||
default=3.0,
|
||||
help="SLO target multiplier: slo_ms = estimated_exec_time_ms * slo_scale (default: 3).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Directory to save generated images/outputs for visual inspection. "
|
||||
"If not set, generated outputs are discarded after metric collection.",
|
||||
)
|
||||
parser.add_argument("--disable-tqdm", action="store_true", help="Disable progress bar.")
|
||||
parser.add_argument(
|
||||
"--enable-negative-prompt",
|
||||
@@ -1358,6 +1458,11 @@ if __name__ == "__main__":
|
||||
default="think",
|
||||
help=("bot_task form field for --endpoint /v1/images/edits (think, recaption, think_recaption, vanilla)."),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--return-stage-metrics",
|
||||
action="store_true",
|
||||
help="Request stage duration metrics from endpoints that support return_stage_metrics.",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
asyncio.run(benchmark(args))
|
||||
File diff suppressed because it is too large.
Load diff
@@ -80,6 +80,17 @@ models:
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTS-Local-Transformer-v1.5: Qwen3 backbone + 1-layer GPT2-style local
|
||||
# depth transformer (n_vq=12), paired with the 48 kHz stereo
|
||||
# MOSS-Audio-Tokenizer-v2 codec. Same voice-clone request shape as the
|
||||
# other MOSS-TTS variants.
|
||||
OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5:
|
||||
supported_tasks: [voice_clone]
|
||||
backend: openai-audio-speech
|
||||
endpoint: /v1/audio/speech
|
||||
task_extra_body:
|
||||
voice_clone: {}
|
||||
|
||||
# MOSS-TTSD synthesizes [S1]/[S2] dialogue. The seed-tts voice_clone dataset
|
||||
# has single-speaker rows; bench piggybacks on it (both speakers share the
|
||||
# same ref_audio). For real dialogue evaluation use the ttsd dataset.
|
||||
|
||||
@@ -2,7 +2,7 @@
|
||||
# The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must
|
||||
# be updated on every rebase to match the new upstream vllm commit/tag.
|
||||
ARG VLLM_BASE_IMAGE=vllm/vllm-openai
|
||||
ARG VLLM_BASE_TAG=v0.22.0
|
||||
ARG VLLM_BASE_TAG=v0.23.0
|
||||
FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG}
|
||||
ARG APP_DIR=/workspace/vllm-omni
|
||||
WORKDIR ${APP_DIR}
|
||||
@@ -16,6 +16,7 @@ RUN apt-get update && \
|
||||
|
||||
RUN uv pip install --system ".[dev]"
|
||||
|
||||
|
||||
RUN ln -sf /usr/bin/python3 /usr/bin/python
|
||||
|
||||
ENTRYPOINT []
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai:v0.23.0
|
||||
FROM ${BASE_IMAGE}
|
||||
|
||||
ARG COMMON_WORKDIR=/app
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.18.0rc1
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.18.0rc1-a3
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
# WORKDIR /vllm-workspace/vllm
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=nightly-main
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
@@ -1,5 +1,5 @@
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/ascend/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=nightly-main-a3
|
||||
ARG VLLM_ASCEND_IMAGE=quay.io/atlas-ci/vllm-ascend
|
||||
ARG VLLM_ASCEND_TAG=v0.23.0-4885d8e-a3
|
||||
FROM ${VLLM_ASCEND_IMAGE}:${VLLM_ASCEND_TAG}
|
||||
|
||||
ARG APP_DIR=/vllm-workspace/vllm-omni
|
||||
|
||||
@@ -1,4 +1,4 @@
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0
|
||||
ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.23.0
|
||||
FROM ${BASE_IMAGE} AS base
|
||||
|
||||
# Declare a variable to know if we want to use the nightly build or the stable build.
|
||||
|
||||
@@ -88,7 +88,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL}
|
||||
ENV UV_INDEX_STRATEGY="unsafe-best-match"
|
||||
ENV UV_LINK_MODE="copy"
|
||||
|
||||
ARG VLLM_VERSION=v0.22.0
|
||||
ARG VLLM_VERSION=v0.23.0
|
||||
RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm
|
||||
WORKDIR /workspace/vllm
|
||||
|
||||
@@ -101,7 +101,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \
|
||||
source /opt/intel/oneapi/ccl/2021.15/env/vars.sh --force && \
|
||||
export CMAKE_PREFIX_PATH="$(python3 -c 'import site; print(site.getsitepackages()[0])'):${CMAKE_PREFIX_PATH}"
|
||||
|
||||
ENV LD_LIBRARY_PATH="$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
ENV LD_LIBRARY_PATH="/opt/intel/oneapi/ccl/2021.15/lib:$LD_LIBRARY_PATH:/usr/local/lib/"
|
||||
|
||||
ENV VLLM_TARGET_DEVICE=xpu
|
||||
ENV VLLM_WORKER_MULTIPROC_METHOD=spawn
|
||||
|
||||
@@ -18,7 +18,6 @@ nav:
|
||||
- Offline Inference:
|
||||
- BAGEL-7B-MoT: user_guide/examples/offline_inference/bagel.md
|
||||
- GLM-Image Multistage End-to-End Inference: user_guide/examples/offline_inference/glm_image.md
|
||||
- Helios Video Generation: user_guide/examples/offline_inference/helios.md
|
||||
- HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md
|
||||
- Image-To-Image: user_guide/examples/offline_inference/image_to_image.md
|
||||
- Image-To-Video: user_guide/examples/offline_inference/image_to_video.md
|
||||
|
||||
Binary file not shown.
|
Before Width: | Height: | Size: 403 KiB After Width: | Height: | Size: 410 KiB |
@@ -129,6 +129,27 @@ The precheck covers five PR types: Bug Fix, Performance, New Model, Diffusion Mo
|
||||
### 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.
|
||||
|
||||
### Automatic skip-ci (docs and pytest skip marks)
|
||||
|
||||
On pull requests and `main` pushes, the bootstrap step in [`.buildkite/pipeline.yml`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/pipeline.yml) runs [`.buildkite/scripts/upload_pipeline.py`](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/scripts/upload_pipeline.py) against the git diff. When every changed file qualifies, **L2 (`ready`) and L3 (`merge-test`) pipelines are not uploaded**, so the default GPU CI jobs are skipped.
|
||||
|
||||
| Change per file | Examples |
|
||||
| --- | --- |
|
||||
| Documentation | `docs/**`, any `*.md`, `mkdocs.yml` |
|
||||
| Pytest skip marks only (under `tests/`) | Add/remove/edit `@pytest.mark.skip`, `@pytest.mark.skipif`, or `pytest.skip(...)`; reformat `pytestmark` only to add a skip/skipif alongside existing `pytest.mark.*` entries |
|
||||
| New skipped test module | New `tests/**/*.py` whose `pytestmark` includes unconditional `pytest.mark.skip` |
|
||||
|
||||
These PR shapes all trigger skip-ci:
|
||||
|
||||
- Documentation only
|
||||
- Qualifying skip-mark edits in `tests/**/*.py` only
|
||||
- **A mix of documentation and qualifying skip-mark test edits**
|
||||
|
||||
Skip-ci does **not** apply when the diff also touches product code (for example `vllm_omni/`), or when test files change assertions, imports, fixtures, or other non-skip logic. If the diff cannot be resolved (non-PR branches outside `main`), CI runs as usual.
|
||||
|
||||
!!! note
|
||||
Skipping L2/L3 does **not** disable the Docker image build step. Nightly (L4) upload can still run when the PR has a `nightly-test` label or on scheduled `main` builds with `NIGHTLY=1`. See [CI levels](./ci/CI_5levels.md) for how bootstrap skip-ci relates to diff-gated E2E jobs.
|
||||
|
||||
### Code Quality
|
||||
|
||||
The PR needs to meet the following code quality standards:
|
||||
|
||||
@@ -11,6 +11,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<tr>
|
||||
<th>Level</th>
|
||||
<th>Scope & Focus</th>
|
||||
<th>Model Coverage Strategy</th>
|
||||
<th>Feature Coverage Strategy</th>
|
||||
<th>Interface Coverage Strategy</th>
|
||||
<th>Tags</th>
|
||||
<th>Time Cost</th>
|
||||
<th>Test Dir</th>
|
||||
<th>Doc</th>
|
||||
@@ -24,6 +28,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td>Contribution Guideline & PR checklist</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>.github/PULL_REQUEST_TEMPLATE.md <a href="../tests_style/"> Test Style (PR Checklist)</a></td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
@@ -32,6 +40,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<td>CI Failure Description</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td><a href="../failures/"> CI Failures</a></td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
@@ -39,6 +51,10 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<tr>
|
||||
<td><strong>L1</strong><br>(Unit & Logic)</td>
|
||||
<td>Unit tests for components like entrypoints, models</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td>/</td>
|
||||
<td><code>core_model and cpu</code></td>
|
||||
<td rowspan="2"><15min</td>
|
||||
<td>/tests/{component_name}/test_xxx</td>
|
||||
<td>
|
||||
@@ -50,10 +66,18 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
</tr>
|
||||
<tr>
|
||||
<td><strong>L2</strong><br>(E2E across models & GPU-required UT)</td>
|
||||
<td>Online & Offline (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
|
||||
<td>Online (basic deployment scenarios):<br>dummy, normal inference function (output format, stream), some instance startup UT</td>
|
||||
<td>High-priority models + online basic scenarios + request success validation</td>
|
||||
<td>High-priority features (using random lightweight models)</td>
|
||||
<td>High-priority interfaces (using random lightweight models)</td>
|
||||
<td><code>core_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
|
||||
<td>
|
||||
<strong>Model tests:</strong><br>
|
||||
/tests/e2e/online_serving/test_{model_name}.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}.py
|
||||
<strong>Feature tests:</strong><br>
|
||||
/tests/{component_name}/test_xxx<br>
|
||||
<strong>Interface tests:</strong><br>
|
||||
/tests/entrypoints/test_xxx
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-1-l1-l2-level-testing-unit-testing-and-basic-end-to-end-verification">Chapter 1</a><br>
|
||||
@@ -65,10 +89,19 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
<tr>
|
||||
<td><strong>L3</strong><br>(Important Perf & Integration & Accuracy)</td>
|
||||
<td>Online & Offline (multiple deployment scenarios):<br>real model, normal inference function, normal accuracy</td>
|
||||
<td>High/medium-priority models with real weights + online/offline key scenarios + basic accuracy validation</td>
|
||||
<td>Medium-priority features (using random lightweight models)</td>
|
||||
<td>Medium-priority interfaces (using random lightweight models)</td>
|
||||
<td><code>advanced_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
|
||||
<td><30min</td>
|
||||
<td>
|
||||
<strong>Model tests:</strong><br>
|
||||
/tests/e2e/online_serving/test_{model_name}.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}.py
|
||||
/tests/e2e/offline_inference/test_{model_name}.py<br>
|
||||
<strong>Feature tests:</strong><br>
|
||||
/tests/{component_name}/test_xxx<br>
|
||||
<strong>Interface tests:</strong><br>
|
||||
/tests/entrypoints/test_xxx
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-2-l3-level-testing-core-integration-performance-and-accuracy-verification">Chapter 2</a><br>
|
||||
@@ -79,18 +112,27 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
</tr>
|
||||
<tr>
|
||||
<td><strong>L4</strong><br>(Perf & Integration & Accuracy)</td>
|
||||
<td>Online & Offline: full functional scenarios + performance test + doc test</td>
|
||||
<td>Online: full functional scenarios + performance test + doc test + accuracy test</td>
|
||||
<td>High-priority models: function, performance, accuracy, and doc testing<br>Medium/low-priority models: function and doc testing</td>
|
||||
<td>Low-priority features (using random lightweight models)</td>
|
||||
<td>Low-priority interfaces (using random lightweight models)</td>
|
||||
<td><code>full_model and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
|
||||
<td><3 hour</td>
|
||||
<td>
|
||||
<strong>Full Function:</strong><br>
|
||||
<strong>Model tests:</strong><br>
|
||||
/tests/e2e/online_serving/test_{model_name}_expansion.py<br>
|
||||
/tests/e2e/offline_inference/test_{model_name}_expansion.py<br>
|
||||
<strong>Feature tests:</strong><br>
|
||||
/tests/{component_name}/test_xxx<br>
|
||||
<strong>Interface tests:</strong><br>
|
||||
/tests/entrypoints/test_xxx<br>
|
||||
<strong>Performance:</strong><br>
|
||||
/tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),<br>
|
||||
and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)<br>
|
||||
<strong>Doc Test:</strong><br>
|
||||
tests/example/online_serving/test_{model_name}.py<br>
|
||||
tests/example/offline_inference/test_{model_name}.py
|
||||
tests/example/offline_inference/test_{model_name}.py<br>
|
||||
<strong>Accuracy Test:</strong><br>
|
||||
/tests/e2e/accuracy/test_{model_name}.py
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-3-l4-level-testing-full-functionality-performance-and-documentation-testing">Chapter 3</a><br>
|
||||
@@ -101,7 +143,11 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
</tr>
|
||||
<tr>
|
||||
<td><strong>L5</strong><br>(Stability & Reliability)</td>
|
||||
<td>Online & Offline: long-term stability test + reliability test</td>
|
||||
<td>Online: long-term stability test + reliability test</td>
|
||||
<td>Long-term stability and reliability testing for high-priority models</td>
|
||||
<td>/</td>
|
||||
<td>Invalid-parameter validation for high-priority interfaces</td>
|
||||
<td><code>slow and hardware_test(H100, L4, etc.) and omni/tts/diffusion</code></td>
|
||||
<td> Depends on reality </td>
|
||||
<td>
|
||||
<strong>Stability:</strong><br>
|
||||
@@ -109,7 +155,7 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari
|
||||
/tests/dfx/stability/tests/test_wan22.json<br>
|
||||
<strong>Reliability:</strong><br>
|
||||
tests/dfx/reliability/test_reliability_{model_key}.py<br>
|
||||
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>)
|
||||
(e.g. <code>test_reliability_qwen3_omni.py</code>, <code>test_reliability_wan22.py</code>, <code>test_reliability_hunyuan_image.py</code>, <code>test_reliability_voxcpm2.py</code>)
|
||||
</td>
|
||||
<td>
|
||||
<a href="#chapter-4-l5-level-testing-stability-and-reliability-testing">Chapter 4</a><br>
|
||||
@@ -695,14 +741,16 @@ L5 level testing focuses on the performance of model services under ***long-runn
|
||||
### 4.2 Testing Content and Scope
|
||||
|
||||
- ***Long-term Stability (Stability) Testing***: Uses JSON under `tests/dfx/stability/tests/` (for example `test_qwen3_omni.json` and `test_wan22.json`) to run the service under moderate load for an extended period (e.g., over 12 hours), monitoring whether metrics like memory/VRAM usage, response time, and throughput degrade over time, and whether the service process remains stable.
|
||||
- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni or `multiprocessing.spawn` for Wan2.2 video workers), **large multimodal chat** or **`/v1/videos`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
|
||||
- ***Reliability Testing***: Uses pytest suites under `tests/dfx/reliability/` to inject controlled faults against a **live** `vllm_omni serve` instance (same **`omni_server` / `omni_server_function`** fixture style as E2E). Current suites emphasize **GPU memory pressure** (CUDA sidecar “memory hog”), **worker / runtime process kill** (`SIGKILL` on `VLLM::Worker` for Qwen3-Omni, `multiprocessing.spawn` for Wan2.2 video workers, or `vLLM-Omni::` for HunyuanImage DiT workers), **large multimodal chat**, **`/v1/videos`**, **`/v1/images/generations`**, or **`/v1/audio/speech`** jobs under OOM, **`/health` → 503** and **fast-fail / non-hanging concurrent** requests after kill, and **OpenAI-style 5xx error contracts** (e.g. text vs text+audio under OOM). **Post-fault recovery** checks exist where enabled (some cases may be `skip` while issues are tracked). See the Reliability `<details>` block in Section 4.4 for file-level responsibilities and CI markers (`slow`, `hardware_test`, POSIX-only kill).
|
||||
|
||||
### 4.3 Test Directory and Execution Files
|
||||
|
||||
- ***Stability Test Configuration***: `tests/dfx/stability/tests/test_qwen3_omni.json`, `tests/dfx/stability/tests/test_wan22.json` (one JSON per model / runner family)
|
||||
- ***Reliability Test Suite*** (`tests/dfx/reliability/`):
|
||||
- `test_reliability_qwen3_omni.py` — Qwen3-Omni chat / multimodal reliability (GPU OOM, process kill, recovery, error contract under `--async-chunk` vs default).
|
||||
- `test_reliability_wan22.py` — Wan2.2 T2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `test_reliability_wan22.py` — Wan2.2 I2V video API reliability (`/v1/videos` under OOM and process kill, recovery).
|
||||
- `test_reliability_hunyuan_image.py` — HunyuanImage-3.0-Instruct DiT-only reliability (`/v1/images/generations` under OOM and process kill; deploy `hunyuan_image3_dit.yaml`, H100 × 4).
|
||||
- `test_reliability_voxcpm2.py` — VoxCPM2 TTS reliability (`/v1/audio/speech` under OOM and process kill; deploy `voxcpm2.yaml`, L4 × 1).
|
||||
- `helpers.py` — Shared primitives used by current suites: raw HTTP probes for `/v1/chat/completions` and `/health`, OpenAI-style error parsing, GPU OOM sidecar (`inject_gpu_oom` / `stop_gpu_oom_hogs`), and `pgrep`-based process-kill injector construction (`make_process_kill_fault_injector`).
|
||||
- `conftest.py` — `fault_injector` and `omni_server_after_fault` / `omni_server_after_fault_function` fixtures to run a callable **after** the server is ready.
|
||||
- `README.md` — Short local run commands for this directory.
|
||||
@@ -800,22 +848,33 @@ Reliability tests are **short fault-injection** integration runs (L5 **(b)** in
|
||||
| `helpers.py` | Shared helpers used by current reliability suites: raw `POST`/`GET` probes (`/v1/chat/completions`, `/health`), OpenAI error parsing (`extract_openai_error_contract_from_bytes`), GPU OOM sidecar lifecycle (`inject_gpu_oom`, `stop_gpu_oom_hogs`), and process-kill injector builder (`make_process_kill_fault_injector`). |
|
||||
| `conftest.py` | Pytest fixtures: indirect `fault_injector`, `omni_server_after_fault` / `omni_server_after_fault_function` (run injector after server is ready, then yield server). |
|
||||
| `test_reliability_qwen3_omni.py` | Qwen3-Omni: OOM vs **text vs text+audio** error contract, large multimodal chat under OOM, concurrent pressure, **SIGKILL** on `VLLM::Worker`, `/health` → 503 + fast-fail + concurrent chat; optional OOM recovery scenario (may be skipped while tracked in issues). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 T2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `test_reliability_wan22.py` | Wan2.2 I2V: large `/v1/videos` under OOM, **SIGKILL** on `multiprocessing.spawn` chain, health / fast-fail / concurrent video requests; optional recovery test (may be skipped). |
|
||||
| `test_reliability_hunyuan_image.py` | HunyuanImage DiT-only: large `/v1/images/generations` under OOM, **SIGKILL** on `vLLM-Omni::` workers and serve/tree targets, health / fast-fail / concurrent image requests; some OOM/recovery cases may be skipped while tracked in issues. |
|
||||
| `test_reliability_voxcpm2.py` | VoxCPM2: `/v1/audio/speech` under OOM (error contract), **SIGKILL** on `VLLM::` workers and serve/tree targets, health / fast-fail / concurrent speech requests; some OOM cases may be skipped while tracked in issues. |
|
||||
| `README.md` | Minimal run / collect examples. |
|
||||
|
||||
#### Parametrization and markers
|
||||
|
||||
- Each test module defines a **`RELIABILITY_SCENARIOS`** list (`test_name`, `server_params`: model, `stage_config_name` or diffusion `server_args`, etc.). **`create_reliability_omni_server_params()`** in `tests/dfx/conftest.py` resolves stage paths (including XPU substitutions where applicable) and builds **`OmniServerParams`** lists consumed by **`@pytest.mark.parametrize(..., indirect=True)`** on `omni_server` or `omni_server_function`.
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** (Qwen3-Omni paths often require **2** cards; Wan2.2 video paths **1** card).
|
||||
- Cases are tagged **`@pytest.mark.slow`** for weekly / selective CI. GPU-heavy suites use **`@hardware_test(res={"cuda": "H100"}, num_cards=...)`** or **`@hardware_test(res={"cuda": "L4"}, num_cards=1)`** (Qwen3-Omni **2**× H100; Wan2.2 **1**× H100; HunyuanImage DiT **4**× H100; VoxCPM2 **1**× L4).
|
||||
- **Process-kill** tests use **`@pytest.mark.skipif(os.name == "nt", ...)`** because injection uses POSIX **`pgrep` / `kill`**.
|
||||
|
||||
#### CI trigger
|
||||
|
||||
Weekly Buildkite (`.buildkite/test-weekly.yml`) runs, for example:
|
||||
Weekly Buildkite (`.buildkite/test-weekly.yml`) runs one step per model suite (trigger: `WEEKLY=1` or PR label `weekly-test`), for example:
|
||||
|
||||
| Buildkite step | Test file | CI hardware |
|
||||
| -------------- | --------- | ----------- |
|
||||
| Reliability Test - qwen3-omni | `test_reliability_qwen3_omni.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - wan22 | `test_reliability_wan22.py` | H100 × 2 (`mithril-h100-pool`) |
|
||||
| Reliability Test - hunyuan-image | `test_reliability_hunyuan_image.py` | H100 × 4 (`mithril-h100-pool`) |
|
||||
| Reliability Test - voxcpm2 | `test_reliability_voxcpm2.py` | L4 × 1 (`gpu_1_queue`) |
|
||||
|
||||
```bash
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m "slow"
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m "slow"
|
||||
```
|
||||
|
||||
#### Local commands
|
||||
@@ -824,6 +883,8 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m "slow"
|
||||
pytest --collect-only tests/dfx/reliability
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
pytest -s -v tests/dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
```
|
||||
|
||||
#### Adding a new model suite
|
||||
@@ -836,7 +897,7 @@ pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow
|
||||
</details>
|
||||
|
||||
- - ***Stability***: `pytest -s -v tests/dfx/stability/scripts/test_stability_qwen3_omni.py` or `pytest -s -v tests/dfx/stability/scripts/test_stability_wan22.py` (or add `test_stability_<model>.py` alongside a matching JSON config)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_qwen3_omni.py -m slow` and/or `pytest -s -v tests/dfx/reliability/test_reliability_wan22.py -m slow` (add `test_reliability_<suite>.py` for new models)
|
||||
- ***Reliability***: `pytest -s -v tests/dfx/reliability/test_reliability_<model>.py -m slow` (current suites: `qwen3_omni`, `wan22`, `hunyuan_image`, `voxcpm2`; add `test_reliability_<suite>.py` and a matching step in `.buildkite/test-weekly.yml` for new models)
|
||||
|
||||
## Summary
|
||||
|
||||
|
||||
@@ -96,15 +96,21 @@ Our test scripts use the pytest framework. First, please use `git clone https://
|
||||
# Stability: Wan2.2 (v1/videos diffusion benchmark loop)
|
||||
pytest -s -v dfx/stability/scripts/test_stability_wan22.py -m slow
|
||||
|
||||
# Reliability: Qwen3-Omni
|
||||
# Reliability: Qwen3-Omni (H100 × 2)
|
||||
pytest -s -v dfx/reliability/test_reliability_qwen3_omni.py -m slow
|
||||
|
||||
# Reliability: Wan2.2
|
||||
# Reliability: Wan2.2 (H100)
|
||||
pytest -s -v dfx/reliability/test_reliability_wan22.py -m slow
|
||||
|
||||
# Reliability: HunyuanImage DiT (H100 × 4)
|
||||
pytest -s -v dfx/reliability/test_reliability_hunyuan_image.py -m slow
|
||||
|
||||
# Reliability: VoxCPM2 (L4)
|
||||
pytest -s -v dfx/reliability/test_reliability_voxcpm2.py -m slow
|
||||
|
||||
```
|
||||
|
||||
The latest L5 commands for CI can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-ready.yml).
|
||||
The latest L5 CI jobs (reliability + invalid-parameter weekly steps) are in [test-weekly.yml](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-weekly.yml).
|
||||
|
||||
You can find more information about markers in the documentation: [marker doc](./tests_markers.md)
|
||||
|
||||
|
||||
@@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr
|
||||
|
||||
### Pre-built wheels
|
||||
|
||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, and 0.23.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source).
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
@@ -94,6 +94,10 @@ Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc
|
||||
|
||||
### Build your own docker image
|
||||
|
||||
=== "NVIDIA CUDA"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/cuda.inc.md:build-docker"
|
||||
|
||||
=== "AMD ROCm"
|
||||
|
||||
--8<-- "docs/getting_started/installation/gpu/rocm.inc.md:build-docker"
|
||||
|
||||
@@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new**
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
uv pip install vllm==0.23.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]'
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
uv pip install vllm==0.23.0 --torch-backend=auto
|
||||
```
|
||||
|
||||
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.
|
||||
The 0.23.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -66,12 +66,12 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.22.0
|
||||
git checkout v0.23.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 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
|
||||
#For CUDA 13.0 (the default for v0.23.0; the wheel filename has no `+cu130` suffix)
|
||||
export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0-cp38-abi3-manylinux_2_28_x86_64.whl
|
||||
```
|
||||
Install vllm with command below (If you have no existing PyTorch).
|
||||
```bash
|
||||
@@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm/vllm-omni:v0.22.0 \
|
||||
vllm/vllm-omni:v0.23.0 \
|
||||
vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -110,3 +110,53 @@ docker run --runtime nvidia --gpus 2 \
|
||||
The CUDA image does not define a default entrypoint, so include `vllm serve ... --omni` after the image name.
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
|
||||
# --8<-- [start:build-docker]
|
||||
|
||||
#### Build docker image
|
||||
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build -f docker/Dockerfile.cuda -t vllm-omni-cuda .
|
||||
```
|
||||
|
||||
If you want to specify the base vLLM version:
|
||||
|
||||
```bash
|
||||
DOCKER_BUILDKIT=1 docker build \
|
||||
-f docker/Dockerfile.cuda \
|
||||
--build-arg BASE_IMAGE=vllm/vllm-openai:v0.22.1 \
|
||||
-t vllm-omni-cuda .
|
||||
```
|
||||
|
||||
#### Launch the docker image
|
||||
|
||||
##### Launch with OpenAI API Server
|
||||
|
||||
!!! note
|
||||
The model `Qwen/Qwen3-Omni-30B-A3B-Instruct` requires significant GPU memory. The example below has been verified on 2 x H100's.
|
||||
|
||||
```bash
|
||||
docker run --runtime nvidia --gpus 2 \
|
||||
-v ${HF_HOME:-$HOME/.cache/huggingface}:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
vllm-omni-cuda \
|
||||
vllm serve --omni --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091
|
||||
```
|
||||
|
||||
By default, this mounts `$HOME/.cache/huggingface` as the model cache directory. To use a custom location, set the `HF_HOME` environment variable before running the command (e.g., `export HF_HOME=/data/models`).
|
||||
|
||||
##### Launch with interactive session for development
|
||||
|
||||
```bash
|
||||
docker run --runtime nvidia --gpus all -it --rm \
|
||||
-v ${HF_HOME:-$HOME/.cache/huggingface}:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
--ipc=host \
|
||||
--entrypoint bash \
|
||||
vllm-omni-cuda
|
||||
```
|
||||
|
||||
# --8<-- [end:build-docker]
|
||||
@@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images.
|
||||
|
||||
vLLM-Omni is built based on vLLM. Please install it with command below.
|
||||
```bash
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
|
||||
```
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
@@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm
|
||||
# --8<-- [start:build-wheel-from-source]
|
||||
|
||||
#### Installation of vLLM
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.22.0 release version of the library
|
||||
If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library
|
||||
|
||||
```bash
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
|
||||
```
|
||||
|
||||
The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
The pre-built 0.23.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead.
|
||||
|
||||
#### Installation of vLLM-Omni
|
||||
Since vllm-omni is rapidly evolving, it's recommended to install it from source
|
||||
@@ -58,7 +58,7 @@ If you want to check, modify or debug with source code of vLLM, install the libr
|
||||
```bash
|
||||
git clone https://github.com/vllm-project/vllm.git
|
||||
cd vllm
|
||||
git checkout v0.22.0
|
||||
git checkout v0.23.0
|
||||
python3 -m pip install -r requirements/rocm.txt
|
||||
python3 setup.py develop
|
||||
```
|
||||
@@ -130,7 +130,7 @@ docker run --rm \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
-p 8091:8091 \
|
||||
vllm/vllm-omni-rocm:v0.22.0 \
|
||||
vllm/vllm-omni-rocm:v0.23.0 \
|
||||
--model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091
|
||||
```
|
||||
|
||||
@@ -149,7 +149,7 @@ docker run --rm -it \
|
||||
-v ~/.cache/huggingface:/root/.cache/huggingface \
|
||||
--env "HF_TOKEN=$HF_TOKEN" \
|
||||
--entrypoint bash \
|
||||
vllm/vllm-omni-rocm:v0.22.0
|
||||
vllm/vllm-omni-rocm:v0.23.0
|
||||
```
|
||||
|
||||
# --8<-- [end:pre-built-images]
|
||||
@@ -19,10 +19,10 @@ uv venv --python 3.12 --seed
|
||||
source .venv/bin/activate
|
||||
|
||||
# On CUDA
|
||||
uv pip install vllm==0.22.0 --torch-backend=auto
|
||||
uv pip install vllm==0.23.0 --torch-backend=auto
|
||||
|
||||
# On ROCm
|
||||
uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722
|
||||
uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722
|
||||
|
||||
git clone https://github.com/vllm-project/vllm-omni.git
|
||||
cd vllm-omni
|
||||
@@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst
|
||||
!!! note
|
||||
It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni.
|
||||
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.23.0` and vLLM Omni `0.23.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue.
|
||||
|
||||
## Offline Inference
|
||||
|
||||
|
||||
@@ -144,6 +144,7 @@ def scan_package(package_name: str = "vllm_omni") -> dict[str, list[str]]:
|
||||
"vllm_omni.entrypoints.async_diffusion",
|
||||
"vllm_omni.entrypoints.openai",
|
||||
"vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts",
|
||||
"vllm_omni.experimental", # optional serving deps not installed in docs build
|
||||
]
|
||||
if any(module_path.startswith(prefix) for prefix in excluded_prefixes):
|
||||
continue
|
||||
|
||||
@@ -6,6 +6,8 @@ vLLM-Omni supports unified multimodal comprehension and generation models across
|
||||
|
||||
If vLLM-Omni natively supports a model, its implementation can be found in <gh-file:vllm_omni/model_executor/models> and <gh-file:vllm_omni/diffusion/models>.
|
||||
|
||||
For deployment recipes (TTS and multimodal), see [recipes.vllm.ai](https://recipes.vllm.ai) and the in-repo [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes) directory.
|
||||
|
||||
## List of Supported Models
|
||||
|
||||
<style>
|
||||
@@ -22,6 +24,7 @@ th {
|
||||
| `MingFlashOmniForConditionalGeneration` + `MingImagePipeline` | Ming-flash-omni-2.0 (omni-speech + imagegen<sup>1</sup>) | `Jonathan1909/Ming-flash-omni-2.0` | ✅︎ | | | |
|
||||
| `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | |
|
||||
| `Gr00tN1d7Pipeline` | GR00T N1.7 | `nvidia/GR00T-N1.7-3B` | ✅︎ | | | |
|
||||
| `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
@@ -33,21 +36,23 @@ th {
|
||||
| `ZImagePipeline` | Z-Image | `Tongyi-MAI/Z-Image-Turbo` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanPipeline` | Wan2.1-T2V, Wan2.2-T2V, Wan2.2-TI2V | `Wan-AI/Wan2.1-T2V-1.3B-Diffusers`, `Wan-AI/Wan2.1-T2V-14B-Diffusers`, `Wan-AI/Wan2.2-T2V-A14B-Diffusers`, `Wan-AI/Wan2.2-TI2V-5B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanImageToVideoPipeline` | Wan2.2-I2V | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `WanSpeechToVideoPipeline` | Wan2.2-S2V | `Wan-AI/Wan2.2-S2V-14B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Cosmos3OmniDiffusersPipeline` | Cosmos3 T2I, T2V, I2V, V2V, T2V with sound, action policy | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22S2VPipeline` | 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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Wan22VACEPipeline` | Wan2.2-VACE | `Pyros13/Wan2.2-VACE-Fun-A14B-Diffusers` | ✅︎ | | | |
|
||||
| `LTX2Pipeline` | LTX-2-T2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoPipeline` | LTX-2-I2V | `Lightricks/LTX-2` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2TwoStagesPipeline` | LTX-2-T2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX2ImageToVideoTwoStagesPipeline` | LTX-2-I2V | `rootonchair/LTX-2-19b-distilled` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23Pipeline` | LTX-2.3-T2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `LTX23ImageToVideoPipeline` | LTX-2.3-I2V | `dg845/LTX-2.3-Diffusers` | ✅︎ | ✅︎ | | |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | | | |
|
||||
| `DreamZeroPipeline` | DreamZero-DROID | `GEAR-Dreams/DreamZero-DROID` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `HeliosPipeline`, `HeliosPyramidPipeline` | Helios | `BestWishYsh/Helios-Base`, `BestWishYsh/Helios-Mid`, `BestWishYsh/Helios-Distilled` | ✅︎ | ✅︎ | ✅︎ | |
|
||||
| `MagiHumanPipeline` | MagiHuman | `SII-GAIR/daVinci-MagiHuman-Base-1080p` | ✅︎ | ✅︎ | | |
|
||||
| `OvisImagePipeline` | Ovis-Image | `OvisAI/Ovis-Image` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `LongcatImagePipeline` | LongCat-Image | `meituan-longcat/LongCat-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `LongCatImageEditPipeline` | LongCat-Image-Edit | `meituan-longcat/LongCat-Image-Edit` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `StableDiffusionXLPipeline` | Stable-Diffusion-XL | `stabilityai/stable-diffusion-xl-base-1.0` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `StableDiffusion3Pipeline` | Stable-Diffusion-3 | `stabilityai/stable-diffusion-3.5-medium` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `CosyVoice3Model` | CosyVoice3 | `FunAudioLLM/Fun-CosyVoice3-0.5B-2512` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniVoiceModel` | OmniVoice | `k2-fsa/OmniVoice` | ✅︎ | | | |
|
||||
@@ -59,16 +64,20 @@ th {
|
||||
| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ |
|
||||
| `SoulXSingerPipeline` | SoulX-Singer (SVS) | `Soul-AILab/SoulX-Singer` | ✅︎ | | | |
|
||||
| `SoulXSingerSVCPipeline` | SoulX-Singer-SVC | `Soul-AILab/SoulX-Singer` (`model-svc.pt`) | ✅︎ | | | |
|
||||
| `AudioXPipeline` | AudioX | `zhangj1an/AudioX` | ✅︎ | ✅︎ | | |
|
||||
| `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` | ✅︎ | ✅︎ | ✅︎ | ✅︎ |
|
||||
| `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-Base | `Qwen/Qwen3-TTS-12Hz-1.7B-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` | ✅︎ | | | |
|
||||
| `HiggsMultimodalQwen3ForConditionalGeneration` | Higgs-Audio v3 (TTS) | `bosonai/higgs-audio-v3-tts-4b` | ✅︎ | | | |
|
||||
| `IndexTTS2TalkerForConditionalGeneration` | IndexTTS-2 | `IndexTeam/IndexTTS-2` | ✅︎ | | | |
|
||||
| `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` | ✅︎ | ✅︎ | | |
|
||||
|
||||
@@ -32,6 +32,20 @@ curl -X POST http://localhost:8000/v1/images/generations \
|
||||
}' | jq -r '.data[0].b64_json' | base64 -d > dragon.png
|
||||
```
|
||||
|
||||
**Using curl save to file:**
|
||||
|
||||
```bash
|
||||
curl -o dragon.png -X POST http://localhost:8000/v1/images/generations \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"prompt": "a dragon laying over the spine of the Green Mountains of Vermont",
|
||||
"size": "1024x1024",
|
||||
"seed": 42,
|
||||
"response_format":"file"
|
||||
}'
|
||||
```
|
||||
|
||||
|
||||
**Using Python:**
|
||||
|
||||
```python
|
||||
@@ -57,6 +71,36 @@ img = Image.open(io.BytesIO(img_bytes))
|
||||
img.save("cat.png")
|
||||
```
|
||||
|
||||
**Using Python save to file:**
|
||||
|
||||
```python
|
||||
import requests
|
||||
import base64
|
||||
from PIL import Image
|
||||
import io
|
||||
import re
|
||||
|
||||
response = requests.post(
|
||||
"http://localhost:8000/v1/images/generations",
|
||||
json={
|
||||
"prompt": "a black and white cat wearing a princess tiara",
|
||||
"size": "1024x1024",
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42,
|
||||
"response_format":"file"
|
||||
}
|
||||
)
|
||||
|
||||
# save to file
|
||||
content_disposition = response.headers.get("Content-Disposition", "")
|
||||
match = re.search(r'filename="?(.+)"?', content_disposition)
|
||||
filename = match.group(1) if match else "save.png"
|
||||
with open(filename, "wb") as f:
|
||||
for chunk in response.iter_content(8192):
|
||||
f.write(chunk)
|
||||
print("saved:", filename)
|
||||
```
|
||||
|
||||
**Using OpenAI SDK:**
|
||||
|
||||
```python
|
||||
|
||||
@@ -9,6 +9,13 @@ vLLM-Omni provides an OpenAI-compatible API for text-to-speech (TTS) generation.
|
||||
|
||||
See the [Supported Models](#supported-models) section below for the full list, including OmniVoice, VoxCPM2, and MOSS-TTS-Nano.
|
||||
|
||||
!!! tip "Deployment recipes"
|
||||
TTS deployment recipes are published at
|
||||
[recipes.vllm.ai](https://recipes.vllm.ai) (e.g.
|
||||
[Qwen3-TTS](https://recipes.vllm.ai/Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice),
|
||||
[Higgs-Audio v3](https://recipes.vllm.ai/bosonai/higgs-audio-v3-tts-4b)).
|
||||
The in-repo runbooks live under [`recipes/`](https://github.com/vllm-project/vllm-omni/tree/main/recipes).
|
||||
|
||||
Each server instance runs a single model (specified at startup via `vllm serve <model> --omni`).
|
||||
|
||||
## Quick Start
|
||||
@@ -226,9 +233,9 @@ Server -> Client:
|
||||
|
||||
| Message | Description |
|
||||
|---------|-------------|
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for a sentence |
|
||||
| `{"type": "audio.start", "sentence_index": 0, "sentence_text": "...", "format": "pcm", "sample_rate": 24000}` | Audio generation starting for the buffered input |
|
||||
| Binary frame | Raw audio bytes (one or more PCM chunks when `stream_audio=true`) |
|
||||
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for a sentence |
|
||||
| `{"type": "audio.done", "sentence_index": 0, "total_bytes": 96000, "error": false}` | Audio complete for the buffered input |
|
||||
| `{"type": "session.done", "total_sentences": N}` | Session complete |
|
||||
| `{"type": "error", "message": "..."}` | Non-fatal error |
|
||||
|
||||
@@ -238,8 +245,7 @@ All REST API parameters are supported, plus:
|
||||
|
||||
| Parameter | Type | Default | Description |
|
||||
|-----------|------|---------|-------------|
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks per sentence over WebSocket |
|
||||
| `split_granularity` | string | "sentence" | Text splitting granularity |
|
||||
| `stream_audio` | bool | false | Stream one or more PCM chunks for the buffered input over WebSocket |
|
||||
|
||||
|
||||
```bash
|
||||
|
||||
@@ -160,7 +160,12 @@ 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.
|
||||
decoded reference frames are used as clean conditioning. Cosmos3 transfer mode
|
||||
also accepts `edge`, `blur`, `depth`, `seg`, or `wsm` control hints plus
|
||||
transfer options such as `control_path`, `control_guidance`,
|
||||
`control_guidance_interval`, `num_video_frames_per_chunk`,
|
||||
`num_conditional_frames`, `show_control_condition`, and `show_input`; see the
|
||||
Cosmos3 recipe for complete examples.
|
||||
|
||||
### Speech-to-Video
|
||||
|
||||
@@ -195,13 +200,16 @@ curl -X POST http://localhost:8091/v1/videos/sync \
|
||||
|
||||
## Storage
|
||||
|
||||
Set `VLLM_OMNI_STORAGE_PATH` to control where asynchronous video outputs are
|
||||
Set `VLLM_OMNI_SERVER_STORAGE__PATH` to control where asynchronous video outputs are
|
||||
stored:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
|
||||
```
|
||||
|
||||
> `VLLM_OMNI_STORAGE_PATH` is deprecated and will be removed in a future release;
|
||||
> use `VLLM_OMNI_SERVER_STORAGE__PATH` instead.
|
||||
|
||||
## Model-Specific Examples
|
||||
|
||||
For complete text-to-video, image-to-video, and model-specific video-to-video
|
||||
|
||||
@@ -193,9 +193,11 @@ Factory function `get_offload_backend()` selects the appropriate backend based o
|
||||
| NextStep11Pipeline | `stepfun-ai/NextStep-1.1` | `NextStepModel` | - | ✓ | `"layers"` |
|
||||
| OvisImagePipeline | `AIDC-AI/Ovis-Image-7B` | `OvisImageTransformer2DModel` | - | ✓ | `"transformer"` |
|
||||
| QwenImagePipeline | `Qwen/Qwen-Image` | `QwenImageTransformer2DModel` | ✓ | ✓ | `"transformer_blocks"` |
|
||||
| StableDiffusionXLPipeline | `stabilityai/stable-diffusion-xl-base-1.0` | `SDXLUNet2DConditionModel` | ✓ | ✓ | `"down_blocks"`, `"up_blocks"` |
|
||||
| StableDiffusion3Pipeline | `stabilityai/stable-diffusion-3.5-medium` | `SD3Transformer2DModel` | - | ✓ | `"transformer_blocks"` |
|
||||
| Wan22I2VPipeline | `Wan-AI/Wan2.2-I2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| Wan22Pipeline | `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | `WanTransformer3DModel` | ✓ | ✓ | `"blocks"` |
|
||||
| SoulXSingerPipeline / SoulXSingerSVCPipeline | `Soul-AILab/SoulX-Singer` | `DiffLlama` (`cfm_decoder.model.diff_estimator`) | ✓ | ✓ | `"layers"` |
|
||||
| BagelPipeline | `ByteDance-Seed/BAGEL-7B-MoT` | `Qwen2MoTModel` | - | ✓ | `"layers"`, `"customized modules"` |
|
||||
|
||||
**Notes:**
|
||||
|
||||
@@ -0,0 +1,50 @@
|
||||
This directory contains auto-tuned Triton kernel configurations for the
|
||||
MoT (Mixture-of-Tokens) GEMM and RMSNorm operators used by BAGEL and other
|
||||
MoT-architecture diffusion models.
|
||||
|
||||
File naming convention:
|
||||
device_name=<GPU_NAME>,dtype=<DTYPE>.json
|
||||
|
||||
For example:
|
||||
device_name=NVIDIA_A100-SXM4-80GB,dtype=w16a16.json
|
||||
|
||||
Each JSON file maps (K, N) matrix shapes to a dictionary of batch sizes (M)
|
||||
and their optimal Triton tile configurations:
|
||||
|
||||
{
|
||||
"3584_4608": { // K=3584, N=4608 (QKV projection)
|
||||
"1024": {
|
||||
"BLOCK_SIZE_M": 64,
|
||||
"BLOCK_SIZE_N": 128,
|
||||
"BLOCK_SIZE_K": 64,
|
||||
"GROUP_SIZE_M": 8,
|
||||
"num_warps": 4,
|
||||
"num_stages": 3
|
||||
},
|
||||
...
|
||||
},
|
||||
"3584_3584": { ... } // K=3584, N=3584 (output projection)
|
||||
}
|
||||
|
||||
Config loading order (3-tier, see ops/mot_gemm.py):
|
||||
1. $VLLM_TUNED_CONFIG_FOLDER/<filename> (env override)
|
||||
2. This directory: vllm_omni/diffusion/layers/mot/configs/<filename>
|
||||
3. Conservative default config (compiles everywhere, sub-optimal perf)
|
||||
|
||||
If no config file matches the current device, a warning is printed with
|
||||
instructions to run the auto-tuning benchmark.
|
||||
|
||||
To generate configs for your hardware:
|
||||
|
||||
python benchmarks/kernels/mot_linear_benchmarks.py \
|
||||
--model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--tp-size 1 --dtype w16a16 --tune \
|
||||
--save-dir vllm_omni/diffusion/layers/mot/configs/
|
||||
|
||||
For multi-GPU tuning (uses Ray for parallel search):
|
||||
|
||||
python benchmarks/kernels/mot_linear_benchmarks.py \
|
||||
--model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--tp-size 2 --tune
|
||||
|
||||
See benchmarks/kernels/mot_linear_benchmarks.py for full options.
|
||||
@@ -130,6 +130,7 @@ The following tables show which models support each feature:
|
||||
| **Qwen-Image-Edit-2509** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ (decode) | ✅ | ❌ | ❌ |
|
||||
| **Qwen-Image-Layered** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **SenseNova-U1** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **Stable-Diffusion-XL** | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Stable-Diffusion3.5** | ❌ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **Z-Image** | ✅ | ✅ | ✅ | ❓ | ✅ (TP=2 only) | ❌ | ✅ | ❌ | ✅ (decode) | ✅ | ❌ |
|
||||
| **ERNIE-Image** | ❌ | ✅ | ✅ | ❓ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
@@ -145,7 +146,7 @@ The following tables show which models support each feature:
|
||||
| Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | Pipeline-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution |
|
||||
|------------------------------|:---------:|:----------:|:---------------------:|:--------------:|:-----------------:|:-----------------:|:------:|:-------------------------:|:--------------------:|:--------------:|:----------------:|
|
||||
| **Wan2.2** | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.2-S2V** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (encode/decode) | ❌ | ❌ |
|
||||
| **Wan2.1-VACE** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
| **LTX-2** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ |
|
||||
| **LTX-2.3** | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ (decode) | ❌ | ❌ |
|
||||
|
||||
@@ -1,203 +0,0 @@
|
||||
# Helios Video Generation
|
||||
|
||||
Source <https://github.com/vllm-project/vllm-omni/tree/main/examples/offline_inference/helios>.
|
||||
|
||||
|
||||
Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
|
||||
|
||||
| Variant | Description | Key Features |
|
||||
|---------|-------------|--------------|
|
||||
| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
|
||||
| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
|
||||
| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/helios
|
||||
```
|
||||
|
||||
### Text-to-Video (T2V)
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_i2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_i2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_i2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Video-to-Video (V2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_v2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_v2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_v2v_distilled.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
|
||||
| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
|
||||
| `--prompt` | — | Text prompt describing the video |
|
||||
| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
|
||||
| `--image-path` | — | Input image (required for `i2v`) |
|
||||
| `--video-path` | — | Input video (required for `v2v`) |
|
||||
| `--height` | `384` | Video height in pixels |
|
||||
| `--width` | `640` | Video width in pixels |
|
||||
| `--num-frames` | `99` | Number of output frames |
|
||||
| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
|
||||
| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--fps` | `16` | Output video frame rate |
|
||||
| `--output` | `helios_output.mp4` | Output file path |
|
||||
|
||||
### Stage 2 / Pyramid Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
|
||||
| `--pyramid-num-stages` | `3` | Number of pyramid stages |
|
||||
| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
|
||||
| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
|
||||
|
||||
### CFG-Zero* Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
|
||||
| `--use-zero-init` | off | Zero init for first steps |
|
||||
| `--zero-steps` | `1` | Number of zero-init steps |
|
||||
|
||||
### Memory & Parallelism
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--vae-use-slicing` | off | Enable VAE slicing |
|
||||
| `--vae-use-tiling` | off | Enable VAE tiling |
|
||||
| `--enforce-eager` | off | Disable torch.compile |
|
||||
| `--enable-cpu-offload` | off | CPU offloading |
|
||||
| `--enable-layerwise-offload` | off | Layerwise offloading |
|
||||
| `--ulysses-degree` | `1` | Ulysses SP degree |
|
||||
| `--ring-degree` | `1` | Ring SP degree |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
|
||||
| `--tensor-parallel-size` | `1` | Tensor parallelism size |
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/helios/end2end.py"
|
||||
``````
|
||||
@@ -0,0 +1,65 @@
|
||||
# AURA Omni Pipeline
|
||||
|
||||
`aura_omni` wires ASR, AURA, and Qwen3-TTS into one vLLM-Omni pipeline:
|
||||
|
||||
```text
|
||||
ASR -> AURA -> Qwen3-TTS Talker -> Code2Wav
|
||||
```
|
||||
|
||||
Qwen3-TTS remains two engine stages so the pipeline reuses the existing native
|
||||
Talker and Code2Wav implementation.
|
||||
|
||||
```bash
|
||||
vllm serve aurateam/AURA \
|
||||
--omni \
|
||||
--deploy-config vllm_omni/deploy/aura_omni.yaml \
|
||||
--served-model-name aurateam/AURA \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
Configure local checkpoints by editing per-stage `model` values in
|
||||
`vllm_omni/deploy/aura_omni.yaml`. The deploy file sets
|
||||
`pipeline: aura_omni`, so the four-stage topology is used even if the
|
||||
command-line `--model` points at one of the component checkpoints.
|
||||
|
||||
Send requests with `"model": "aurateam/AURA"`. The ASR, AURA, and Qwen3-TTS
|
||||
checkpoint paths are internal stage models from the deploy YAML, not the
|
||||
OpenAI-facing served model name.
|
||||
|
||||
The AURA stage can emit `<|silent|>`. Silent outputs are treated as a gate:
|
||||
they produce no Qwen3-TTS Talker input, so no audio is synthesized for that
|
||||
turn.
|
||||
|
||||
## GPU Utilization Recommendation
|
||||
|
||||
`gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` controls how much
|
||||
VRAM each stage can reserve. Start with this split for a single GPU:
|
||||
|
||||
- Stage 0 (ASR): `0.10`
|
||||
- Stage 1 (AURA): `0.40`
|
||||
- Stage 2 (Qwen3-TTS Talker): `0.20`
|
||||
- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
|
||||
|
||||
## TTS Modes
|
||||
|
||||
`aura_omni` can pass AURA text to Qwen3-TTS in two task modes:
|
||||
|
||||
- `Base`: voice clone from `tts_ref_audio` with ICL enabled in the AURA
|
||||
pipeline. Provide both `tts_ref_audio` and `tts_ref_text`. Set
|
||||
`tts_x_vector_only_mode=true` to disable ICL and use speaker embedding only.
|
||||
- `CustomVoice`: predefined speaker mode. Use a Qwen3-TTS CustomVoice
|
||||
checkpoint for stages 2 and 3 in `aura_omni.yaml`, then pass
|
||||
`tts_task_type=CustomVoice` and `tts_speaker`.
|
||||
|
||||
By default, AURA responses are passed to Qwen3-TTS as text. Set
|
||||
`additional_information.tts_pass_token_ids=true` to pass AURA-generated
|
||||
assistant token ids directly instead. Even when token passthrough is disabled,
|
||||
the stage processor uses AURA token ids when available to estimate the Talker
|
||||
prefill length, so scheduling does not rely on raw character length.
|
||||
|
||||
The example client exposes this as:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-pass-token-ids
|
||||
```
|
||||
@@ -19,7 +19,7 @@ both stages on a single GPU, override per stage:
|
||||
|
||||
```bash
|
||||
vllm serve zai-org/GLM-Image --omni --port 8091 \
|
||||
--stage-0-devices 0 --stage-1-devices 0
|
||||
--stage-overrides '{"0": {"devices": "0"}, "1": {"devices": "0"}}'
|
||||
```
|
||||
|
||||
## API Calls
|
||||
|
||||
@@ -84,16 +84,26 @@ curl -X POST http://localhost:8091/v1/videos/sync \
|
||||
Generated video files are stored on local disk by the async video API.
|
||||
Local file storage behavior can be controlled via the following environment variables:
|
||||
|
||||
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
|
||||
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
|
||||
- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files
|
||||
(default: `/tmp/storage`)
|
||||
- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open
|
||||
operations (default: `4`)
|
||||
- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in
|
||||
seconds
|
||||
- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in
|
||||
seconds for enforcing file TTL; defaults to `300` when TTL is set
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
|
||||
export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
|
||||
export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
|
||||
export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400
|
||||
export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300
|
||||
```
|
||||
|
||||
See also: [Configuration Options](../../../configuration/README.md)
|
||||
|
||||
## API Calls
|
||||
|
||||
### Method 1: Using curl
|
||||
|
||||
@@ -90,16 +90,26 @@ curl -X POST http://localhost:8091/v1/videos/sync \
|
||||
Generated video files are stored on local disk by the async video API.
|
||||
Local file storage behavior can be controlled via the following environment variables:
|
||||
|
||||
- `VLLM_OMNI_STORAGE_PATH`: directory used for generated files (default: `/tmp/storage`)
|
||||
- `VLLM_OMNI_STORAGE_MAX_CONCURRENCY`: max concurrent save/delete operations (default: `4`)
|
||||
- `VLLM_OMNI_SERVER_STORAGE__PATH`: directory used for generated files
|
||||
(default: `/tmp/storage`)
|
||||
- `VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY`: max concurrent save/delete/open
|
||||
operations (default: `4`)
|
||||
- `VLLM_OMNI_SERVER_STORAGE__FILE_TTL`: optional TTL for generated files in
|
||||
seconds
|
||||
- `VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL`: optional sweep interval in
|
||||
seconds for enforcing file TTL; defaults to `300` when TTL is set
|
||||
|
||||
Example:
|
||||
|
||||
```bash
|
||||
export VLLM_OMNI_STORAGE_PATH=/var/tmp/vllm-omni-videos
|
||||
export VLLM_OMNI_STORAGE_MAX_CONCURRENCY=8
|
||||
export VLLM_OMNI_SERVER_STORAGE__PATH=/var/tmp/vllm-omni-videos
|
||||
export VLLM_OMNI_SERVER_STORAGE__FILE_CONCURRENCY=8
|
||||
export VLLM_OMNI_SERVER_STORAGE__FILE_TTL=86400
|
||||
export VLLM_OMNI_SERVER_STORAGE__TTL_SWEEP_INTERVAL=300
|
||||
```
|
||||
|
||||
See also: [Configuration Options](../../../configuration/README.md)
|
||||
|
||||
## API Calls
|
||||
|
||||
### Method 1: Using curl
|
||||
|
||||
@@ -65,7 +65,7 @@ More failure modes (for example, network jitter and network interruption) will b
|
||||
|
||||
## Source of Conclusions
|
||||
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on `Qwen3-Omni` and `Wan2.2`.
|
||||
The behaviors and conclusions above are summarized from current fault injection validation results on **`Qwen3-Omni`**, **`Wan2.2`**, **`HunyuanImage-3.0-Instruct`** (DiT-only, `/v1/images/generations`), and **`VoxCPM2`** (`/v1/audio/speech`). Automated coverage lives under `tests/dfx/reliability/` (`test_reliability_qwen3_omni.py`, `test_reliability_wan22.py`, `test_reliability_hunyuan_image.py`, `test_reliability_voxcpm2.py`) and runs weekly via `.buildkite/test-weekly.yml`.
|
||||
|
||||
### Example: `SIGTERM` fault injection log (Qwen3-Omni)
|
||||
|
||||
|
||||
@@ -85,6 +85,7 @@ warmup_quack_fp8([(14040, 2048, 6144), (14040, 2048, 2048)])
|
||||
| FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B` | Yes | Yes | All layers | None | |
|
||||
| HunyuanImage-3.0 | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | Yes | Yes | All layers; use the Hunyuan stage config for multi-stage runs | None | |
|
||||
| HunyuanVideo-1.5 | `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v`, `720p_t2v`, `480p_i2v` | Yes | Yes | All layers | None | |
|
||||
| Cosmos3 | `nvidia/Cosmos3-Nano`, `nvidia/Cosmos3-Super` | Yes | Not validated | All layers | None | |
|
||||
|
||||
### Multi-Stage Omni/TTS Model (Qwen3-Omni, Qwen3-TTS)
|
||||
|
||||
|
||||
@@ -89,11 +89,54 @@ and
|
||||
| Model | Scope | Status |
|
||||
|-------|-------|--------|
|
||||
| Qwen3-Omni | Thinker language-model stage | ModelOpt FP8 checkpoint path |
|
||||
| Qwen3-Omni | Thinker language-model stage (W4A4 NVFP4) | Validated; see [Qwen3-Omni NVFP4 W4A4](#qwen3-omni-nvfp4-w4a4-thinker) below |
|
||||
| Qwen3-TTS | TTS language-model stage | Not validated |
|
||||
|
||||
Audio encoder, vision encoder, talker, and code2wav stages stay in BF16 unless
|
||||
a model-specific guide documents otherwise.
|
||||
|
||||
#### Qwen3-Omni NVFP4 W4A4 (thinker)
|
||||
|
||||
vLLM-Omni serves ModelOpt NVFP4 W4A4 quantizations of the
|
||||
Qwen3-Omni-30B-A3B-Instruct thinker language model. The thinker text body
|
||||
(attention + MoE experts) is quantized to NVFP4 with FP8 per-tensor input
|
||||
scales; the audio encoder, vision encoder, talker, and code2wav stay in BF16.
|
||||
|
||||
| Variant | HF checkpoint | Hardware |
|
||||
|---------|---------------|----------|
|
||||
| W4A4 NVFP4 (full thinker) | `YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip` | sm_100+ (Blackwell, FlashInfer FP4 GEMM) |
|
||||
|
||||
Calibration uses ModelOpt `mtq.NVFP4_DEFAULT_CFG` with the `awq_clip` algorithm
|
||||
on 1024 ultrachat samples chat-templated through the Qwen3-Omni tokenizer.
|
||||
Excluded modules: `*audio_tower*`, `*visual*`, `*talker*`, `*code2wav*`,
|
||||
`*lm_head*`, `*mlp.gate*`. See `scripts/nvfp4/calibrate.py` for the reference
|
||||
recipe.
|
||||
|
||||
!!! note "ModelOpt 0.44 NaN regression workaround"
|
||||
ModelOpt 0.44's float32 -> FP8 E4M3 cast of per-block weight scales
|
||||
occasionally emits literal NaN bytes (E4M3 encoding 0x7F / 0xFF) for
|
||||
blocks whose pre-cast scale rounds above the FP8 max of 448 after the
|
||||
global-scale division. A single NaN byte in any `weight_scale`
|
||||
propagates through the FP4 GEMM and collapses the served model output
|
||||
to `!!!!`. vLLM-Omni's `vllm_omni.patch` installs a defensive override
|
||||
of `ModelOptNvFp4LinearMethod.process_weights_after_loading` that
|
||||
clamps these bytes to the FP8 E4M3 max at load time. The override
|
||||
self-extinguishes once vllm-omni's vllm pin moves to a release
|
||||
containing the upstream fix. Set `VLLM_OMNI_SKIP_NVFP4_NAN_CLAMP=1`
|
||||
to disable the override for diagnostics.
|
||||
|
||||
Serving:
|
||||
|
||||
```bash
|
||||
vllm serve YihongJin/Qwen3-Omni-30B-A3B-Instruct-NVFP4-W4A4-full-thinker-awqclip \
|
||||
--omni --port 8000
|
||||
```
|
||||
|
||||
> Do **not** pass `--enforce-eager` for production / benchmarks. CUDA graphs
|
||||
> amortize launch overhead and unlock the FP4 throughput wins; with
|
||||
> `--enforce-eager` set, W4A4 TPOT degrades ~10x relative to the CUDA-graph
|
||||
> configuration.
|
||||
|
||||
### Multi-Stage Diffusion Model
|
||||
|
||||
ModelOpt checkpoints must be routed to the stage whose checkpoint contains the
|
||||
|
||||
@@ -0,0 +1,61 @@
|
||||
# AURA Omni
|
||||
|
||||
This example runs the native AURA Omni pipeline offline:
|
||||
|
||||
```text
|
||||
Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav
|
||||
```
|
||||
|
||||
The first stage consumes speech and produces a transcript. The AURA stage then combines the transcript with the video frames from the original request and returns text or `<|silent|>`. Non-silent responses are passed to Qwen3-TTS as AURA-generated token ids.
|
||||
|
||||
## Run
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/aura_omni
|
||||
bash run_single_prompt.sh
|
||||
```
|
||||
|
||||
Use local media:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--audio-path /path/to/input.wav \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--modalities text,audio
|
||||
```
|
||||
|
||||
Base voice clone mode is the default. It uses the AURA reference audio unless
|
||||
you override it:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--tts-task-type Base \
|
||||
--tts-ref-audio /data/yrr/rein_test/shuhan.mp3 \
|
||||
--tts-x-vector-only-mode
|
||||
```
|
||||
|
||||
CustomVoice mode requires stages 2 and 3 in `vllm_omni/deploy/aura_omni.yaml`
|
||||
to point at a Qwen3-TTS CustomVoice checkpoint:
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--tts-task-type CustomVoice \
|
||||
--tts-speaker Vivian
|
||||
```
|
||||
|
||||
For local checkpoints, edit the stage `model` entries in
|
||||
`vllm_omni/deploy/aura_omni.yaml` or pass a copied deploy config with
|
||||
`--deploy-config`.
|
||||
|
||||
## GPU Utilization Recommendation
|
||||
|
||||
Set `gpu_memory_utilization` in `vllm_omni/deploy/aura_omni.yaml` by stage.
|
||||
Suggested starting point:
|
||||
|
||||
- Stage 0 (ASR): `0.10`
|
||||
- Stage 1 (AURA): `0.4`
|
||||
- Stage 2 (Qwen3-TTS Talker): `0.20`
|
||||
- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
|
||||
|
||||
Generated text and audio are written to `--output-dir`
|
||||
(default: `output_aura_omni`).
|
||||
@@ -0,0 +1,236 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline AURA Omni example.
|
||||
|
||||
The native AURA pipeline is:
|
||||
|
||||
Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav
|
||||
|
||||
Stage 0 consumes audio and emits a transcript. The stage input processor then
|
||||
combines that transcript with the original visual payload for AURA.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from typing import NamedTuple
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import vllm
|
||||
from vllm import SamplingParams
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.assets.video import VideoAsset, video_to_ndarrays
|
||||
from vllm.multimodal.media.audio import load_audio
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
default_qwen3_tts_ref_audio_path,
|
||||
)
|
||||
|
||||
SEED = 42
|
||||
DEFAULT_MODEL = "aurateam/AURA"
|
||||
DEFAULT_DEPLOY_CONFIG = "vllm_omni/deploy/aura_omni.yaml"
|
||||
|
||||
|
||||
class QueryResult(NamedTuple):
|
||||
inputs: dict
|
||||
limit_mm_per_prompt: dict[str, int]
|
||||
|
||||
|
||||
def _load_audio(audio_path: str | None, sampling_rate: int) -> tuple[np.ndarray, int]:
|
||||
if audio_path:
|
||||
if not os.path.exists(audio_path):
|
||||
raise FileNotFoundError(f"Audio file not found: {audio_path}")
|
||||
audio_signal, sr = load_audio(audio_path, sr=sampling_rate)
|
||||
return audio_signal.astype(np.float32), int(sr)
|
||||
audio_signal, sr = AudioAsset("mary_had_lamb").audio_and_sample_rate
|
||||
return audio_signal.astype(np.float32), int(sr)
|
||||
|
||||
|
||||
def _load_video(video_path: str | None, num_frames: int) -> list[np.ndarray]:
|
||||
if video_path:
|
||||
if not os.path.exists(video_path):
|
||||
raise FileNotFoundError(f"Video file not found: {video_path}")
|
||||
return video_to_ndarrays(video_path, num_frames=num_frames)
|
||||
return VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays
|
||||
|
||||
|
||||
def build_query(args) -> QueryResult:
|
||||
prompt = (
|
||||
"<|im_start|>user\n"
|
||||
"<|audio_start|><|audio_pad|><|audio_end|>"
|
||||
"Please transcribe this speech and use the video context to decide whether a reply is needed.<|im_end|>\n"
|
||||
"<|im_start|>assistant\n"
|
||||
)
|
||||
audio_data = _load_audio(args.audio_path, args.sampling_rate)
|
||||
video_frames = _load_video(args.video_path, args.num_frames)
|
||||
return QueryResult(
|
||||
inputs={
|
||||
"prompt": prompt,
|
||||
"multi_modal_data": {
|
||||
"audio": audio_data,
|
||||
"video": video_frames,
|
||||
},
|
||||
"additional_information": {
|
||||
"aura_system_prompt": args.aura_system_prompt,
|
||||
"tts_task_type": args.tts_task_type,
|
||||
"tts_language": args.tts_language,
|
||||
"tts_speaker": args.tts_speaker,
|
||||
"tts_instruct": args.tts_instruct,
|
||||
"tts_ref_audio": args.tts_ref_audio,
|
||||
"tts_ref_text": args.tts_ref_text,
|
||||
"tts_x_vector_only_mode": args.tts_x_vector_only_mode,
|
||||
"tts_pass_token_ids": args.tts_pass_token_ids,
|
||||
},
|
||||
},
|
||||
limit_mm_per_prompt={"audio": 1, "video": 1},
|
||||
)
|
||||
|
||||
|
||||
def build_sampling_params() -> list[SamplingParams]:
|
||||
asr_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=256,
|
||||
detokenize=True,
|
||||
seed=SEED,
|
||||
)
|
||||
aura_sampling = SamplingParams(
|
||||
temperature=0.5,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=256,
|
||||
detokenize=True,
|
||||
repetition_penalty=1.0,
|
||||
seed=SEED,
|
||||
)
|
||||
talker_sampling = SamplingParams(
|
||||
temperature=0.9,
|
||||
top_k=50,
|
||||
max_tokens=4096,
|
||||
detokenize=False,
|
||||
repetition_penalty=1.05,
|
||||
stop_token_ids=[2150],
|
||||
seed=SEED,
|
||||
)
|
||||
code2wav_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
top_p=1.0,
|
||||
top_k=-1,
|
||||
max_tokens=65536,
|
||||
detokenize=True,
|
||||
repetition_penalty=1.0,
|
||||
seed=SEED,
|
||||
)
|
||||
return [asr_sampling, aura_sampling, talker_sampling, code2wav_sampling]
|
||||
|
||||
|
||||
def save_stage_output(stage_outputs, output_dir: str) -> None:
|
||||
output = stage_outputs.request_output
|
||||
request_id = output.request_id
|
||||
if stage_outputs.final_output_type == "text":
|
||||
text_output = output.outputs[0].text
|
||||
out_txt = os.path.join(output_dir, f"{request_id}.txt")
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
f.write(str(text_output).strip() + "\n")
|
||||
print(f"Request ID: {request_id}, text saved to {out_txt}")
|
||||
print(text_output)
|
||||
return
|
||||
|
||||
if stage_outputs.final_output_type == "audio":
|
||||
mm = output.outputs[0].multimodal_output
|
||||
if not mm or "audio" not in mm:
|
||||
print(f"Request ID: {request_id}, no audio output")
|
||||
return
|
||||
audio = mm["audio"]
|
||||
if isinstance(audio, list):
|
||||
import torch
|
||||
|
||||
audio = torch.cat([(item if hasattr(item, "detach") else torch.tensor(item)).flatten() for item in audio])
|
||||
sample_rate = int(mm.get("sr", 24000))
|
||||
audio_numpy = audio.float().detach().cpu().numpy().flatten()
|
||||
out_wav = os.path.join(output_dir, f"{request_id}.wav")
|
||||
sf.write(out_wav, audio_numpy, samplerate=sample_rate, format="WAV")
|
||||
print(f"Request ID: {request_id}, audio saved to {out_wav}")
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20)
|
||||
os.makedirs(args.output_dir, exist_ok=True)
|
||||
|
||||
query = build_query(args)
|
||||
prompts = [query.inputs for _ in range(args.num_prompts)]
|
||||
if args.modalities:
|
||||
modalities = [item.strip() for item in args.modalities.split(",") if item.strip()]
|
||||
for prompt in prompts:
|
||||
prompt["modalities"] = modalities
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
omni_kwargs["model"] = args.model
|
||||
omni = Omni(**omni_kwargs)
|
||||
sampling_params_list = build_sampling_params()[: omni.num_stages]
|
||||
|
||||
try:
|
||||
for stage_outputs in omni.generate(prompts, sampling_params_list):
|
||||
save_stage_output(stage_outputs, args.output_dir)
|
||||
finally:
|
||||
omni.close()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="AURA Omni offline inference example")
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL, help="Top-level model name/path.")
|
||||
parser.add_argument("--deploy-config", default=DEFAULT_DEPLOY_CONFIG, help="AURA Omni deploy config.")
|
||||
parser.add_argument("--audio-path", "-a", default=None, help="Local audio file. Uses a built-in asset by default.")
|
||||
parser.add_argument("--video-path", "-v", default=None, help="Local video file. Uses a built-in asset by default.")
|
||||
parser.add_argument("--num-frames", type=int, default=16, help="Number of video frames to sample.")
|
||||
parser.add_argument("--sampling-rate", type=int, default=16000, help="Audio sampling rate for ASR input.")
|
||||
parser.add_argument("--num-prompts", type=int, default=1, help="Number of repeated prompts.")
|
||||
parser.add_argument("--modalities", default="text,audio", help="Comma-separated final modalities.")
|
||||
parser.add_argument("--output-dir", default="output_aura_omni", help="Directory for generated files.")
|
||||
parser.add_argument(
|
||||
"--aura-system-prompt",
|
||||
default=(
|
||||
"You are receiving a live video stream where the final frame is the present moment. "
|
||||
"Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in Chinese."
|
||||
),
|
||||
help="System prompt for the AURA stage.",
|
||||
)
|
||||
parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"], help="Qwen3-TTS task.")
|
||||
parser.add_argument("--tts-language", default="Chinese", help="Qwen3-TTS language.")
|
||||
parser.add_argument("--tts-speaker", default="Vivian", help="Qwen3-TTS speaker.")
|
||||
parser.add_argument("--tts-instruct", default="", help="Optional Qwen3-TTS style instruction.")
|
||||
parser.add_argument(
|
||||
"--tts-ref-audio",
|
||||
default=default_qwen3_tts_ref_audio_path(),
|
||||
help="Base-mode reference audio path/URL.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-ref-text",
|
||||
default=DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
help="Base-mode reference audio transcript.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-x-vector-only-mode",
|
||||
action="store_true",
|
||||
help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-pass-token-ids",
|
||||
action="store_true",
|
||||
help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.",
|
||||
)
|
||||
parser.add_argument("--init-timeout", type=int, default=2000, help="Overall initialization timeout.")
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=2000, help="Per-stage initialization timeout.")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable stats logging.")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,8 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
python end2end.py \
|
||||
--model aurateam/AURA \
|
||||
--deploy-config vllm_omni/deploy/aura_omni.yaml \
|
||||
--modalities text,audio \
|
||||
"$@"
|
||||
@@ -1,391 +0,0 @@
|
||||
# BAGEL-7B-MoT
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Architecture
|
||||
|
||||
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
|
||||
|
||||
| Topology | Stages | Description |
|
||||
| :------- | :----- | :---------- |
|
||||
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
|
||||
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
|
||||
|
||||
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
|
||||
|
||||
## Quick Start
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
# Default two-stage mode (auto-detected)
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
|
||||
# Single-stage mode
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
|
||||
|
||||
## Modality Control
|
||||
|
||||
Control the mode using the `--modality` argument:
|
||||
|
||||
| Modality | Input | Output | Description |
|
||||
| :------- | :---- | :----- | :---------- |
|
||||
| `text2img` | Text | Image | Generate images from text prompts |
|
||||
| `img2img` | Image + Text | Image | Transform images using text guidance |
|
||||
| `img2text` | Image + Text | Text | Generate text descriptions from images |
|
||||
| `text2text` | Text | Text | Pure text generation (language model mode) |
|
||||
|
||||
### Text to Image (text2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--steps 50
|
||||
```
|
||||
|
||||
### Image to Image (img2img)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2img \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Let the woman wear a blue dress"
|
||||
```
|
||||
|
||||
### Image to Text (img2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2text \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Describe this image in detail"
|
||||
```
|
||||
|
||||
### Text to Text (text2text)
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "What is the capital of France?"
|
||||
|
||||
# Load prompts from a text file (one prompt per line):
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--txt-prompts /path/to/prompts.txt
|
||||
```
|
||||
|
||||
## Think Mode
|
||||
|
||||
Think mode enables the model to generate `<think>...</think>` planning/reasoning tokens before producing the final output. This improves generation quality for complex prompts.
|
||||
|
||||
- **Two-stage**: The Thinker (AR) stage decodes think tokens, then transfers the augmented KV cache to the DiT stage for image generation.
|
||||
- **Single-stage**: The DiT's internal LLM generates think tokens in-place before proceeding to denoise.
|
||||
|
||||
```bash
|
||||
# Think + text2img: plan before generating
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A futuristic city with flying cars" \
|
||||
--think \
|
||||
--max-think-tokens 1000
|
||||
|
||||
# Think + img2img: reason about the edit
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2img \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "Make it look like a watercolor painting" \
|
||||
--think
|
||||
|
||||
# Think + img2text: reason before describing
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality img2text \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompts "What is happening in this image?" \
|
||||
--think
|
||||
|
||||
# Think + text2text: chain-of-thought reasoning
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2text \
|
||||
--prompts "Solve: 23 * 47" \
|
||||
--think
|
||||
```
|
||||
|
||||
Think mode parameters:
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--think` | `False` | Enable thinking mode |
|
||||
| `--max-think-tokens` | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | `False` | Enable sampling (vs. greedy) for text generation |
|
||||
| `--text-temperature` | `0.3` | Temperature for text generation sampling |
|
||||
|
||||
## Classifier-Free Guidance (CFG)
|
||||
|
||||
CFG controls the trade-off between prompt fidelity and diversity. These parameters apply to image generation modalities (`text2img`, `img2img`).
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A photorealistic portrait" \
|
||||
--cfg-text-scale 6.0 \
|
||||
--cfg-img-scale 2.0 \
|
||||
--negative-prompt "blurry, low quality, distorted" \
|
||||
--cfg-interval 0.4 1.0 \
|
||||
--cfg-renorm-type global \
|
||||
--cfg-renorm-min 0.0
|
||||
```
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | `4.0` | Text CFG scale (higher = more prompt-adherent) |
|
||||
| `--cfg-img-scale` | `1.5` | Image CFG scale (for img2img) |
|
||||
| `--negative-prompt` | `None` | Negative prompt for CFG conditioning |
|
||||
| `--cfg-interval` | pipeline default | CFG active interval `[start, end]` as fractions of total timesteps |
|
||||
| `--cfg-renorm-type` | `None` | Renormalization type: `global`, `text_channel`, `channel` |
|
||||
| `--cfg-renorm-min` | `None` | Minimum renormalization value |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size: `1` = batched (single GPU), `2` = 2-branch parallel, `3` = full 3-GPU parallel |
|
||||
|
||||
## Deployment Topologies
|
||||
|
||||
### Two-Stage (Default)
|
||||
|
||||
The default topology auto-detected from the model. No extra flags needed.
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
|
||||
The pipeline is defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml). Stage 0 (Thinker) and Stage 1 (DiT) share GPU 0 by default. For dual-GPU setups, customize the deploy YAML and set `devices: "1"` for stage 1.
|
||||
|
||||
### Single-Stage
|
||||
|
||||
Pass the single-stage deploy config via `--deploy-config`:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration details. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
|
||||
|
||||
### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, customize the deploy YAML (see [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)) and set per-stage `tensor_parallel_size` and `devices`:
|
||||
|
||||
```yaml
|
||||
# Example: TP=2 on GPUs 0,1 for the Thinker stage
|
||||
stages:
|
||||
- stage_id: 0
|
||||
tensor_parallel_size: 2
|
||||
devices: "0,1"
|
||||
```
|
||||
|
||||
### VAE Patch Parallelism
|
||||
|
||||
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) splits Bagel VAE **decode/encode** tiles across multiple GPUs on the **DiT stage**, reducing **per-GPU peak memory during VAE decode**. Use it when high-resolution `text2img` or `img2img` hits VAE OOM or large decode spikes.
|
||||
|
||||
**Bagel-specific notes:**
|
||||
|
||||
- Implemented in `BagelPipeline` via `DistributedAutoEncoder` (DiT stage only).
|
||||
- **Single-stage** is the simplest path: one DiT process with TP + VAE patch parallel.
|
||||
- **Two-stage**: enable on **stage 1 (DiT)** only; stage 0 (Thinker) keeps encoder-only `VAEEncoder` and does not use VAE patch parallel.
|
||||
- You need a DiT `world_size` ≥ `vae_patch_parallel_size` (typically `tensor_parallel_size=2` on that stage). VAE PP reuses the DiT process group; it is not a standalone second-GPU VAE worker.
|
||||
|
||||
**Single-stage via deploy YAML** (recommended for `end2end.py`):
|
||||
|
||||
```yaml
|
||||
pipeline: bagel_single_stage
|
||||
async_chunk: false
|
||||
|
||||
stages:
|
||||
- stage_id: 0
|
||||
max_num_batched_tokens: 32768
|
||||
max_num_seqs: 1
|
||||
enforce_eager: true
|
||||
trust_remote_code: true
|
||||
enable_prefix_caching: false
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
default_sampling_params:
|
||||
seed: 52
|
||||
```
|
||||
|
||||
```bash
|
||||
cd examples/offline_inference/bagel
|
||||
|
||||
CUDA_VISIBLE_DEVICES=0,1 python end2end.py \
|
||||
--model /path/to/BAGEL-7B-MoT \
|
||||
--deploy-config /path/to/bagel_single_stage_vae_pp.yaml \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--steps 10 \
|
||||
--output ./out_vae_pp
|
||||
```
|
||||
|
||||
**Single-stage via `Omni` kwargs** (same flags as online serving):
|
||||
|
||||
```python
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni = Omni(
|
||||
model="ByteDance-Seed/BAGEL-7B-MoT",
|
||||
deploy_config="vllm_omni/deploy/bagel_single_stage.yaml",
|
||||
tensor_parallel_size=2,
|
||||
vae_patch_parallel_size=2,
|
||||
vae_use_tiling=True,
|
||||
)
|
||||
# Then call omni.generate(...) as in end2end.py
|
||||
```
|
||||
|
||||
**Two-stage (VAE PP on DiT only):**
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
# AR Thinker — no vae_patch_parallel here
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
```
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--deploy-config /path/to/bagel_vae_pp.yaml \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat"
|
||||
```
|
||||
|
||||
**Startup log checks:**
|
||||
|
||||
```text
|
||||
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
|
||||
```
|
||||
|
||||
| Setting | Role |
|
||||
| :------ | :--- |
|
||||
| `parallel_config.tensor_parallel_size` | DiT world size / TP (must be ≥ `vae_patch_parallel_size`) |
|
||||
| `parallel_config.vae_patch_parallel_size` | Number of ranks for distributed VAE tiles (`1` = off) |
|
||||
| `vae_use_tiling` | Enable spatial tiling (auto-enabled when `vae_patch_parallel_size > 1`) |
|
||||
|
||||
#### Hybrid Sharded Data Parallel (HSDP)
|
||||
|
||||
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
||||
|
||||
1. **Enable HSDP**: Set `use_hsdp: true`.
|
||||
2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`).
|
||||
3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups.
|
||||
4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`).
|
||||
|
||||
Example configuration for HSDP across 4 GPUs:
|
||||
```yaml
|
||||
- stage_id: 1
|
||||
devices: "0,1,2,3"
|
||||
parallel_config:
|
||||
use_hsdp: true
|
||||
hsdp_shard_size: 4
|
||||
hsdp_replicate_size: 1
|
||||
```
|
||||
|
||||
Then pass the custom deploy YAML:
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--deploy-config /path/to/custom_bagel.yaml
|
||||
```
|
||||
|
||||
### FP8 Quantization
|
||||
|
||||
```bash
|
||||
python end2end.py --model ByteDance-Seed/BAGEL-7B-MoT \
|
||||
--modality text2img \
|
||||
--prompts "A cute cat" \
|
||||
--quantization fp8
|
||||
```
|
||||
|
||||
## Command Line Reference
|
||||
|
||||
### Core Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--model` | string | `ByteDance-Seed/BAGEL-7B-MoT` | Model path or HuggingFace name |
|
||||
| `--modality` | choice | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--prompts` | list | `None` | Input text prompts |
|
||||
| `--txt-prompts` | string | `None` | Path to text file with one prompt per line |
|
||||
| `--image-path` | string | `None` | Input image path (required for `img2img`/`img2text`) |
|
||||
| `--output` | string | `.` | Output directory for saved images |
|
||||
| `--steps` | int | `50` | Number of diffusion inference steps |
|
||||
| `--seed` | int | `None` | Random seed for reproducibility |
|
||||
|
||||
### Think Mode Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--think` | flag | `False` | Enable `<think>...</think>` planning/reasoning |
|
||||
| `--max-think-tokens` | int | `1000` | Maximum tokens for think generation |
|
||||
| `--do-sample` | flag | `False` | Use sampling instead of greedy decoding |
|
||||
| `--text-temperature` | float | `0.3` | Sampling temperature for text generation |
|
||||
|
||||
### CFG Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--cfg-text-scale` | float | `4.0` | Text CFG guidance scale |
|
||||
| `--cfg-img-scale` | float | `1.5` | Image CFG guidance scale |
|
||||
| `--negative-prompt` | string | `None` | Negative prompt for CFG |
|
||||
| `--cfg-parallel-size` | int | `1` | CFG parallel GPU count (1, 2, or 3) |
|
||||
| `--cfg-interval` | float[2] | pipeline default | CFG active window `[start, end]` |
|
||||
| `--cfg-renorm-type` | string | `None` | `global`, `text_channel`, or `channel` |
|
||||
| `--cfg-renorm-min` | float | `None` | Minimum renormalization value |
|
||||
|
||||
### Engine Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| :------- | :--- | :------ | :---------- |
|
||||
| `--deploy-config` | string | `None` | Path to deploy YAML (auto-detected if omitted) |
|
||||
| `--worker-backend` | choice | `process` | `process` or `ray` |
|
||||
| `--ray-address` | string | `None` | Ray cluster address |
|
||||
| `--quantization` | string | `None` | Quantization method (e.g. `fp8`) |
|
||||
| `--log-stats` | flag | `False` | Enable statistics logging |
|
||||
| `--init-timeout` | int | `300` | Initialization timeout (seconds) |
|
||||
| `--batch-timeout` | int | `5` | Batch timeout (seconds) |
|
||||
| `--enable-diffusion-pipeline-profiler` | flag | `False` | Profile diffusion stage durations |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
|
||||
|
||||
**Two-stage VRAM usage:**
|
||||
|
||||
| Stage | VRAM |
|
||||
| :---- | :--- |
|
||||
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
|
||||
| Stage 1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
|
||||
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
|
||||
@@ -1,278 +0,0 @@
|
||||
import argparse
|
||||
import os
|
||||
|
||||
from vllm_omni.inputs.data import OmniPromptType
|
||||
from vllm_omni.model_executor.stage_input_processors.bagel import (
|
||||
GEN_THINK_SYSTEM_PROMPT,
|
||||
VLM_THINK_SYSTEM_PROMPT,
|
||||
)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="ByteDance-Seed/BAGEL-7B-MoT",
|
||||
help="Path to merged model directory.",
|
||||
)
|
||||
parser.add_argument("--prompts", nargs="+", default=None, help="Input text prompts.")
|
||||
parser.add_argument(
|
||||
"--txt-prompts",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to a .txt file with one prompt per line (preferred).",
|
||||
)
|
||||
parser.add_argument("--prompt-type", default="text", choices=["text"])
|
||||
|
||||
parser.add_argument(
|
||||
"--modality",
|
||||
default="text2img",
|
||||
choices=["text2img", "img2img", "img2text", "text2text"],
|
||||
help="Modality mode to control stage execution.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--image-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to input image for img2img.",
|
||||
)
|
||||
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
default=".",
|
||||
help="Output directory to save images.",
|
||||
)
|
||||
|
||||
# OmniLLM init args
|
||||
parser.add_argument("--log-stats", action="store_true", default=False)
|
||||
parser.add_argument("--init-sleep-seconds", type=int, default=20)
|
||||
parser.add_argument("--batch-timeout", type=int, default=5)
|
||||
parser.add_argument("--init-timeout", type=int, default=300)
|
||||
parser.add_argument("--shm-threshold-bytes", type=int, default=65536)
|
||||
parser.add_argument("--worker-backend", type=str, default="process", choices=["process", "ray"])
|
||||
parser.add_argument("--ray-address", type=str, default=None)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Path to deploy YAML. If unset, auto-loads vllm_omni/deploy/bagel.yaml based on the HF model_type.",
|
||||
)
|
||||
parser.add_argument("--steps", type=int, default=50, help="Number of inference steps.")
|
||||
|
||||
parser.add_argument("--cfg-text-scale", type=float, default=4.0, help="Text CFG scale (default: 4.0)")
|
||||
parser.add_argument("--cfg-img-scale", type=float, default=1.5, help="Image CFG scale (default: 1.5)")
|
||||
parser.add_argument(
|
||||
"--negative-prompt", type=str, default=None, help="Negative prompt for CFG (default: empty prompt)"
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-parallel-size",
|
||||
type=int,
|
||||
default=1,
|
||||
choices=[1, 2, 3],
|
||||
help="CFG parallel size: 1=batched (single GPU), 2=parallel with 2 branches (text CFG only), 3=parallel (3 GPUs).",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=None, help="Random seed for generation.")
|
||||
parser.add_argument(
|
||||
"--cfg-interval",
|
||||
type=float,
|
||||
nargs=2,
|
||||
default=None,
|
||||
help="CFG interval [start, end] (default: pipeline default)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--cfg-renorm-type", type=str, default=None, help="CFG renorm type: global, text_channel, channel"
|
||||
)
|
||||
parser.add_argument("--cfg-renorm-min", type=float, default=None, help="CFG renorm min")
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Quantization method (e.g. 'fp8').",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--think",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable thinking mode: AR stage decodes <think>...</think> planning tokens before image generation.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--max-think-tokens",
|
||||
type=int,
|
||||
default=1000,
|
||||
help="Maximum number of tokens for thinking text generation (default: 1000).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--do-sample",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable sampling for text generation (default: greedy).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--text-temperature",
|
||||
type=float,
|
||||
default=0.3,
|
||||
help="Temperature for text generation sampling (default: 0.3).",
|
||||
)
|
||||
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
os.makedirs(args.output, exist_ok=True)
|
||||
model_name = args.model
|
||||
prompts: list[OmniPromptType] = []
|
||||
try:
|
||||
if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
|
||||
with open(args.txt_prompts, encoding="utf-8") as f:
|
||||
lines = [ln.strip() for ln in f.readlines()]
|
||||
prompts = [ln for ln in lines if ln != ""]
|
||||
print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}")
|
||||
else:
|
||||
prompts = args.prompts
|
||||
except Exception as e:
|
||||
print(f"[Error] Failed to load prompts: {e}")
|
||||
raise
|
||||
|
||||
if not prompts:
|
||||
prompts = ["A cute cat"]
|
||||
print(f"[Info] No prompts provided, using default: {prompts}")
|
||||
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
omni_kwargs.pop("model", None)
|
||||
deploy_config = args.deploy_config
|
||||
if args.think and deploy_config is None:
|
||||
deploy_config = "vllm_omni/deploy/bagel_think.yaml"
|
||||
print(f"[Info] Think mode enabled, using deploy config: {deploy_config}")
|
||||
if deploy_config:
|
||||
omni_kwargs["deploy_config"] = deploy_config
|
||||
|
||||
if args.quantization:
|
||||
omni_kwargs["quantization_config"] = args.quantization
|
||||
|
||||
omni = Omni(model=model_name, **omni_kwargs)
|
||||
|
||||
formatted_prompts = []
|
||||
for p in prompts:
|
||||
if args.modality == "img2img":
|
||||
if not args.image_path or not os.path.exists(args.image_path):
|
||||
raise ValueError(f"img2img requires --image-path pointing to an existing file, got: {args.image_path}")
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
|
||||
final_prompt_text = f"{think_prefix}<|fim_middle|><|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"img2img": loaded_image},
|
||||
"modalities": ["img2img"],
|
||||
}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "img2text":
|
||||
if args.image_path:
|
||||
loaded_image = Image.open(args.image_path).convert("RGB")
|
||||
think_prefix = f"<|im_start|>system\n{VLM_THINK_SYSTEM_PROMPT}<|im_end|>\n" if args.think else ""
|
||||
final_prompt_text = (
|
||||
f"{think_prefix}<|im_start|>user\n<|image_pad|>\n{p}<|im_end|>\n<|im_start|>assistant\n"
|
||||
)
|
||||
prompt_dict = {
|
||||
"prompt": final_prompt_text,
|
||||
"multi_modal_data": {"image": loaded_image},
|
||||
"modalities": ["text"],
|
||||
}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
elif args.modality == "text2text":
|
||||
think_prefix = f"<|im_start|>{VLM_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
|
||||
final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|><|im_start|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["text"]}
|
||||
formatted_prompts.append(prompt_dict)
|
||||
else:
|
||||
think_prefix = f"<|im_start|>{GEN_THINK_SYSTEM_PROMPT}<|im_end|>" if args.think else ""
|
||||
final_prompt_text = f"{think_prefix}<|im_start|>{p}<|im_end|>"
|
||||
prompt_dict = {"prompt": final_prompt_text, "modalities": ["image"]}
|
||||
if args.negative_prompt is not None:
|
||||
prompt_dict["negative_prompt"] = args.negative_prompt
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
params_list = omni.default_sampling_params_list
|
||||
# Bagel exposes 1 sampling param set for single-stage (DiT-only) and
|
||||
# 2 for two-stage (Thinker + DiT). This heuristic may need updating
|
||||
# if future pipelines break that 1:1 mapping.
|
||||
is_single_stage = len(params_list) == 1
|
||||
|
||||
diffusion_params_idx = 0 if is_single_stage else (1 if len(params_list) > 1 else 0)
|
||||
diffusion_params = params_list[diffusion_params_idx]
|
||||
|
||||
if args.modality in ("text2img", "img2img"):
|
||||
diffusion_params.num_inference_steps = args.steps # type: ignore
|
||||
diffusion_params.cfg_parallel_size = args.cfg_parallel_size # type: ignore
|
||||
if args.seed is not None:
|
||||
diffusion_params.seed = args.seed # type: ignore
|
||||
|
||||
extra = getattr(diffusion_params, "extra_args", {}) or {}
|
||||
extra["cfg_text_scale"] = args.cfg_text_scale
|
||||
extra["cfg_img_scale"] = args.cfg_img_scale
|
||||
if args.cfg_interval is not None:
|
||||
extra["cfg_interval"] = tuple(args.cfg_interval)
|
||||
if args.cfg_renorm_type is not None:
|
||||
extra["cfg_renorm_type"] = args.cfg_renorm_type
|
||||
if args.cfg_renorm_min is not None:
|
||||
extra["cfg_renorm_min"] = args.cfg_renorm_min
|
||||
if args.negative_prompt is not None:
|
||||
extra["negative_prompt"] = args.negative_prompt
|
||||
|
||||
needs_text_gen = is_single_stage and (args.think or args.modality in ("text2text", "img2text"))
|
||||
if needs_text_gen:
|
||||
if args.think:
|
||||
extra["think"] = True
|
||||
extra["max_think_tokens"] = args.max_think_tokens
|
||||
extra["do_sample"] = args.do_sample
|
||||
extra["text_temperature"] = args.text_temperature
|
||||
diffusion_params.extra_args = extra # type: ignore
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
|
||||
img_idx = 0
|
||||
for req_output in omni_outputs:
|
||||
# 2-stage think mode: text output from thinker stage
|
||||
ro = getattr(req_output, "request_output", None)
|
||||
if ro and getattr(ro, "outputs", None):
|
||||
txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
|
||||
if txt:
|
||||
if args.think:
|
||||
print(f"[Think]\n{txt}")
|
||||
else:
|
||||
print(f"[Output] Text:\n{txt}")
|
||||
|
||||
# Single-stage DiT: text from custom_output
|
||||
custom = getattr(req_output, "_custom_output", {}) or {}
|
||||
if custom.get("think_text"):
|
||||
print(f"[Think]\n{custom['think_text']}")
|
||||
if custom.get("text_output"):
|
||||
print(f"[Output] Text:\n{custom['text_output']}")
|
||||
|
||||
images = getattr(req_output, "images", None)
|
||||
if not images:
|
||||
continue
|
||||
|
||||
for j, img in enumerate(images):
|
||||
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
|
||||
img.save(save_path)
|
||||
print(f"[Output] Saved image to {save_path}")
|
||||
img_idx += 1
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,14 +0,0 @@
|
||||
{
|
||||
"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"
|
||||
}
|
||||
@@ -0,0 +1,129 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""DreamZero AR client frame schedule aligned with upstream ``test_client_AR.py``."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import logging
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
|
||||
# Frame schedule constants (matching upstream debug_inference.py / test_client_AR.py).
|
||||
ACTION_HORIZON = 24
|
||||
DEFAULT_ACTION_DIM = 8
|
||||
RELATIVE_OFFSETS = [-23, -16, -8, 0]
|
||||
DEFAULT_NUM_AR_CHUNKS = 15
|
||||
|
||||
CAMERA_FILES = {
|
||||
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
|
||||
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
|
||||
"observation/wrist_image_left": "wrist_image_left.mp4",
|
||||
}
|
||||
|
||||
|
||||
def build_frame_schedule(
|
||||
total_frames: int,
|
||||
num_chunks: int,
|
||||
*,
|
||||
logger: logging.Logger | None = None,
|
||||
) -> list[list[int]]:
|
||||
"""Build the frame index schedule for multi-frame chunks.
|
||||
|
||||
Args:
|
||||
total_frames: Number of frames available in each camera video.
|
||||
num_chunks: Number of 4-frame chunks to schedule after the initial frame.
|
||||
|
||||
Returns:
|
||||
A list of frame-index lists. Each inner list has four indices. The
|
||||
returned list may be shorter than ``num_chunks`` when the videos run
|
||||
out of frames (upstream stops early instead of erroring).
|
||||
"""
|
||||
chunks: list[list[int]] = []
|
||||
current_frame = 23
|
||||
for _ in range(num_chunks):
|
||||
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
|
||||
if indices[-1] >= total_frames:
|
||||
if logger is not None:
|
||||
logger.info(
|
||||
"Frame %s >= %s, stopping at %s chunks",
|
||||
indices[-1],
|
||||
total_frames,
|
||||
len(chunks),
|
||||
)
|
||||
break
|
||||
chunks.append(indices)
|
||||
current_frame += ACTION_HORIZON
|
||||
return chunks
|
||||
|
||||
|
||||
def make_obs_from_video(
|
||||
camera_frames: dict[str, np.ndarray],
|
||||
frame_indices: list[int],
|
||||
*,
|
||||
prompt: str,
|
||||
session_id: str,
|
||||
) -> dict[str, Any]:
|
||||
"""Build an observation dict from real video frames.
|
||||
|
||||
For one frame each image key is ``(H, W, 3)``. For four frames each key is
|
||||
``(4, H, W, 3)``.
|
||||
"""
|
||||
obs: dict[str, Any] = {}
|
||||
for camera_key, all_frames in camera_frames.items():
|
||||
selected = all_frames[frame_indices]
|
||||
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
|
||||
|
||||
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
||||
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
||||
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
||||
obs["prompt"] = prompt
|
||||
obs["session_id"] = session_id
|
||||
return obs
|
||||
|
||||
|
||||
def build_ar_observations(
|
||||
camera_frames: dict[str, np.ndarray],
|
||||
*,
|
||||
prompt: str,
|
||||
session_id: str,
|
||||
num_chunks: int = DEFAULT_NUM_AR_CHUNKS,
|
||||
repeat_chunk_observations: bool = False,
|
||||
logger: logging.Logger | None = None,
|
||||
) -> list[dict[str, Any]]:
|
||||
"""Build the AR observation sequence used by upstream ``test_client_AR.py``.
|
||||
|
||||
Step 0 sends frame ``[0]``. Each subsequent step sends one 4-frame chunk.
|
||||
|
||||
``num_chunks`` counts only the 4-frame chunks after the initial frame, not
|
||||
the total number of inferences.
|
||||
"""
|
||||
if num_chunks < 0:
|
||||
raise ValueError("num_chunks must be non-negative")
|
||||
|
||||
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
||||
observations = [
|
||||
make_obs_from_video(
|
||||
camera_frames,
|
||||
[0],
|
||||
prompt=prompt,
|
||||
session_id=session_id,
|
||||
)
|
||||
]
|
||||
|
||||
chunk_schedule = build_frame_schedule(total_frames, num_chunks, logger=logger)
|
||||
if repeat_chunk_observations and chunk_schedule and len(chunk_schedule) < num_chunks:
|
||||
while len(chunk_schedule) < num_chunks:
|
||||
chunk_schedule.append(chunk_schedule[-1])
|
||||
|
||||
for frame_indices in chunk_schedule:
|
||||
observations.append(
|
||||
make_obs_from_video(
|
||||
camera_frames,
|
||||
frame_indices,
|
||||
prompt=prompt,
|
||||
session_id=session_id,
|
||||
)
|
||||
)
|
||||
return observations
|
||||
@@ -11,6 +11,11 @@ from pathlib import Path
|
||||
import cv2
|
||||
import numpy as np
|
||||
import torch
|
||||
from client_schedule import (
|
||||
CAMERA_FILES,
|
||||
DEFAULT_NUM_AR_CHUNKS,
|
||||
build_ar_observations,
|
||||
)
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni import Omni
|
||||
@@ -20,13 +25,6 @@ from vllm_omni.outputs import OmniRequestOutput
|
||||
WORKER_EXTENSION = "vllm_omni.diffusion.models.dreamzero.video_export_worker.DreamZeroVideoExportWorkerExtension"
|
||||
ASSET_REPO_ID = "YangshenDeng/vllm-omni-dreamzero-assets"
|
||||
DEFAULT_SESSION_PREFIX = "dreamzero-export"
|
||||
RELATIVE_OFFSETS = [-23, -16, -8, 0]
|
||||
ACTION_HORIZON = 24
|
||||
CAMERA_FILES = {
|
||||
"observation/exterior_image_0_left": "exterior_image_1_left.mp4",
|
||||
"observation/exterior_image_1_left": "exterior_image_2_left.mp4",
|
||||
"observation/wrist_image_left": "wrist_image_left.mp4",
|
||||
}
|
||||
|
||||
|
||||
def _parse_args() -> argparse.Namespace:
|
||||
@@ -49,6 +47,17 @@ def _parse_args() -> argparse.Namespace:
|
||||
default="Move the pan forward and use the brush in the middle of the plates to brush the inside of the pan",
|
||||
)
|
||||
parser.add_argument("--session-id", default=None)
|
||||
parser.add_argument(
|
||||
"--num-chunks",
|
||||
type=int,
|
||||
default=DEFAULT_NUM_AR_CHUNKS,
|
||||
help="Number of 4-frame chunks to send after the initial frame (default: 15).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--repeat-chunk-observations",
|
||||
action="store_true",
|
||||
help="Repeat the last valid chunk observation when bundled assets run out of frames.",
|
||||
)
|
||||
parser.add_argument("--save-input-video", action="store_true")
|
||||
parser.add_argument("--save-gif", action="store_true")
|
||||
parser.add_argument("--save-actions", action="store_true")
|
||||
@@ -84,57 +93,25 @@ def _load_camera_frames(video_dir: Path) -> dict[str, np.ndarray]:
|
||||
return camera_frames
|
||||
|
||||
|
||||
def _build_frame_schedule(total_frames: int, num_chunks: int) -> list[list[int]]:
|
||||
chunks: list[list[int]] = []
|
||||
current_frame = 23
|
||||
for _ in range(num_chunks):
|
||||
indices = [max(current_frame + offset, 0) for offset in RELATIVE_OFFSETS]
|
||||
if indices[-1] >= total_frames:
|
||||
break
|
||||
chunks.append(indices)
|
||||
current_frame += ACTION_HORIZON
|
||||
return chunks
|
||||
|
||||
|
||||
def _make_obs_from_video(
|
||||
camera_frames: dict[str, np.ndarray],
|
||||
frame_indices: list[int],
|
||||
*,
|
||||
def _build_observations(
|
||||
video_dir: Path,
|
||||
prompt: str,
|
||||
session_id: str,
|
||||
) -> dict:
|
||||
obs: dict = {}
|
||||
for camera_key, all_frames in camera_frames.items():
|
||||
selected = all_frames[frame_indices]
|
||||
obs[camera_key] = selected[0] if len(frame_indices) == 1 else selected
|
||||
|
||||
obs["observation/joint_position"] = np.zeros(7, dtype=np.float32)
|
||||
obs["observation/cartesian_position"] = np.zeros(6, dtype=np.float32)
|
||||
obs["observation/gripper_position"] = np.zeros(1, dtype=np.float32)
|
||||
obs["prompt"] = prompt
|
||||
obs["session_id"] = session_id
|
||||
return obs
|
||||
|
||||
|
||||
def _build_observations(video_dir: Path, prompt: str, session_id: str) -> tuple[dict[str, np.ndarray], list[dict]]:
|
||||
*,
|
||||
num_chunks: int,
|
||||
repeat_chunk_observations: bool,
|
||||
) -> tuple[dict[str, np.ndarray], list[dict]]:
|
||||
camera_frames = _load_camera_frames(video_dir)
|
||||
total_frames = min(frames.shape[0] for frames in camera_frames.values())
|
||||
chunks = _build_frame_schedule(total_frames, 1)
|
||||
observations = [
|
||||
_make_obs_from_video(camera_frames, [0], prompt=prompt, session_id=session_id),
|
||||
]
|
||||
if chunks:
|
||||
observations.append(
|
||||
_make_obs_from_video(
|
||||
camera_frames,
|
||||
chunks[0],
|
||||
prompt=prompt,
|
||||
session_id=session_id,
|
||||
)
|
||||
)
|
||||
observations = build_ar_observations(
|
||||
camera_frames,
|
||||
prompt=prompt,
|
||||
session_id=session_id,
|
||||
num_chunks=num_chunks,
|
||||
repeat_chunk_observations=repeat_chunk_observations,
|
||||
)
|
||||
if len(observations) < 2:
|
||||
raise RuntimeError("Need at least two DreamZero example observations to export a prediction video.")
|
||||
return camera_frames, observations[:2]
|
||||
raise RuntimeError("Need at least two DreamZero observations to export a prediction video.")
|
||||
return camera_frames, observations
|
||||
|
||||
|
||||
def _extract_latents(output: OmniRequestOutput) -> torch.Tensor:
|
||||
@@ -229,6 +206,26 @@ def _run_generation(
|
||||
return omni, outputs
|
||||
|
||||
|
||||
def _decode_accumulated_with_worker(omni: Omni, session_id: str) -> np.ndarray:
|
||||
"""Decode server-side accumulated latents when the pipeline supports it."""
|
||||
stage_client = omni.engine.stage_clients[0]
|
||||
engine = getattr(stage_client, "_engine", None)
|
||||
if engine is None:
|
||||
raise RuntimeError("DreamZero export requires inline diffusion stage access.")
|
||||
|
||||
decoded = engine.executor.collective_rpc(
|
||||
"decode_accumulated_video_latents_to_uint8",
|
||||
args=(session_id,),
|
||||
unique_reply_rank=0,
|
||||
exec_all_ranks=True,
|
||||
)
|
||||
if isinstance(decoded, torch.Tensor):
|
||||
decoded = decoded.numpy()
|
||||
if not isinstance(decoded, np.ndarray):
|
||||
raise TypeError(f"Unexpected decoded output type: {type(decoded)!r}")
|
||||
return decoded
|
||||
|
||||
|
||||
def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
|
||||
stage_client = omni.engine.stage_clients[0]
|
||||
engine = getattr(stage_client, "_engine", None)
|
||||
@@ -248,6 +245,15 @@ def _decode_with_worker(omni: Omni, full_latents: torch.Tensor) -> np.ndarray:
|
||||
return decoded
|
||||
|
||||
|
||||
def _decode_video(omni: Omni, session_id: str, outputs: list[OmniRequestOutput]) -> np.ndarray:
|
||||
"""Decode full AR rollout by concatenating per-step latents from all outputs."""
|
||||
latent_steps = [_extract_latents(output) for output in outputs]
|
||||
latent_t = [step.shape[2] for step in latent_steps]
|
||||
print(f"Decoding {len(latent_steps)} latent steps, latent_t={latent_t}, total_t={sum(latent_t)}")
|
||||
full_latents = torch.cat(latent_steps, dim=2)
|
||||
return _decode_with_worker(omni, full_latents)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = _parse_args()
|
||||
session_id = args.session_id or f"{DEFAULT_SESSION_PREFIX}-{uuid.uuid4()}"
|
||||
@@ -256,6 +262,12 @@ def main() -> None:
|
||||
video_dir=args.video_dir,
|
||||
prompt=args.prompt,
|
||||
session_id=session_id,
|
||||
num_chunks=args.num_chunks,
|
||||
repeat_chunk_observations=args.repeat_chunk_observations,
|
||||
)
|
||||
print(
|
||||
f"Running {len(observations)} inferences "
|
||||
f"(1 initial + {len(observations) - 1} chunks; requested {args.num_chunks} post-initial chunks)"
|
||||
)
|
||||
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
@@ -273,17 +285,19 @@ def main() -> None:
|
||||
deploy_config_path=args.deploy_config,
|
||||
observations=observations,
|
||||
)
|
||||
latent_steps = [_extract_latents(output) for output in outputs]
|
||||
full_latents = torch.cat(latent_steps, dim=2)
|
||||
frames = _decode_with_worker(omni, full_latents)
|
||||
frames = _decode_video(omni, session_id, outputs)
|
||||
finally:
|
||||
if omni is not None:
|
||||
omni.close()
|
||||
|
||||
expected_frames = 8 * len(observations) + 1
|
||||
mp4_path = args.output_dir / f"{args.output_stem}.mp4"
|
||||
|
||||
_write_mp4(mp4_path, frames, fps=args.fps)
|
||||
print(f"SAVED_MP4={mp4_path}")
|
||||
print(
|
||||
f"SAVED_MP4={mp4_path} frames={len(frames)} "
|
||||
f"expected≈{expected_frames} duration≈{len(frames) / max(args.fps, 1):.1f}s"
|
||||
)
|
||||
|
||||
if args.save_gif:
|
||||
gif_path = args.output_dir / f"{args.output_stem}.gif"
|
||||
@@ -294,8 +308,10 @@ def main() -> None:
|
||||
npz_path = args.output_dir / f"{args.output_stem}_actions.npz"
|
||||
np.savez(
|
||||
npz_path,
|
||||
step0=np.asarray(outputs[0].multimodal_output.get("actions")),
|
||||
step1=np.asarray(outputs[1].multimodal_output.get("actions")),
|
||||
**{
|
||||
f"step{index}": np.asarray(output.multimodal_output.get("actions"))
|
||||
for index, output in enumerate(outputs)
|
||||
},
|
||||
)
|
||||
print(f"SAVED_ACTIONS={npz_path}")
|
||||
|
||||
|
||||
@@ -1,193 +0,0 @@
|
||||
# Helios Video Generation
|
||||
|
||||
Helios is a text-to-video (T2V), image-to-video (I2V), and video-to-video (V2V) diffusion model. This example demonstrates end-to-end video generation using vLLM-Omni with three model variants:
|
||||
|
||||
| Variant | Description | Key Features |
|
||||
|---------|-------------|--------------|
|
||||
| **Helios-Base** | Base model, Stage 1 only | Single-stage denoising, `guidance_scale=5.0` |
|
||||
| **Helios-Mid** | Mid model, Stage 2 pyramid | Multi-stage pyramid denoising, CFG-Zero* support |
|
||||
| **Helios-Distilled** | Distilled model, Stage 2+3 | Few-step inference with DMD, `guidance_scale=1.0` |
|
||||
|
||||
## Setup
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/helios
|
||||
```
|
||||
|
||||
### Text-to-Video (T2V)
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type t2v \
|
||||
--prompt "A dynamic time-lapse video showing the rapidly moving scenery from the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_i2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_i2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "A towering emerald wave surges forward, its crest curling with raw power and energy." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_i2v_distilled.mp4
|
||||
```
|
||||
|
||||
### Video-to-Video (V2V)
|
||||
|
||||
**Helios-Base**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_v2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 + CFG-Zero*):
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1 \
|
||||
--output helios_v2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled**:
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "A bright yellow Lamborghini speeds along a curving mountain road." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk \
|
||||
--output helios_v2v_distilled.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | `BestWishYsh/Helios-Base` | Model ID or local path |
|
||||
| `--sample-type` | `t2v` | Generation mode: `t2v`, `i2v`, or `v2v` |
|
||||
| `--prompt` | — | Text prompt describing the video |
|
||||
| `--negative-prompt` | *(see source)* | Negative prompt for CFG (includes anti-static terms) |
|
||||
| `--image-path` | — | Input image (required for `i2v`) |
|
||||
| `--video-path` | — | Input video (required for `v2v`) |
|
||||
| `--height` | `384` | Video height in pixels |
|
||||
| `--width` | `640` | Video width in pixels |
|
||||
| `--num-frames` | `99` | Number of output frames |
|
||||
| `--num-inference-steps` | `50` | Denoising steps (Stage 1 only) |
|
||||
| `--guidance-scale` | `5.0` | CFG scale (`1.0` for Distilled) |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--fps` | `16` | Output video frame rate |
|
||||
| `--output` | `helios_output.mp4` | Output file path |
|
||||
|
||||
### Stage 2 / Pyramid Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--is-enable-stage2` | off | Enable pyramid multi-stage denoising |
|
||||
| `--pyramid-num-stages` | `3` | Number of pyramid stages |
|
||||
| `--pyramid-num-inference-steps-list` | `10 10 10` | Steps per pyramid stage |
|
||||
| `--is-amplify-first-chunk` | off | DMD amplification (Distilled) |
|
||||
|
||||
### CFG-Zero* Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--use-cfg-zero-star` | off | Enable CFG-Zero* guidance (Mid) |
|
||||
| `--use-zero-init` | off | Zero init for first steps |
|
||||
| `--zero-steps` | `1` | Number of zero-init steps |
|
||||
|
||||
### Memory & Parallelism
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--vae-use-slicing` | off | Enable VAE slicing |
|
||||
| `--vae-use-tiling` | off | Enable VAE tiling |
|
||||
| `--enforce-eager` | off | Disable torch.compile |
|
||||
| `--enable-cpu-offload` | off | CPU offloading |
|
||||
| `--enable-layerwise-offload` | off | Layerwise offloading |
|
||||
| `--ulysses-degree` | `1` | Ulysses SP degree |
|
||||
| `--ring-degree` | `1` | Ring SP degree |
|
||||
| `--cfg-parallel-size` | `1` | CFG parallel size (1 or 2) |
|
||||
| `--tensor-parallel-size` | `1` | Tensor parallelism size |
|
||||
@@ -1,355 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Helios video generation example supporting T2V, I2V, and V2V.
|
||||
|
||||
Usage (T2V, Helios-Base, Stage 1 only):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--height 384 --width 640 --num-frames 99 \
|
||||
--num-inference-steps 50 --guidance-scale 5.0
|
||||
|
||||
Usage (I2V, Helios-Base):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type i2v \
|
||||
--image-path /path/to/image.jpg \
|
||||
--prompt "Description of desired animation." \
|
||||
--guidance-scale 5.0
|
||||
|
||||
Usage (V2V, Helios-Base):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Base --sample-type v2v \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--prompt "Description of desired transformation." \
|
||||
--guidance-scale 5.0
|
||||
|
||||
Usage (Helios-Mid, Stage 2 + CFG-Zero*):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Mid --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--guidance-scale 5.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 20 20 20 \
|
||||
--use-cfg-zero-star --use-zero-init --zero-steps 1
|
||||
|
||||
Usage (Helios-Distilled, Stage 2 pyramid + DMD):
|
||||
python end2end.py \
|
||||
--model /path/to/Helios-Distilled --sample-type t2v \
|
||||
--prompt "A serene lakeside sunrise with mist over the water." \
|
||||
--guidance-scale 1.0 \
|
||||
--is-enable-stage2 \
|
||||
--pyramid-num-inference-steps-list 2 2 2 \
|
||||
--is-amplify-first-chunk
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def load_image_as_tensor(image_path: str, height: int, width: int) -> torch.Tensor:
|
||||
"""Load an image and return a (1, C, H, W) tensor normalized to [-1, 1]."""
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
|
||||
image = Image.open(image_path).convert("RGB").resize((width, height))
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
return transform(image).unsqueeze(0)
|
||||
|
||||
|
||||
def load_video_as_tensor(video_path: str, height: int, width: int) -> torch.Tensor:
|
||||
"""Load a video and return a (1, C, T, H, W) tensor normalized to [-1, 1]."""
|
||||
from torchvision import transforms
|
||||
from torchvision.io import read_video
|
||||
|
||||
video_frames, _, _ = read_video(video_path, output_format="TCHW")
|
||||
video_frames = video_frames.float() / 255.0
|
||||
|
||||
transform = transforms.Compose(
|
||||
[
|
||||
transforms.Resize((height, width), antialias=True),
|
||||
transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5]),
|
||||
]
|
||||
)
|
||||
video_frames = torch.stack([transform(f) for f in video_frames])
|
||||
return video_frames.permute(1, 0, 2, 3).unsqueeze(0)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video with Helios (T2V / I2V / V2V).")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="BestWishYsh/Helios-Base",
|
||||
help="Helios model ID or local path (e.g. Helios-Base, Helios-Mid, Helios-Distilled).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-type",
|
||||
choices=["t2v", "i2v", "v2v"],
|
||||
default="t2v",
|
||||
help="Generation mode: t2v (text-to-video), i2v (image-to-video), v2v (video-to-video).",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
||||
parser.add_argument(
|
||||
"--negative-prompt",
|
||||
default=(
|
||||
"Bright tones, overexposed, static, blurred details, subtitles, style, works, paintings, "
|
||||
"images, static, overall gray, worst quality, low quality, JPEG compression residue, ugly, "
|
||||
"incomplete, extra fingers, poorly drawn hands, poorly drawn faces, deformed, disfigured, "
|
||||
"misshapen limbs, fused fingers, still picture, messy background, three legs, many people "
|
||||
"in the background, walking backwards"
|
||||
),
|
||||
help="Negative prompt.",
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
|
||||
parser.add_argument("--height", type=int, default=384, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=640, help="Video width.")
|
||||
parser.add_argument("--num-frames", type=int, default=99, help="Number of video frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps (Stage 1 only).")
|
||||
parser.add_argument("--output", type=str, default="helios_output.mp4", help="Output video path.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Frames per second for the output video.")
|
||||
|
||||
# I2V / V2V
|
||||
parser.add_argument("--image-path", type=str, default=None, help="Input image path for I2V mode.")
|
||||
parser.add_argument("--video-path", type=str, default=None, help="Input video path for V2V mode.")
|
||||
|
||||
# Stage 2 (pyramid multi-stage denoising)
|
||||
parser.add_argument(
|
||||
"--is-enable-stage2",
|
||||
action="store_true",
|
||||
help="Enable pyramid multi-stage denoising (Stage 2). Required for Helios-Distilled.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pyramid-num-stages",
|
||||
type=int,
|
||||
default=3,
|
||||
help="Number of pyramid stages for Stage 2.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--pyramid-num-inference-steps-list",
|
||||
type=int,
|
||||
nargs="+",
|
||||
default=[10, 10, 10],
|
||||
help="Inference steps per pyramid stage.",
|
||||
)
|
||||
|
||||
# DMD
|
||||
parser.add_argument(
|
||||
"--is-amplify-first-chunk",
|
||||
action="store_true",
|
||||
help="Enable DMD amplification for the first chunk (Helios-Distilled).",
|
||||
)
|
||||
|
||||
# CFG Zero Star
|
||||
parser.add_argument(
|
||||
"--use-cfg-zero-star",
|
||||
action="store_true",
|
||||
help="Enable CFG Zero Star guidance (recommended for Helios-Mid).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-zero-init",
|
||||
action="store_true",
|
||||
help="Use zero initialization for the first denoising steps with CFG-Zero*.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--zero-steps",
|
||||
type=int,
|
||||
default=1,
|
||||
help="Number of initial denoising steps using zero prediction (default: 1).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable diffusion pipeline profiler to display stage durations.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--quantization",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Quantization method (e.g. 'fp8').",
|
||||
)
|
||||
|
||||
# Memory & parallelism
|
||||
parser.add_argument("--vae-use-slicing", action="store_true", help="Enable VAE slicing.")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true", help="Enable VAE tiling.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
parser.add_argument("--enable-cpu-offload", action="store_true", help="Enable CPU offloading.")
|
||||
parser.add_argument("--enable-layerwise-offload", action="store_true", help="Enable layerwise offloading.")
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
|
||||
parser.add_argument("--ring-degree", type=int, default=1, help="Ring SP degree.")
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
|
||||
parser.add_argument("--tensor-parallel-size", type=int, default=1, help="Tensor parallelism size.")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
enable_layerwise_offload=args.enable_layerwise_offload,
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
enable_cpu_offload=args.enable_cpu_offload,
|
||||
parallel_config=parallel_config,
|
||||
enforce_eager=args.enforce_eager,
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
quantization=args.quantization,
|
||||
)
|
||||
|
||||
# Validate I2V / V2V arguments
|
||||
if args.sample_type == "i2v" and not args.image_path:
|
||||
raise ValueError("--image-path is required for I2V mode.")
|
||||
if args.sample_type == "v2v" and not args.video_path:
|
||||
raise ValueError("--video-path is required for V2V mode.")
|
||||
|
||||
# Build extra_args for Helios-specific parameters
|
||||
extra_args = {}
|
||||
if args.is_enable_stage2:
|
||||
extra_args["is_enable_stage2"] = True
|
||||
extra_args["pyramid_num_stages"] = args.pyramid_num_stages
|
||||
extra_args["pyramid_num_inference_steps_list"] = args.pyramid_num_inference_steps_list
|
||||
if args.is_amplify_first_chunk:
|
||||
extra_args["is_amplify_first_chunk"] = True
|
||||
if args.use_cfg_zero_star:
|
||||
extra_args["use_cfg_zero_star"] = True
|
||||
if args.use_zero_init:
|
||||
extra_args["use_zero_init"] = True
|
||||
extra_args["zero_steps"] = args.zero_steps
|
||||
|
||||
if args.sample_type == "i2v":
|
||||
image_tensor = load_image_as_tensor(args.image_path, args.height, args.width)
|
||||
extra_args["image"] = image_tensor
|
||||
elif args.sample_type == "v2v":
|
||||
video_tensor = load_video_as_tensor(args.video_path, args.height, args.width)
|
||||
extra_args["video"] = video_tensor
|
||||
|
||||
# Print generation configuration
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Helios Generation Configuration:")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Sample type: {args.sample_type.upper()}")
|
||||
if args.sample_type == "i2v":
|
||||
print(f" Image: {args.image_path}")
|
||||
elif args.sample_type == "v2v":
|
||||
print(f" Video: {args.video_path}")
|
||||
print(f" Prompt: {args.prompt[:80]}{'...' if len(args.prompt) > 80 else ''}")
|
||||
print(f" Video size: {args.width}x{args.height}, {args.num_frames} frames")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Guidance scale: {args.guidance_scale}")
|
||||
if args.is_enable_stage2:
|
||||
print(f" Stage 2: enabled (stages={args.pyramid_num_stages}, steps={args.pyramid_num_inference_steps_list})")
|
||||
if args.is_amplify_first_chunk:
|
||||
print(" DMD amplify first chunk: enabled")
|
||||
if args.use_cfg_zero_star:
|
||||
print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
|
||||
else:
|
||||
if args.use_cfg_zero_star:
|
||||
print(f" CFG Zero Star: enabled (zero_init={args.use_zero_init}, zero_steps={args.zero_steps})")
|
||||
print(" Stage 2: disabled (Stage 1 only)")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
frames = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
generator=generator,
|
||||
guidance_scale=args.guidance_scale,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_frames=args.num_frames,
|
||||
extra_args=extra_args,
|
||||
),
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
print(f"Total generation time: {generation_time:.4f} seconds ({generation_time * 1000:.2f} ms)")
|
||||
|
||||
# Extract video frames from OmniRequestOutput
|
||||
if isinstance(frames, list) and len(frames) > 0:
|
||||
first_item = frames[0]
|
||||
|
||||
if hasattr(first_item, "final_output_type"):
|
||||
if first_item.final_output_type != "image":
|
||||
raise ValueError(
|
||||
f"Unexpected output type '{first_item.final_output_type}', expected 'image' for video generation."
|
||||
)
|
||||
|
||||
if hasattr(first_item, "is_pipeline_output") and first_item.is_pipeline_output:
|
||||
inner_output = first_item.request_output
|
||||
if isinstance(inner_output, OmniRequestOutput) and hasattr(inner_output, "images"):
|
||||
frames = inner_output.images[0] if inner_output.images else None
|
||||
if frames is None:
|
||||
raise ValueError("No video frames found in output.")
|
||||
elif hasattr(first_item, "images") and first_item.images:
|
||||
frames = first_item.images
|
||||
else:
|
||||
raise ValueError("No video frames found in OmniRequestOutput.")
|
||||
|
||||
# Unwrap batch list from postprocess_video: [numpy(T,H,W,C)] -> numpy(T,H,W,C)
|
||||
if isinstance(frames, list) and len(frames) > 0 and isinstance(frames[0], np.ndarray):
|
||||
frames = frames[0]
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
try:
|
||||
from diffusers.utils import export_to_video
|
||||
except ImportError:
|
||||
raise ImportError("diffusers is required for export_to_video.")
|
||||
|
||||
if isinstance(frames, torch.Tensor):
|
||||
video_tensor = frames.detach().cpu()
|
||||
if video_tensor.dim() == 5:
|
||||
if video_tensor.shape[1] in (3, 4):
|
||||
video_tensor = video_tensor[0].permute(1, 2, 3, 0)
|
||||
else:
|
||||
video_tensor = video_tensor[0]
|
||||
elif video_tensor.dim() == 4 and video_tensor.shape[0] in (3, 4):
|
||||
video_tensor = video_tensor.permute(1, 2, 3, 0)
|
||||
if video_tensor.is_floating_point():
|
||||
video_tensor = video_tensor.clamp(-1, 1) * 0.5 + 0.5
|
||||
video_array = video_tensor.float().numpy()
|
||||
elif isinstance(frames, np.ndarray):
|
||||
video_array = frames
|
||||
else:
|
||||
video_array = frames
|
||||
|
||||
if isinstance(video_array, np.ndarray) and video_array.ndim == 5:
|
||||
video_array = video_array[0]
|
||||
|
||||
export_to_video(video_array, str(output_path), fps=args.fps)
|
||||
print(f"Saved generated video to {output_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -74,6 +74,12 @@ def parse_args():
|
||||
parser.add_argument("--sys-type", type=str, default=None, help="Override system prompt type.")
|
||||
parser.add_argument("--deploy-config", type=str, default=None, help="Custom deploy YAML path.")
|
||||
parser.add_argument("--stage-configs-path", type=str, default=None, help="Custom legacy stage config YAML path.")
|
||||
parser.add_argument(
|
||||
"--stream",
|
||||
action="store_true",
|
||||
default=False,
|
||||
help="Enable streaming CoT display; print AR text token-by-token in real time.",
|
||||
)
|
||||
parser.add_argument("--log-stats", action="store_true", default=False)
|
||||
parser.add_argument("--init-timeout", type=int, default=300, help="Initialization timeout in seconds.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
@@ -222,6 +228,8 @@ def main():
|
||||
prompt_dict["modalities"] = ["text"]
|
||||
formatted_prompts.append(prompt_dict)
|
||||
|
||||
from vllm.sampling_params import RequestOutputKind
|
||||
|
||||
params_list = list(omni.default_sampling_params_list)
|
||||
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
@@ -253,6 +261,10 @@ def main():
|
||||
sp.width = args.width
|
||||
elif hasattr(sp, "stop_token_ids"):
|
||||
sp.stop_token_ids = ar_stop_token_ids
|
||||
# When --stream is set, request DELTA output from the AR stage
|
||||
# so we can display CoT text token-by-token in real time.
|
||||
if args.stream and hasattr(sp, "output_kind"):
|
||||
sp.output_kind = RequestOutputKind.DELTA
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print("HunyuanImage-3.0 Generation Configuration:")
|
||||
@@ -281,22 +293,32 @@ def main():
|
||||
print(f" Prompts: {prompts}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
omni_outputs = list(omni.generate(prompts=formatted_prompts, sampling_params_list=params_list))
|
||||
# When --stream is set, print AR CoT text token-by-token in real time.
|
||||
# Otherwise, collect and print the full AR text once when stage 0 finishes.
|
||||
omni_outputs = omni.generate(
|
||||
prompts=formatted_prompts,
|
||||
sampling_params_list=params_list,
|
||||
py_generator=True,
|
||||
use_tqdm=False,
|
||||
)
|
||||
img_idx = 0
|
||||
for req_output in omni_outputs:
|
||||
ro = getattr(req_output, "request_output", None)
|
||||
txt = ""
|
||||
if ro and getattr(ro, "outputs", None):
|
||||
txt = "".join(getattr(o, "text", "") or "" for o in ro.outputs)
|
||||
if not txt:
|
||||
ar_text = getattr(req_output, "custom_output", {}).get("ar_generated_text")
|
||||
if isinstance(ar_text, list):
|
||||
txt = "\n".join(text for text in ar_text if text)
|
||||
else:
|
||||
txt = ar_text or ""
|
||||
if txt:
|
||||
print(f"[Output] Text:\n{txt}")
|
||||
stage_id = getattr(req_output, "stage_id", None)
|
||||
|
||||
# AR stage text — each CompletionOutput.text is already a delta when
|
||||
# output_kind=DELTA, so we can print it directly (matching the pattern
|
||||
# in serving_chat.py).
|
||||
if stage_id == 0 and ro and getattr(ro, "outputs", None):
|
||||
for o in ro.outputs:
|
||||
text = getattr(o, "text", "") or ""
|
||||
if text:
|
||||
print(text, end="", flush=True)
|
||||
# Non-streaming: one shot with full text — emit a trailing newline.
|
||||
if not args.stream:
|
||||
print(flush=True)
|
||||
|
||||
# Collect images from diffusion stage
|
||||
images = getattr(req_output, "images", None)
|
||||
if not images and ro and hasattr(ro, "images"):
|
||||
images = ro.images
|
||||
@@ -304,7 +326,7 @@ def main():
|
||||
for j, img in enumerate(images):
|
||||
save_path = os.path.join(args.output, f"output_{img_idx}_{j}.png")
|
||||
img.save(save_path)
|
||||
print(f"[Output] Saved image to {save_path}")
|
||||
print(f"\n[Output] Saved image to {save_path}")
|
||||
img_idx += 1
|
||||
|
||||
|
||||
|
||||
@@ -97,9 +97,16 @@ import torch
|
||||
from PIL import Image
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.model_extras import (
|
||||
build_image_to_image_prompt,
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
should_init_extra_args_for_non_diffusion_stages,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
@@ -187,6 +194,12 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument(
|
||||
"--guidance-scale-2", type=float, default=None, help="image guidance scale for image-to-image generation."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-args",
|
||||
type=parse_profiler_config,
|
||||
default=None,
|
||||
help="JSON object copied to OmniDiffusionSamplingParams.extra_args, e.g. '{\"cfg_text_scale\": 4.0}'.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
type=str,
|
||||
@@ -468,6 +481,8 @@ def main():
|
||||
enable_diffusion_pipeline_profiler=args.enable_diffusion_pipeline_profiler,
|
||||
profiler_config=args.profiler_config,
|
||||
)
|
||||
model_class_name = get_model_class_name(omni)
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
print("Pipeline loaded")
|
||||
|
||||
profiler_enabled = args.profiler_config is not None
|
||||
@@ -478,6 +493,8 @@ def main():
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Cache backend: {args.cache_backend if args.cache_backend else 'None (no acceleration)'}")
|
||||
if args.height is not None or args.width is not None:
|
||||
print(f" Output size: {args.width or 'auto'}x{args.height or 'auto'}")
|
||||
if isinstance(input_image, list):
|
||||
print(f" Number of input images: {len(input_image)}")
|
||||
for idx, img in enumerate(input_image):
|
||||
@@ -487,7 +504,6 @@ def main():
|
||||
print(
|
||||
f" Parallel configuration: ulysses_degree={args.ulysses_degree}, ring_degree={args.ring_degree}, cfg_parallel_size={args.cfg_parallel_size}, tensor_parallel_size={args.tensor_parallel_size}, enable_expert_parallel: {args.enable_expert_parallel}"
|
||||
)
|
||||
print(f" Output image size: {args.width or 'default'}x{args.height or 'default'}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
@@ -496,10 +512,20 @@ def main():
|
||||
print("[Profiler] Starting profiling...")
|
||||
omni.start_profile()
|
||||
|
||||
# Generate edited image
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
width=args.width,
|
||||
prompt_dict = build_image_to_image_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
input_image=input_image,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
)
|
||||
|
||||
extra_args_from_cli = dict(args.extra_args or {})
|
||||
if args.negative_prompt is not None:
|
||||
extra_args_from_cli.setdefault("negative_prompt", args.negative_prompt)
|
||||
|
||||
diffusion_params = OmniDiffusionSamplingParams(
|
||||
generator=generator,
|
||||
true_cfg_scale=args.cfg_scale,
|
||||
guidance_scale=args.guidance_scale,
|
||||
@@ -507,17 +533,44 @@ def main():
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
num_outputs_per_prompt=args.num_outputs_per_prompt,
|
||||
layers=args.layers,
|
||||
resolution=args.resolution,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
)
|
||||
if args.resolution:
|
||||
sampling_params.resolution = args.resolution
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"multi_modal_data": {"image": input_image},
|
||||
},
|
||||
sampling_params,
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(
|
||||
diffusion_params,
|
||||
declared_extra_body_params,
|
||||
extra_args_from_cli,
|
||||
)
|
||||
else:
|
||||
diffusion_params.extra_args.update({k: v for k, v in extra_args_from_cli.items() if v is not None})
|
||||
|
||||
# Build per-stage sampling params for multi-stage models
|
||||
init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(
|
||||
model_class_name,
|
||||
)
|
||||
defaults = list(omni.default_sampling_params_list or [])
|
||||
sampling_params_list = [clone_sampling_params(p) for p in defaults]
|
||||
if not sampling_params_list:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
diffusion_replaced = False
|
||||
for idx, params in enumerate(sampling_params_list):
|
||||
if isinstance(params, OmniDiffusionSamplingParams):
|
||||
merged_extra = dict(getattr(params, "extra_args", {}) or {})
|
||||
merged_extra.update(diffusion_params.extra_args)
|
||||
diffusion_params.extra_args = merged_extra
|
||||
sampling_params_list[idx] = diffusion_params
|
||||
diffusion_replaced = True
|
||||
elif init_non_diffusion and hasattr(params, "extra_args"):
|
||||
if params.extra_args is None:
|
||||
params.extra_args = {}
|
||||
|
||||
if not diffusion_replaced and len(sampling_params_list) == 1:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -544,17 +597,16 @@ def main():
|
||||
if not outputs:
|
||||
raise ValueError("No output generated from omni.generate()")
|
||||
|
||||
# Extract images from OmniRequestOutput
|
||||
# omni.generate() returns list[OmniRequestOutput], extract images from request_output.images
|
||||
first_output = outputs[0]
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
images = None
|
||||
for output in outputs:
|
||||
images = getattr(output, "images", None)
|
||||
if images:
|
||||
break
|
||||
req_out = getattr(output, "request_output", None)
|
||||
images = getattr(req_out, "images", None) if req_out is not None else None
|
||||
if images:
|
||||
break
|
||||
|
||||
req_out = first_output.request_output
|
||||
if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images' key")
|
||||
|
||||
images = req_out.images
|
||||
if not images:
|
||||
raise ValueError("No images found in request_output")
|
||||
|
||||
|
||||
@@ -1,6 +1,7 @@
|
||||
# Image-To-Video
|
||||
|
||||
This example demonstrates how to generate videos from images using Wan2.2 Image-to-Video models with vLLM-Omni's offline inference API.
|
||||
This shared example generates videos from images with VACE, Wan2.2, LTX2,
|
||||
HunyuanVideo-1.5, Cosmos3, and other compatible pipelines.
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -48,19 +49,38 @@ python image_to_video.py \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
# Cosmos3 bundles example frames under assets/ (any RGB image works too):
|
||||
python image_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--image /path/to/Cosmos3-Nano/assets/example_i2v_input.jpg \
|
||||
--prompt "The scene comes to life with smooth, natural motion." \
|
||||
--negative-prompt "blurry, distorted, low quality" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}' \
|
||||
--output cosmos3_i2v.mp4
|
||||
```
|
||||
|
||||
Key arguments:
|
||||
|
||||
- `--model`: Model ID (I2V-A14B for MoE, TI2V-5B for unified T2V+I2V).
|
||||
- `--image`: Path to input image (required).
|
||||
- `--image`: Path to the first-frame or source image.
|
||||
- `--last-image`: Optional last-frame condition for models such as VACE.
|
||||
- `--mask-image`: Optional inpainting mask. White pixels are regenerated and black pixels are preserved.
|
||||
- `--reference-image`: Optional reference image; repeat it to provide multiple references.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3.
|
||||
- `--prompt`: Text description of desired motion/animation.
|
||||
- `--height/--width`: Output resolution (auto-calculated from image if not set). Dimensions should be multiples of 16.
|
||||
- `--num-frames`: Number of frames (default 81).
|
||||
- `--num-frames`: Number of frames (default: model-specific — Wan 81, LTX2 121, Cosmos3 189).
|
||||
- `--guidance-scale` and `--guidance-scale-high`: CFG scale (applied to low/high-noise stages for MoE).
|
||||
- `--negative-prompt`: Optional list of artifacts to suppress.
|
||||
- `--boundary-ratio`: Boundary split ratio for two-stage MoE models.
|
||||
- `--flow-shift`: Scheduler flow shift (5.0 for 720p, 12.0 for 480p).
|
||||
- `--flow-shift`: Scheduler flow shift (default: model-specific — Wan/LTX2 5.0, Cosmos3 10.0).
|
||||
- `--sample-solver`: Wan2.2 sampling solver. Use `unipc` for the default multistep solver, or `euler` for Lightning/Distill checkpoints.
|
||||
- `--num-inference-steps`: Number of denoising steps (default 50).
|
||||
- `--num-inference-steps`: Number of denoising steps (default: model-specific — Wan 50, LTX2 40, Cosmos3 35).
|
||||
- `--fps`: Frames per second for the saved MP4 (requires `diffusers` export_to_video).
|
||||
- `--output`: Path to save the generated video.
|
||||
- `--vae-use-slicing`: Enable VAE slicing for memory optimization.
|
||||
@@ -76,5 +96,98 @@ Key arguments:
|
||||
|
||||
> ℹ️ If you encounter OOM errors, try using `--vae-use-slicing` and `--vae-use-tiling` to reduce memory usage.
|
||||
|
||||
## Wan2.1 VACE Conditional Tasks
|
||||
|
||||
The shared script selects the VACE conditioning structure from the media inputs.
|
||||
No explicit mode parameter is required: the script constructs the source video,
|
||||
mask, or reference images consumed by the VACE pipeline.
|
||||
Download the same Hugging Face assets used by the original VACE example:
|
||||
|
||||
```bash
|
||||
wget -O astronaut.jpg https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/astronaut.jpg
|
||||
wget -O vace_first_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_first_frame.png
|
||||
wget -O vace_last_frame.png https://huggingface.co/datasets/huggingface/documentation-images/resolve/main/diffusers/flf2v_input_last_frame.png
|
||||
```
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_i2v_output.mp4
|
||||
```
|
||||
|
||||
### Video-to-Last-Frame (V2LF)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--last-image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked, otherworldly egg on the surface of the moon" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_v2lf_output.mp4
|
||||
```
|
||||
|
||||
### First-Last-Frame-to-Video (FLF2V)
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image vace_first_frame.png \
|
||||
--last-image vace_last_frame.png \
|
||||
--prompt "CG animation style, a small blue bird takes off from a branch and lands on another branch" \
|
||||
--seed 42 --height 512 --width 512 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_flf2v_output.mp4
|
||||
```
|
||||
|
||||
### Inpainting
|
||||
|
||||
Create a mask matching the original VACE example: a 160-pixel-wide white
|
||||
vertical stripe marks the region to regenerate.
|
||||
|
||||
```bash
|
||||
python - <<'PY'
|
||||
from PIL import Image
|
||||
|
||||
mask = Image.new("L", (832, 480), 0)
|
||||
mask.paste(255, (336, 0, 496, 480))
|
||||
mask.save("vace_center_mask.png")
|
||||
PY
|
||||
```
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--image astronaut.jpg \
|
||||
--mask-image vace_center_mask.png \
|
||||
--prompt "Shrek, the ogre, walks out of a building in a happy mood" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_inpaint_output.mp4
|
||||
```
|
||||
|
||||
### Reference-to-Video (R2V)
|
||||
|
||||
Repeat `--reference-image` to provide more than one reference image.
|
||||
|
||||
```bash
|
||||
python image_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--reference-image astronaut.jpg \
|
||||
--prompt "Camera slowly zooms out from the character walking in a garden" \
|
||||
--seed 42 --height 480 --width 832 --num-frames 81 \
|
||||
--num-inference-steps 30 --guidance-scale 5.0 --flow-shift 5.0 \
|
||||
--vae-use-tiling --output vace_r2v_output.mp4
|
||||
```
|
||||
|
||||
The VACE T2V command is documented in the shared
|
||||
[`text_to_video.py`](../text_to_video/text_to_video.md#wan21-vace-t2v) example.
|
||||
|
||||
For Wan2.2 LightX2V-converted local Diffusers directories and related LoRA
|
||||
assets, see the [LoRA guide](../../../docs/user_guide/diffusion/lora.md#wan22-lightx2v-offline-assembly).
|
||||
@@ -2,13 +2,14 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, or HunyuanVideo-1.5.
|
||||
Image-to-Video generation example using Wan2.2 I2V/TI2V models, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE.
|
||||
|
||||
Supports:
|
||||
- Wan2.2-I2V-A14B-Diffusers: MoE model with CLIP image encoder
|
||||
- Wan2.2-TI2V-5B-Diffusers: Unified T2V+I2V model (dense 5B)
|
||||
- LTX2 image-to-video pipeline
|
||||
- HunyuanVideo-1.5 I2V: SigLIP + VAE dual image conditioning
|
||||
- Wan2.1 VACE: first/last-frame, inpainting, and reference conditioning
|
||||
|
||||
Usage:
|
||||
# Wan I2V-A14B (MoE)
|
||||
@@ -30,9 +31,16 @@ Usage:
|
||||
python image_to_video.py --model hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_i2v \
|
||||
--image input.jpg --prompt "A cat playing with yarn" \
|
||||
--flow-shift 5.0 --guidance-scale 6.0
|
||||
|
||||
# Cosmos3 I2V (image conditioning)
|
||||
python image_to_video.py --model nvidia/Cosmos3-Nano \
|
||||
--image input.jpg --prompt "The scene comes to life with smooth, natural motion." \
|
||||
--num-frames 189 --num-inference-steps 35 --guidance-scale 6.0 --fps 24 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false}'
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -43,39 +51,59 @@ import PIL.Image
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import (
|
||||
build_image_to_video_prompt,
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
)
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate a video from an image (Wan2.2, LTX2, HunyuanVideo-1.5).")
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a video from one or more images (Wan2.2, LTX2, HunyuanVideo-1.5, or Wan2.1 VACE)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.2-I2V-A14B-Diffusers",
|
||||
help="Diffusers I2V model ID or local path (Wan2.2 or HunyuanVideo-1.5).",
|
||||
help="Diffusers I2V model ID or local path (Wan2.2, HunyuanVideo-1.5, or Wan2.1 VACE).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model-class-name",
|
||||
default=None,
|
||||
help="Override model class name (e.g., LTX2ImageToVideoPipeline).",
|
||||
)
|
||||
parser.add_argument("--image", required=True, help="Path to input image.")
|
||||
parser.add_argument("--image", help="Path to the first-frame or source image.")
|
||||
parser.add_argument("--last-image", help="Path to a last-frame condition (used by models such as VACE).")
|
||||
parser.add_argument("--mask-image", help="Path to an inpainting mask (used by models such as VACE).")
|
||||
parser.add_argument(
|
||||
"--reference-image",
|
||||
action="append",
|
||||
default=None,
|
||||
help="Path to a reference image. Repeat to provide multiple references.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="", help="Text prompt describing the desired motion.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG scale.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
"--guidance-scale-high", type=float, default=None, help="Optional separate CFG for high-noise (MoE only)."
|
||||
)
|
||||
@@ -83,8 +111,10 @@ def parse_args() -> argparse.Namespace:
|
||||
"--height", type=int, default=None, help="Video height (auto-calculated from image if not set)."
|
||||
)
|
||||
parser.add_argument("--width", type=int, default=None, help="Video width (auto-calculated from image if not set).")
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=50, help="Sampling steps.")
|
||||
parser.add_argument("--num-frames", type=int, default=None, help="Number of frames. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
"--num-inference-steps", type=int, default=None, help="Sampling steps. Default: model-specific."
|
||||
)
|
||||
parser.add_argument("--boundary-ratio", type=float, default=0.875, help="Boundary split ratio for MoE models.")
|
||||
parser.add_argument(
|
||||
"--frame-rate",
|
||||
@@ -93,7 +123,10 @@ def parse_args() -> argparse.Namespace:
|
||||
help="Optional generation frame rate (used by models like LTX2). Defaults to --fps.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--flow-shift", type=float, default=5.0, help="Scheduler flow_shift (5.0 for 720p, 12.0 for 480p)."
|
||||
"--flow-shift",
|
||||
type=float,
|
||||
default=None,
|
||||
help="Scheduler flow_shift. Default: model-specific (Wan 5.0, Cosmos3 10.0).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sample-solver",
|
||||
@@ -244,6 +277,19 @@ def parse_args() -> argparse.Namespace:
|
||||
default=None,
|
||||
help='JSON profiler config for torch/cuda profiling, e.g. \'{"profiler":"torch","torch_profiler_dir":"./perf"}\'.',
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=functools.partial(parse_json_object, flag_name="--extra-body"),
|
||||
default=None,
|
||||
help=(
|
||||
"Model-specific generation params as a JSON object. Keys are filtered "
|
||||
"against the model's declared extra_body_params (see vllm_omni/model_extras), "
|
||||
"so unknown keys for the chosen model are silently dropped. "
|
||||
'Cosmos3 V2V example: \'{"condition_frame_indexes_vision": [0, 1], '
|
||||
'"condition_video_keep": "first", "flow_shift": 10.0, '
|
||||
'"max_sequence_length": 4096, "guardrails": false}\'.'
|
||||
),
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
@@ -267,31 +313,63 @@ def main():
|
||||
model_name = str(args.model).lower() if args.model is not None else ""
|
||||
model_class_name = args.model_class_name
|
||||
is_ltx2 = "ltx2" in model_name or (model_class_name and "ltx2" in model_class_name.lower())
|
||||
is_cosmos = "cosmos" in model_name or (model_class_name is not None and "cosmos" in model_class_name.lower())
|
||||
if model_class_name is None and is_ltx2:
|
||||
model_class_name = "LTX2ImageToVideoPipeline"
|
||||
|
||||
# Load input image
|
||||
image = PIL.Image.open(args.image).convert("RGB")
|
||||
image = PIL.Image.open(args.image).convert("RGB") if args.image else None
|
||||
last_image = PIL.Image.open(args.last_image).convert("RGB") if args.last_image else None
|
||||
mask_image = PIL.Image.open(args.mask_image).convert("L") if args.mask_image else None
|
||||
reference_images = (
|
||||
[PIL.Image.open(path).convert("RGB") for path in args.reference_image] if args.reference_image else None
|
||||
)
|
||||
dimension_image = image or last_image or (reference_images[0] if reference_images else None)
|
||||
if dimension_image is None:
|
||||
raise ValueError("Provide --image, --last-image, or at least one --reference-image.")
|
||||
|
||||
fps = args.fps if args.fps is not None else (24 if is_ltx2 else 16)
|
||||
# Per-model generation defaults, applied only when the matching flag is omitted.
|
||||
# Cosmos3 would otherwise silently inherit the Wan2.2 defaults (wrong size/steps/shift).
|
||||
if is_cosmos:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = (
|
||||
24,
|
||||
6.0,
|
||||
189,
|
||||
35,
|
||||
10.0,
|
||||
1280 * 720,
|
||||
16,
|
||||
)
|
||||
elif is_ltx2:
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 24, 4.0, 121, 40, 5.0, 512 * 768, 32
|
||||
else: # Wan2.2 / HunyuanVideo-1.5
|
||||
d_fps, d_guidance, d_num_frames, d_steps, d_flow_shift, d_max_area, d_mod = 16, 5.0, 81, 50, 5.0, 480 * 832, 16
|
||||
|
||||
fps = args.fps if args.fps is not None else d_fps
|
||||
frame_rate = args.frame_rate if args.frame_rate is not None else float(fps)
|
||||
guidance_scale = args.guidance_scale if args.guidance_scale is not None else (4.0 if is_ltx2 else 5.0)
|
||||
num_frames = args.num_frames if args.num_frames is not None else (121 if is_ltx2 else 81)
|
||||
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else (40 if is_ltx2 else 50)
|
||||
guidance_scale = args.guidance_scale if args.guidance_scale is not None else d_guidance
|
||||
num_frames = args.num_frames if args.num_frames is not None else d_num_frames
|
||||
num_inference_steps = args.num_inference_steps if args.num_inference_steps is not None else d_steps
|
||||
flow_shift = args.flow_shift if args.flow_shift is not None else d_flow_shift
|
||||
|
||||
# Calculate dimensions if not provided
|
||||
# Calculate dimensions if not provided (model-aware max area).
|
||||
height = args.height
|
||||
width = args.width
|
||||
if height is None or width is None:
|
||||
# Default to 480P area for Wan2.2 I2V, 512x768 area for LTX2
|
||||
max_area = 512 * 768 if is_ltx2 else 480 * 832
|
||||
mod_value = 32 if is_ltx2 else 16
|
||||
calc_height, calc_width = calculate_dimensions(image, max_area=max_area, mod_value=mod_value)
|
||||
calc_height, calc_width = calculate_dimensions(dimension_image, max_area=d_max_area, mod_value=d_mod)
|
||||
height = height or calc_height
|
||||
width = width or calc_width
|
||||
|
||||
# Resize image to target dimensions
|
||||
image = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
media_inputs: dict[str, Any] = {}
|
||||
if image is not None:
|
||||
media_inputs["image"] = image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if last_image is not None:
|
||||
media_inputs["last_image"] = last_image.resize((width, height), PIL.Image.Resampling.LANCZOS)
|
||||
if mask_image is not None:
|
||||
media_inputs["mask"] = mask_image.resize((width, height), PIL.Image.Resampling.NEAREST)
|
||||
if reference_images is not None:
|
||||
media_inputs["reference_images"] = [
|
||||
reference.resize((width, height), PIL.Image.Resampling.LANCZOS) for reference in reference_images
|
||||
]
|
||||
|
||||
# Configure cache based on backend type
|
||||
cache_config = None
|
||||
@@ -342,7 +420,7 @@ def main():
|
||||
vae_use_slicing=args.vae_use_slicing,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
flow_shift=args.flow_shift,
|
||||
flow_shift=flow_shift,
|
||||
diffusion_kv_cache_dtype=args.diffusion_kv_cache_dtype,
|
||||
diffusion_kv_cache_skip_steps=args.diffusion_kv_cache_skip_steps,
|
||||
diffusion_kv_cache_skip_layers=args.diffusion_kv_cache_skip_layers,
|
||||
@@ -357,7 +435,13 @@ def main():
|
||||
)
|
||||
if args.quantization is not None:
|
||||
omni_kwargs["quantization"] = args.quantization
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni) or model_class_name
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -367,8 +451,8 @@ def main():
|
||||
print(f"\n{'=' * 60}")
|
||||
print("Generation Configuration:")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Inference steps: {args.num_inference_steps}")
|
||||
print(f" Frames: {args.num_frames}")
|
||||
print(f" Inference steps: {num_inference_steps}")
|
||||
print(f" Frames: {num_frames}")
|
||||
print(f" Solver: {args.sample_solver}")
|
||||
print(f" diffusion_kv_cache_dtype(config): {args.diffusion_kv_cache_dtype}")
|
||||
print(f" diffusion_kv_cache_skip_steps(config): {args.diffusion_kv_cache_skip_steps}")
|
||||
@@ -381,30 +465,43 @@ def main():
|
||||
print(f" Video size: {args.width}x{args.height}")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
prompt_dict = build_image_to_video_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
media_inputs=media_inputs,
|
||||
height=height,
|
||||
width=width,
|
||||
num_frames=num_frames,
|
||||
)
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
guidance_scale=guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
extra_args={
|
||||
"sample_solver": args.sample_solver,
|
||||
"flow_shift": flow_shift,
|
||||
},
|
||||
)
|
||||
|
||||
# Route model-specific knobs through extra_body, filtered against the model's
|
||||
# declared extra_body_params. Models without a declaration only forward explicit
|
||||
# --extra-body JSON (e.g. Cosmos3 V2V's condition_frame_indexes_vision).
|
||||
extra_body = dict(args.extra_body or {})
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
|
||||
elif extra_body:
|
||||
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
# omni.generate() returns Generator[OmniRequestOutput, None, None]
|
||||
frames = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"multi_modal_data": {"image": image},
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=height,
|
||||
width=width,
|
||||
generator=generator,
|
||||
guidance_scale=guidance_scale,
|
||||
guidance_scale_2=args.guidance_scale_high,
|
||||
boundary_ratio=args.boundary_ratio,
|
||||
num_inference_steps=num_inference_steps,
|
||||
num_frames=num_frames,
|
||||
frame_rate=frame_rate,
|
||||
extra_args={
|
||||
"sample_solver": args.sample_solver,
|
||||
"flow_shift": args.flow_shift,
|
||||
},
|
||||
),
|
||||
)
|
||||
frames = omni.generate(prompt_dict, sampling_params)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
|
||||
@@ -1,72 +0,0 @@
|
||||
# MagiHuman Generation
|
||||
|
||||
MagiHuman is an advanced, omni-modality model that generates both high-quality video and lip-synced audio from a text prompt.
|
||||
|
||||
Because MagiHuman is a very large model featuring a powerful DiT MoE backbone and a ~9B parameter T5Gemma text encoder, it natively supports **Tensor Parallelism (TP)** in vLLM-Omni to run efficiently across multi-GPU setups, reducing device memory bottlenecks.
|
||||
|
||||
## Setup
|
||||
|
||||
### Install MagiCompiler (recommended)
|
||||
|
||||
MagiHuman relies on [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for custom-op registration used by the DiT attention kernels. While the pipeline can fall back to stub implementations, installing MagiCompiler is **strongly recommended** for correct behaviour.
|
||||
|
||||
```bash
|
||||
# Clone the repo
|
||||
git clone https://github.com/SandAI-org/MagiCompiler.git
|
||||
cd MagiCompiler
|
||||
|
||||
# System dependencies (optional, for FX graph visualization; Debian/Ubuntu)
|
||||
sudo apt update && sudo apt install -y graphviz
|
||||
|
||||
# Python dependencies
|
||||
pip install -r requirements.txt
|
||||
|
||||
# Install MagiCompiler
|
||||
pip install . # end users (recommended)
|
||||
# pip install -e . # developers (editable install)
|
||||
```
|
||||
|
||||
### Hardware requirements
|
||||
|
||||
Ensure your hardware has enough VRAM. For a standard node with 80GB GPUs, running with `--tensor-parallel-size 4` is recommended to shard both the MoE weights and the T5Gemma text encoder across 4 GPUs, reducing the per-GPU peak VRAM overhead significantly (by roughly ~13.5GB per GPU compared to single-device inference).
|
||||
|
||||
Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) for further details on allocating memory.
|
||||
|
||||
## Run Examples
|
||||
|
||||
Get into the example folder:
|
||||
```bash
|
||||
cd examples/offline_inference/magi_human
|
||||
```
|
||||
|
||||
### End-to-End Generation (Text to Video+Audio)
|
||||
|
||||
Generate a video with synchronized speech natively generated by the model.
|
||||
|
||||
```bash
|
||||
python end2end.py \
|
||||
--model /proj-tango-pvc/users/zhipeng.wang/workspace/models/daVinci-MagiHuman \
|
||||
--prompt "A young woman with long, wavy golden blonde hair..." \
|
||||
--tensor-parallel-size 4 \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
## Common Parameters
|
||||
|
||||
| Parameter | Default | Description |
|
||||
|-----------|---------|-------------|
|
||||
| `--model` | *(Required)* | Local model path or HuggingFace ID |
|
||||
| `--prompt` | *(built-in demo prompt)* | Highly detailed text prompt dictating visual look and dialogue text |
|
||||
| `--tensor-parallel-size` | `4` | Tensor parallelism size (Number of GPUs) |
|
||||
| `--height` | `256` | Initial resolution height |
|
||||
| `--width` | `448` | Initial resolution width |
|
||||
| `--num-inference-steps` | `8` | Denoising steps |
|
||||
| `--seed` | `52` | Random seed |
|
||||
| `--output` | `output_magihuman.mp4` | Output video with audio path |
|
||||
|
||||
## Example materials
|
||||
|
||||
??? abstract "end2end.py"
|
||||
``````py
|
||||
--8<-- "examples/offline_inference/magi_human/end2end.py"
|
||||
``````
|
||||
@@ -1,122 +0,0 @@
|
||||
import argparse
|
||||
|
||||
from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="End-to-end inference script for MagiHuman.")
|
||||
parser.add_argument("--model", type=str, required=True, help="Path or ID of the MagiHuman model.")
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
type=str,
|
||||
default="",
|
||||
help="Text prompt containing visual description, dialogue, and background sound.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tensor-parallel-size", "-tp", type=int, default=4, help="Tensor parallel size (number of GPUs)."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output", type=str, default="output_magihuman.mp4", help="Path to save the generated mp4 file."
|
||||
)
|
||||
parser.add_argument("--height", type=int, default=256, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=448, help="Video width.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=8, help="Number of denoising steps.")
|
||||
parser.add_argument("--seed", type=int, default=52, help="Random seed for generation.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
|
||||
print(f"Initializing MagiHuman pipeline with TP={args.tensor_parallel_size}...")
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
init_timeout=1200,
|
||||
tensor_parallel_size=args.tensor_parallel_size,
|
||||
devices=list(range(args.tensor_parallel_size)),
|
||||
)
|
||||
|
||||
prompt = args.prompt
|
||||
if not prompt:
|
||||
prompt = (
|
||||
"A young woman with long, wavy golden blonde hair and bright blue eyes, "
|
||||
"wearing a fitted ivory silk blouse with a delicate lace collar, sits "
|
||||
"stationary in front of a softly lit, blurred warm-toned interior. Her "
|
||||
"overall disposition is warm, composed, and gently confident. The camera "
|
||||
"holds a static medium close-up, framing her from the shoulders up, "
|
||||
"with shallow depth of field keeping her face in sharp focus. Soft "
|
||||
"directional key light falls from the upper left, casting a gentle "
|
||||
"highlight along her cheekbone and nose bridge. She draws a quiet breath, "
|
||||
"the levator labii superiors relaxing as her lips part. She speaks in "
|
||||
"clear, warm, unhurried American English: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\" "
|
||||
"Her jaw descends smoothly on each stressed syllable; the orbicularis oris "
|
||||
"shapes each vowel with precision. A faint, genuine smile engages the "
|
||||
"zygomaticus major, lifting her lip corners fractionally. Her brows rest "
|
||||
"in a soft, neutral arch throughout. She maintains steady, forward-facing "
|
||||
"eye contact. Head position remains level; no torso displacement occurs.\n\n"
|
||||
"Dialogue:\n"
|
||||
"<Young blonde woman, American English>: "
|
||||
"\"The most beautiful things in life aren't things at all — "
|
||||
"they're moments, feelings, and the people who make you feel truly alive.\"\n\n"
|
||||
"Background Sound:\n"
|
||||
"<Soft, warm indoor ambience with a faint distant piano melody>"
|
||||
)
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"seconds": 5,
|
||||
"sr_height": 1080,
|
||||
"sr_width": 1920,
|
||||
"sr_num_inference_steps": 5,
|
||||
},
|
||||
)
|
||||
|
||||
print(f"Generating with prompt: {prompt[:80]}...")
|
||||
outputs = omni.generate(
|
||||
prompts=[prompt],
|
||||
sampling_params_list=[sampling_params],
|
||||
)
|
||||
|
||||
print(f"Generation complete. Output type: {type(outputs)}")
|
||||
if outputs:
|
||||
first = outputs[0]
|
||||
|
||||
if hasattr(first, "images") and first.images:
|
||||
video_frames = first.images[0]
|
||||
print(f"Video frames: shape={video_frames.shape}, dtype={video_frames.dtype}")
|
||||
|
||||
audio_waveform = None
|
||||
mm = first.multimodal_output or {}
|
||||
if mm:
|
||||
audio_waveform = mm.get("audio")
|
||||
if audio_waveform is not None:
|
||||
print(f"Audio waveform: shape={audio_waveform.shape}, dtype={audio_waveform.dtype}")
|
||||
|
||||
output_fps = float(mm.get("fps", 25))
|
||||
output_sr = int(mm.get("audio_sample_rate", 24000))
|
||||
print(f"Using fps={output_fps}, audio_sample_rate={output_sr} from model output")
|
||||
|
||||
video_bytes = mux_video_audio_bytes(
|
||||
video_frames,
|
||||
audio_waveform,
|
||||
fps=output_fps,
|
||||
audio_sample_rate=output_sr,
|
||||
)
|
||||
with open(args.output, "wb") as f:
|
||||
f.write(video_bytes)
|
||||
print(f"Saved MP4 ({len(video_bytes)} bytes) to {args.output}")
|
||||
print("SUCCESS: MagiHuman pipeline generation completed.")
|
||||
else:
|
||||
print("WARNING: No outputs returned.")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -145,7 +145,7 @@ python end2end.py \
|
||||
## Reducing GPU Memory Usage
|
||||
|
||||
For hardware with limited VRAM, enable **module-wise CPU offload** with
|
||||
`--enable-cpu-offload`. The pipeline implements `SupportsModuleOffload`, so the
|
||||
`--enable-cpu-offload`. The pipeline implements `SupportsComponentDiscovery`, so the
|
||||
vision encoder (`vision_model`) and the Qwen3 LLM (`language_model`) are
|
||||
swapped between CPU and GPU on demand:
|
||||
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import functools
|
||||
import json
|
||||
import time
|
||||
from pathlib import Path
|
||||
@@ -10,10 +11,18 @@ from typing import Any
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import logger
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.entrypoints.openai.stage_params import clone_sampling_params
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.lora.request import LoRARequest
|
||||
from vllm_omni.lora.utils import stable_lora_int_id
|
||||
from vllm_omni.model_extras import (
|
||||
build_text_to_image_prompt,
|
||||
get_extra_body_params,
|
||||
get_model_class_name,
|
||||
should_init_extra_args_for_non_diffusion_stages,
|
||||
)
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
@@ -30,16 +39,20 @@ def is_nextstep_model(model_name: str) -> bool:
|
||||
return False
|
||||
|
||||
|
||||
def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
def parse_json_object(value: str, flag_name: str = "argument") -> dict[str, Any]:
|
||||
"""Parse a CLI value as a JSON object, attributing errors to ``flag_name``."""
|
||||
try:
|
||||
config = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--profiler-config must be valid JSON: {e}") from e
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be valid JSON: {e}") from e
|
||||
if not isinstance(config, dict):
|
||||
raise argparse.ArgumentTypeError("--profiler-config must be a JSON object")
|
||||
raise argparse.ArgumentTypeError(f"{flag_name} must be a JSON object")
|
||||
return config
|
||||
|
||||
|
||||
parse_profiler_config = functools.partial(parse_json_object, flag_name="--profiler-config")
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Generate an image with supported diffusion models.")
|
||||
parser.add_argument(
|
||||
@@ -259,6 +272,19 @@ def parse_args() -> argparse.Namespace:
|
||||
action="store_true",
|
||||
help="[NextStep-1.1 only] Apply layer normalization to sampled tokens.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=functools.partial(parse_json_object, flag_name="--extra-body"),
|
||||
default=None,
|
||||
help=(
|
||||
"Model-specific generation params as a JSON object, e.g. "
|
||||
'\'{"timestep_shift": 3.0, "cfg_text_scale": 4.0, "cfg_interval": [0.4, 1.0]}\'. '
|
||||
"Each key is filtered against the model's declared extra_body_params "
|
||||
"(see vllm_omni/model_extras), so unknown keys for the chosen model are "
|
||||
"silently dropped. Values here take precedence over the equivalent "
|
||||
"model-specific flags above."
|
||||
),
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
@@ -404,7 +430,13 @@ def main():
|
||||
if use_nextstep:
|
||||
# NextStep-1.1 requires explicit pipeline class
|
||||
omni_kwargs["model_class_name"] = "NextStep11Pipeline"
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni)
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -445,34 +477,79 @@ def main():
|
||||
)
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
extra_args = {
|
||||
|
||||
prompt_dict = build_text_to_image_prompt(
|
||||
model_class_name=model_class_name,
|
||||
prompt=args.prompt,
|
||||
negative_prompt=args.negative_prompt,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
)
|
||||
|
||||
diffusion_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
seed=args.seed,
|
||||
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_images_per_prompt,
|
||||
)
|
||||
|
||||
# Base layer: backward-compatible model-specific flags. New model params should
|
||||
# instead be declared in vllm_omni/model_extras and passed via --extra-body, so
|
||||
# this dict does not need to grow per-model.
|
||||
user_extra = {
|
||||
"cfg_scale": args.cfg_scale,
|
||||
"cfg_text_scale": args.cfg_scale,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
"timestep_shift": args.timesteps_shift,
|
||||
"timesteps_shift": args.timesteps_shift,
|
||||
"cfg_schedule": args.cfg_schedule,
|
||||
"use_norm": args.use_norm,
|
||||
"use_system_prompt": args.use_system_prompt,
|
||||
"system_prompt": args.system_prompt,
|
||||
}
|
||||
if lora_request:
|
||||
extra_args["lora_request"] = lora_request
|
||||
extra_args["lora_scale"] = args.lora_scale
|
||||
# Override layer: generic JSON passthrough wins over the flags above. Keys are
|
||||
# still filtered against the model's declared extra_body_params downstream.
|
||||
if args.extra_body:
|
||||
user_extra.update(args.extra_body)
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(diffusion_params, declared_extra_body_params, user_extra)
|
||||
else:
|
||||
diffusion_params.extra_args.update({k: v for k, v in user_extra.items() if v is not None})
|
||||
|
||||
outputs = omni.generate(
|
||||
{
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
},
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
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_images_per_prompt,
|
||||
extra_args=extra_args,
|
||||
),
|
||||
if lora_request:
|
||||
diffusion_params.extra_args["lora_request"] = lora_request
|
||||
diffusion_params.extra_args["lora_scale"] = args.lora_scale
|
||||
|
||||
# Build per-stage sampling params for multi-stage models (e.g. BAGEL),
|
||||
# or wrap single diffusion params for single-stage models.
|
||||
init_non_diffusion = should_init_extra_args_for_non_diffusion_stages(
|
||||
model_class_name,
|
||||
)
|
||||
defaults = list(omni.default_sampling_params_list or [])
|
||||
sampling_params_list = [clone_sampling_params(p) for p in defaults]
|
||||
if not sampling_params_list:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
diffusion_replaced = False
|
||||
for idx, params in enumerate(sampling_params_list):
|
||||
if isinstance(params, OmniDiffusionSamplingParams):
|
||||
sampling_params_list[idx] = diffusion_params
|
||||
diffusion_replaced = True
|
||||
elif init_non_diffusion and hasattr(params, "extra_args"):
|
||||
if params.extra_args is None:
|
||||
params.extra_args = {}
|
||||
if args.seed is not None and hasattr(params, "seed"):
|
||||
params.seed = args.seed
|
||||
|
||||
if not diffusion_replaced and len(sampling_params_list) == 1:
|
||||
sampling_params_list = [diffusion_params]
|
||||
|
||||
outputs = omni.generate(prompt_dict, sampling_params_list=sampling_params_list)
|
||||
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
@@ -502,15 +579,16 @@ def main():
|
||||
raise ValueError("No output generated from omni.generate()")
|
||||
logger.info(f"Outputs: {outputs}")
|
||||
|
||||
first_output = outputs[0]
|
||||
if not hasattr(first_output, "request_output") or not first_output.request_output:
|
||||
raise ValueError("No request_output found in OmniRequestOutput")
|
||||
images = None
|
||||
for output in outputs:
|
||||
images = getattr(output, "images", None)
|
||||
if images:
|
||||
break
|
||||
req_out = getattr(output, "request_output", None)
|
||||
images = getattr(req_out, "images", None) if req_out is not None else None
|
||||
if images:
|
||||
break
|
||||
|
||||
req_out = first_output.request_output
|
||||
if not hasattr(req_out, "images"):
|
||||
raise ValueError("Invalid request_output structure or missing 'images'.")
|
||||
|
||||
images = req_out.images
|
||||
if not images:
|
||||
raise ValueError("No images found in request_output")
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ list of supported architectures across all modalities, see
|
||||
| OmniVoice | `k2-fsa/OmniVoice` | 2 (gen + dec) | ✓ | — | voice design, language hint | 24 kHz |
|
||||
| Qwen3-TTS | `Qwen/Qwen3-TTS-12Hz-1.7B-{CustomVoice,VoiceDesign,Base}` | 2 (talker + code2wav) | ✓ (Base) | ✓ | 3 task variants | 24 kHz |
|
||||
| VoxCPM2 | `openbmb/VoxCPM2` | single (native AR) | ✓ | ✓ (online) | continuation | 48 kHz |
|
||||
| IndexTTS-2 | `IndexTeam/IndexTTS-2` | 2 (AR talker + S2Mel DiT + BigVGAN) | ✓ (required) | — | emotion control (`--emo-audio`, `--emo-text`, `--emo-vector`) | 22.05 kHz |
|
||||
| Voxtral TTS | `mistralai/Voxtral-4B-TTS-2603` | varies | ✓ | ✓ | voice presets | 24 kHz |
|
||||
|
||||
## Common Quick Start
|
||||
@@ -369,6 +370,20 @@ python examples/offline_inference/text_to_speech/qwen3_tts/end2end.py \
|
||||
```
|
||||
Streaming requires `async_chunk: true` in the stage config.
|
||||
|
||||
### Word Timestamps
|
||||
Generate a WAV offline and a JSON sidecar with word-level timestamps from
|
||||
`Qwen/Qwen3-ForcedAligner-0.6B`:
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B \
|
||||
--text "Hello world." \
|
||||
--output-dir /tmp/qwen3_tts_timestamps
|
||||
```
|
||||
The script writes `qwen3_tts_word_timestamps.wav` and
|
||||
`qwen3_tts_word_timestamps.json`. On machines without a local CUDA toolkit,
|
||||
set `VLLM_USE_FLASHINFER_SAMPLER=0` to avoid FlashInfer sampler JIT.
|
||||
|
||||
### Batched decoding
|
||||
The Code2Wav stage supports batched decoding through the SpeechTokenizer. Pass multiple prompts via `--txt-prompts` and set `--batch-size` accordingly. To raise `max_num_seqs` on either stage, point `--stage-configs-path` at a stage configs YAML with the desired values (see `vllm_omni/model_executor/stage_configs/` for templates):
|
||||
```bash
|
||||
@@ -422,6 +437,50 @@ Streaming is exposed through the online OpenAI Speech API (`stream=true`). See [
|
||||
|
||||
---
|
||||
|
||||
## IndexTTS-2
|
||||
|
||||
2-stage TTS pipeline (GPT AR talker + S2Mel CFM DiT + BigVGAN vocoder) at 22.05 kHz. Every request requires reference audio for zero-shot voice cloning. Supports emotion conditioning via audio, text, or 8-dim vector.
|
||||
|
||||
### Quick start
|
||||
```bash
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/reference.wav
|
||||
```
|
||||
|
||||
### Emotion control
|
||||
```bash
|
||||
# Emotion from reference audio
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
|
||||
# Emotion from 8-dim vector (happy angry sad afraid disgusted melancholy surprised calm)
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-vector 1.0 0.0 0.0 0.0 0.0 0.0 0.0 0.0
|
||||
|
||||
# Emotion from text description
|
||||
python examples/offline_inference/text_to_speech/indextts2/end2end.py \
|
||||
--model IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-text "happy and excited"
|
||||
```
|
||||
|
||||
### Notes
|
||||
- `--ref-audio` is **required** — IndexTTS-2 does not support text-only synthesis.
|
||||
- Stage 0 (AR Talker): GPT-2 generates mel codes from text + reference audio.
|
||||
- Stage 1 (S2Mel + BigVGAN): CFM DiT converts mel codes to waveform at 22.05 kHz.
|
||||
- Deploy config: `vllm_omni/deploy/indextts2.yaml`. Stage 1 runs with `enforce_eager: true` (DiT has dynamic shapes).
|
||||
|
||||
---
|
||||
|
||||
## Voxtral TTS
|
||||
|
||||
Voxtral-4B-TTS (Mistral). Has its own argparse surface; uses voice presets and the `mistral_common` `SpeechRequest` protocol.
|
||||
@@ -462,3 +521,100 @@ Available voice presets are listed on the HF model card (`mistralai/Voxtral-4B-T
|
||||
- `--num-prompts N` replicates the prompt for performance measurement.
|
||||
- `--concurrency M` requires `--streaming` and must evenly divide `--num-prompts`.
|
||||
- Run `--help` for the full argument surface.
|
||||
|
||||
---
|
||||
|
||||
## SoulX-Singer
|
||||
|
||||
Singing voice synthesis (SVS) and conversion (SVC) at 24 kHz. Script: `soulxsinger/end2end.py`. Deploy: `vllm_omni/deploy/soulxsinger_svs.yaml` or `soulxsinger_svc.yaml`.
|
||||
|
||||
### Prerequisites
|
||||
|
||||
Download DiT and preprocess weights, then set up separate SVS / SVC view directories. Copy `soulxsinger/utils/phoneme/phone_set.json` from upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) into the model weights dir as `phoneme/phone_set.json` — HuggingFace does not ship it.
|
||||
|
||||
```bash
|
||||
# 1. DiT weights
|
||||
export BASE=path/to/SoulX-Singer
|
||||
export PREPROCESS=path/to/SoulX-Singer-Preprocess
|
||||
export SVC_DIR=path/to/SoulX-Singer-svc
|
||||
|
||||
huggingface-cli download Soul-AILab/SoulX-Singer --local-dir "$BASE"
|
||||
|
||||
# 2. Preprocess weights (required)
|
||||
huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess --local-dir "$PREPROCESS"
|
||||
export SOULX_PREPROCESS_WEIGHTS_DIR="$PREPROCESS"
|
||||
|
||||
# 3. SVS / SVC view directories
|
||||
mkdir -p "$SVC_DIR"
|
||||
cp $BASE/{config.yaml,README.md,assets} $SVC_DIR
|
||||
mv $BASE/model-svc.pt $SVC_DIR/model-svc.pt
|
||||
|
||||
cat > "$BASE/config.json" <<'EOF'
|
||||
{
|
||||
"model_type": "soulxsinger",
|
||||
"architectures": ["SoulXSingerPipeline"],
|
||||
"max_num_seqs": 1
|
||||
}
|
||||
EOF
|
||||
|
||||
cat > "$SVC_DIR/config.json" <<'EOF'
|
||||
{
|
||||
"model_type": "soulxsinger",
|
||||
"architectures": ["SoulXSingerSVCPipeline"],
|
||||
"max_num_seqs": 1
|
||||
}
|
||||
EOF
|
||||
```
|
||||
|
||||
`config.yaml` hyper-parameters live under `$BASE`; each view's `config.json` `architectures` field is the single source of truth for SVS vs SVC. Point `--model` at the matching directory (`$BASE` for SVS, `$SVC_DIR` for SVC). Deploy YAML is chosen automatically from `config.json`; optional `--svs` / `--svc` only assert the mode matches.
|
||||
|
||||
**Online preprocess** is the default: pass `--prompt-audio` and `--target-audio`, and the worker runs vocal separation, F0, and (for SVS) lyrics/MIDI before DiT. Install only what your run needs:
|
||||
|
||||
```bash
|
||||
pip install "BS-RoFormer" # vocal sep + F0 on GPU — SVS and SVC
|
||||
```
|
||||
|
||||
Mandarin SVS also needs FunASR and Chinese G2P; `ffmpeg` must be on `PATH`:
|
||||
|
||||
```bash
|
||||
# install optional dependencies:
|
||||
pip install -e ".[soulx-svs]"
|
||||
```
|
||||
|
||||
English SVS adds NeMo ASR and NLTK data; pass `--language English`:
|
||||
|
||||
```bash
|
||||
pip install "nemo_toolkit[asr]==2.6.1" lhotse==1.32.2
|
||||
python -c "import nltk; nltk.download('cmudict'); nltk.download('averaged_perceptron_tagger_eng')"
|
||||
```
|
||||
|
||||
**Precomputed metadata** is the alternative: pass both `--prompt-metadata-path` and `--target-metadata-path` and skip online ASR/ROSVOT — none of the packages above are required. JSON can be produced by integrated preprocess on a prior run, or by upstream [SoulX-Singer](https://github.com/Soul-AILab/SoulX-Singer) `preprocess/` scripts if you prefer to run that outside vLLM-Omni.
|
||||
|
||||
### Quick start
|
||||
|
||||
```bash
|
||||
# SVS — default demo audio: tests/assets/soulxsinger/zh_prompt.mp3 + music.mp3
|
||||
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
|
||||
--model "$BASE" \
|
||||
--preprocess-weights-dir "$PREPROCESS" \
|
||||
--control score \
|
||||
--num-inference-steps 32 \
|
||||
-o output.wav
|
||||
|
||||
python examples/offline_inference/text_to_speech/soulxsinger/end2end.py \
|
||||
--model "$SVC_DIR" \
|
||||
--preprocess-weights-dir "$PREPROCESS" \
|
||||
--svc \
|
||||
--num-inference-steps 32 \
|
||||
-o output_svc.wav
|
||||
```
|
||||
|
||||
`SOULX_PREPROCESS_WEIGHTS_DIR` makes `--preprocess-weights-dir` optional. Long SVS targets are handled in one request. See `end2end.py --help` for `--pitch-shift`, `--vocal-sep`, `--auto-shift`, and language/control options.
|
||||
|
||||
### Notes
|
||||
|
||||
- Output: 24 kHz mono WAV; batch only.
|
||||
- Defaults match upstream: `--guidance-scale 3.0`, `--seed 42`, `--auto-shift` on.
|
||||
- SVS `--control`: `score` or `melody`. MIDI / lyric QC: upstream `midi_editor` only.
|
||||
|
||||
---
|
||||
@@ -0,0 +1,187 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Offline inference example for IndexTTS2 via vLLM-Omni.
|
||||
|
||||
Two-stage pipeline: GPT AR (Stage 0) → S2Mel + BigVGAN (Stage 1).
|
||||
Output is 22050 Hz mono WAV.
|
||||
|
||||
Usage:
|
||||
python end2end.py \
|
||||
--model /path/to/IndexTeam/IndexTTS-2 \
|
||||
--text "你好,这是一个语音合成测试。" \
|
||||
--ref-audio /path/to/ref.wav
|
||||
|
||||
# With emotion audio:
|
||||
python end2end.py \
|
||||
--model /path/to/IndexTeam/IndexTTS-2 \
|
||||
--text "今天天气真好!" \
|
||||
--ref-audio /path/to/ref.wav \
|
||||
--emo-audio /path/to/happy.wav
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from vllm import SamplingParams
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
|
||||
from vllm_omni import Omni # noqa: E402
|
||||
from vllm_omni.model_executor.models.indextts2.prompt_utils import ( # noqa: E402
|
||||
build_indextts2_prefill_prompt_ids,
|
||||
)
|
||||
|
||||
|
||||
def build_request(
|
||||
model: str,
|
||||
text: str,
|
||||
ref_audio_path: str | None = None,
|
||||
emo_audio_path: str | None = None,
|
||||
emo_text: str | None = None,
|
||||
emo_vector: list[float] | None = None,
|
||||
emo_alpha: float | None = None,
|
||||
use_emo_text: bool = False,
|
||||
use_random: bool = False,
|
||||
) -> dict:
|
||||
additional: dict = {"text": [text]}
|
||||
if ref_audio_path:
|
||||
additional["voice"] = [str(ref_audio_path)]
|
||||
if emo_audio_path:
|
||||
additional["emo_audio"] = [str(emo_audio_path)]
|
||||
if emo_text:
|
||||
additional["emo_text"] = [emo_text]
|
||||
if emo_vector is not None:
|
||||
additional["emo_vector"] = [emo_vector]
|
||||
if emo_alpha is not None:
|
||||
additional["emo_alpha"] = [emo_alpha]
|
||||
if use_emo_text:
|
||||
additional["use_emo_text"] = [True]
|
||||
if use_random:
|
||||
additional["use_random"] = [True]
|
||||
return {
|
||||
"prompt_token_ids": build_indextts2_prefill_prompt_ids(model, text),
|
||||
"additional_information": additional,
|
||||
}
|
||||
|
||||
|
||||
def save_audio(waveform: torch.Tensor, path: str, sample_rate: int = 22050) -> None:
|
||||
audio_np = waveform.float().numpy()
|
||||
sf.write(path, audio_np, sample_rate)
|
||||
print(f" Saved {path} ({audio_np.shape}, {sample_rate} Hz)")
|
||||
|
||||
|
||||
def extract_audio(mm: dict) -> tuple[torch.Tensor | None, int]:
|
||||
audio = mm.get("audio")
|
||||
if audio is None:
|
||||
audio = mm.get("model_outputs")
|
||||
if isinstance(audio, list):
|
||||
chunks = [chunk.reshape(-1) for chunk in audio if isinstance(chunk, torch.Tensor) and chunk.numel() > 0]
|
||||
audio = torch.cat(chunks, dim=0) if chunks else None
|
||||
|
||||
sr_val = mm.get("sr")
|
||||
if isinstance(sr_val, list):
|
||||
sr_val = sr_val[-1] if sr_val else None
|
||||
if hasattr(sr_val, "item"):
|
||||
sample_rate = int(sr_val.item())
|
||||
else:
|
||||
sample_rate = int(sr_val) if sr_val is not None else 22050
|
||||
|
||||
return audio if isinstance(audio, torch.Tensor) else None, sample_rate
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
deploy_config=args.deploy_config,
|
||||
stage_init_timeout=args.stage_init_timeout,
|
||||
)
|
||||
|
||||
# Stage 0: GPT AR (text → mel codes)
|
||||
gpt_sampling = SamplingParams(
|
||||
temperature=0.8,
|
||||
top_p=0.8,
|
||||
top_k=30,
|
||||
max_tokens=1500,
|
||||
repetition_penalty=10.0,
|
||||
stop_token_ids=[8193],
|
||||
seed=args.seed if args.seed is not None else 42,
|
||||
detokenize=False,
|
||||
)
|
||||
# Stage 1: S2Mel + BigVGAN (non-AR, params mostly ignored)
|
||||
s2mel_sampling = SamplingParams(
|
||||
temperature=0.0,
|
||||
max_tokens=65536,
|
||||
detokenize=True,
|
||||
)
|
||||
sampling_params = [gpt_sampling, s2mel_sampling]
|
||||
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
print(f"Synthesizing: {args.text!r}")
|
||||
if args.ref_audio:
|
||||
print(f" ref_audio: {args.ref_audio}")
|
||||
inputs = build_request(
|
||||
model=args.model,
|
||||
text=args.text,
|
||||
ref_audio_path=args.ref_audio,
|
||||
emo_audio_path=args.emo_audio,
|
||||
emo_text=args.emo_text,
|
||||
emo_vector=args.emo_vector,
|
||||
emo_alpha=args.emo_alpha,
|
||||
use_emo_text=args.use_emo_text,
|
||||
use_random=args.use_random,
|
||||
)
|
||||
|
||||
for i, omni_out in enumerate(omni.generate(inputs, sampling_params_list=sampling_params)):
|
||||
mm = omni_out.multimodal_output
|
||||
if not mm:
|
||||
print(f" [req {i}] No multimodal output.")
|
||||
continue
|
||||
audio, sr = extract_audio(mm)
|
||||
if audio is None:
|
||||
print(f" [req {i}] No waveform in multimodal_output.")
|
||||
continue
|
||||
out_path = str(output_dir / f"output_{i}.wav")
|
||||
save_audio(audio.cpu(), out_path, sr)
|
||||
|
||||
print("Done.")
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="IndexTTS2 offline inference")
|
||||
parser.add_argument("--model", required=True, help="HF model path for IndexTTS2.")
|
||||
parser.add_argument("--text", default="你好,这是IndexTTS2语音合成测试。")
|
||||
parser.add_argument("--ref-audio", required=True, help="Reference audio for voice cloning.")
|
||||
parser.add_argument("--emo-audio", default=None, help="Emotion reference audio.")
|
||||
parser.add_argument("--emo-text", default=None, help="Emotion description text.")
|
||||
parser.add_argument(
|
||||
"--emo-vector",
|
||||
type=float,
|
||||
nargs=8,
|
||||
default=None,
|
||||
help="8-dim emotion vector: happy angry sad afraid disgusted melancholic surprised calm.",
|
||||
)
|
||||
parser.add_argument("--emo-alpha", type=float, default=None, help="Emotion weight in [0, 1].")
|
||||
parser.add_argument("--use-emo-text", action="store_true", help="Infer emotion vector from emo-text or text.")
|
||||
parser.add_argument("--use-random", action="store_true", help="Use random emotion prototypes.")
|
||||
parser.add_argument("--seed", type=int, default=None)
|
||||
parser.add_argument(
|
||||
"--output-dir",
|
||||
default=os.path.join(
|
||||
os.environ.get("XDG_CACHE_HOME", os.path.join(os.path.expanduser("~"), ".cache")),
|
||||
"indextts2_output",
|
||||
),
|
||||
)
|
||||
parser.add_argument("--deploy-config", default=None)
|
||||
parser.add_argument("--stage-init-timeout", type=int, default=600)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -1,11 +1,12 @@
|
||||
"""Offline inference example for MOSS-TTS variants.
|
||||
|
||||
Covers all five models:
|
||||
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
|
||||
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
|
||||
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
|
||||
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
|
||||
- OpenMOSS-Team/MOSS-VoiceGenerator(1.7B, voice design)
|
||||
Covers six models:
|
||||
- OpenMOSS-Team/MOSS-TTS (8B, general TTS)
|
||||
- OpenMOSS-Team/MOSS-TTSD-v1.0 (8B, dialogue TTS)
|
||||
- OpenMOSS-Team/MOSS-SoundEffect (8B, sound effect synthesis)
|
||||
- OpenMOSS-Team/MOSS-TTS-Realtime (1.7B, low-latency streaming)
|
||||
- OpenMOSS-Team/MOSS-VoiceGenerator (1.7B, voice design)
|
||||
- OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5 (n_vq=12, 48 kHz stereo)
|
||||
|
||||
Usage
|
||||
-----
|
||||
@@ -210,12 +211,18 @@ def run_tts(args: argparse.Namespace) -> None:
|
||||
"OpenMOSS-Team/MOSS-SoundEffect": "moss_sound_effect",
|
||||
"OpenMOSS-Team/MOSS-TTS-Realtime": "moss_tts_realtime",
|
||||
"OpenMOSS-Team/MOSS-VoiceGenerator": "moss_voice_generator",
|
||||
"OpenMOSS-Team/MOSS-TTS-Local-Transformer-v1.5": "moss_tts_local",
|
||||
}
|
||||
deploy_config = deploy_map.get(args.model, "moss_tts")
|
||||
|
||||
is_sound_effect = "SoundEffect" in args.model
|
||||
is_voice_gen = "VoiceGenerator" in args.model
|
||||
is_realtime = "Realtime" in args.model
|
||||
# Local-v1.5 ships its own AutoProcessor (processor_config.json +
|
||||
# processing_moss_tts.py), so it reuses the same _build_unified_codes
|
||||
# path as the Delay variants below — just a different n_vq/codec.
|
||||
is_local = "Local" in args.model
|
||||
output_sr = 48000 if is_local else 24000
|
||||
|
||||
ref_audio = None
|
||||
builder_kwargs: dict = {}
|
||||
@@ -314,27 +321,33 @@ def run_tts(args: argparse.Namespace) -> None:
|
||||
wav = mm.get("model_outputs")
|
||||
if wav is None:
|
||||
continue
|
||||
# Local-v1.5's codec emits stereo chunks shaped (C, T); keep the
|
||||
# channel axis intact (don't flatten) so concatenation below stays
|
||||
# along the time axis. Mono variants stay (T,) as before.
|
||||
if isinstance(wav, list):
|
||||
non_empty = [w.reshape(-1) for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
|
||||
non_empty = [w for w in wav if isinstance(w, torch.Tensor) and w.numel() > 0]
|
||||
if non_empty and t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.extend(non_empty)
|
||||
elif isinstance(wav, torch.Tensor) and wav.numel() > 0:
|
||||
if t_first_audio is None:
|
||||
t_first_audio = _time.monotonic()
|
||||
chunks.append(wav.reshape(-1))
|
||||
chunks.append(wav)
|
||||
|
||||
t_gen_end = _time.monotonic()
|
||||
if not chunks:
|
||||
print("No audio in output — check model loading logs.")
|
||||
return
|
||||
|
||||
wav_t = torch.cat(chunks).float().cpu()
|
||||
sf.write(args.output, wav_t.numpy(), 24000)
|
||||
print(f"Saved {len(wav_t) / 24000:.2f}s of audio to {args.output}")
|
||||
cat_dim = -1 if chunks[0].dim() == 2 else 0
|
||||
wav_t = torch.cat(chunks, dim=cat_dim).float().cpu()
|
||||
n_samples = wav_t.shape[-1]
|
||||
# soundfile wants (T,) mono or (T, C) multi-channel.
|
||||
sf.write(args.output, (wav_t.transpose(0, 1) if wav_t.dim() == 2 else wav_t).numpy(), output_sr)
|
||||
print(f"Saved {n_samples / output_sr:.2f}s of audio to {args.output}")
|
||||
ttfa_ms = (t_first_audio - t_gen_start) * 1000.0 if t_first_audio is not None else float("nan")
|
||||
gen_total_ms = (t_gen_end - t_gen_start) * 1000.0
|
||||
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={len(wav_t) / 24000:.2f}")
|
||||
print(f"TTFA_MS={ttfa_ms:.1f} GEN_TOTAL_MS={gen_total_ms:.1f} AUDIO_S={n_samples / output_sr:.2f}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
|
||||
@@ -0,0 +1,172 @@
|
||||
"""Offline Qwen3-TTS example with word-level timestamps.
|
||||
|
||||
This demo runs Qwen3-TTS offline, aligns the synthesized audio in memory with
|
||||
the same shared forced aligner used by the streaming server path, then writes
|
||||
the WAV and JSON sidecar:
|
||||
|
||||
python examples/offline_inference/text_to_speech/qwen3_tts/word_timestamps.py \
|
||||
--model Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice \
|
||||
--forced-aligner Qwen/Qwen3-ForcedAligner-0.6B
|
||||
|
||||
On machines without a local CUDA toolkit, set ``VLLM_USE_FLASHINFER_SAMPLER=0``
|
||||
to avoid FlashInfer sampler JIT during warmup.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import asyncio
|
||||
import json
|
||||
import os
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
os.environ.setdefault("VLLM_WORKER_MULTIPROC_METHOD", "spawn")
|
||||
# The in-process offline aligner consumes token_classify outputs. Keeping V1
|
||||
# multiprocessing off avoids msgspec serialization issues for token-wise logits.
|
||||
os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
import torch
|
||||
from end2end import _estimate_prompt_len
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni import Omni
|
||||
from vllm_omni.engine.arg_utils import nullify_stage_engine_defaults
|
||||
from vllm_omni.utils.forced_aligner import align, build_forced_aligner_config
|
||||
|
||||
|
||||
def _default_stage_config() -> str:
|
||||
repo_root = Path(__file__).resolve().parents[4]
|
||||
return str(repo_root / "vllm_omni" / "deploy" / "qwen3_tts.yaml")
|
||||
|
||||
|
||||
def _build_custom_voice_input(args: Any) -> dict[str, Any]:
|
||||
additional_information = {
|
||||
"task_type": ["CustomVoice"],
|
||||
"text": [args.text],
|
||||
"language": [args.language],
|
||||
"speaker": [args.speaker],
|
||||
"instruct": [args.instructions],
|
||||
"max_new_tokens": [args.max_new_tokens],
|
||||
}
|
||||
return {
|
||||
"prompt_token_ids": [0] * _estimate_prompt_len(additional_information, args.model),
|
||||
"additional_information": additional_information,
|
||||
}
|
||||
|
||||
|
||||
def _audio_tensor_and_sample_rate(mm: dict[str, Any]) -> tuple[torch.Tensor, int]:
|
||||
audio_data = mm["audio"]
|
||||
sr_raw = mm["sr"]
|
||||
sr_val = sr_raw[-1] if isinstance(sr_raw, list) and sr_raw else sr_raw
|
||||
sample_rate = int(sr_val.item()) if hasattr(sr_val, "item") else int(sr_val)
|
||||
audio_tensor = torch.cat(audio_data, dim=-1) if isinstance(audio_data, list) else audio_data
|
||||
return audio_tensor.float().detach().cpu().flatten(), sample_rate
|
||||
|
||||
|
||||
def _float_audio_to_pcm16_bytes(audio: torch.Tensor) -> bytes:
|
||||
samples = audio.numpy()
|
||||
pcm = (np.clip(samples, -1.0, 1.0) * 32767.0).astype("<i2")
|
||||
return pcm.tobytes()
|
||||
|
||||
|
||||
async def _run_alignment(args: Any, audio: torch.Tensor, sample_rate: int) -> list[dict[str, Any]] | None:
|
||||
aligner_config = build_forced_aligner_config(args)
|
||||
if aligner_config is None:
|
||||
raise ValueError("--forced-aligner or --forced-aligner-config is required")
|
||||
timestamps = await align(
|
||||
audio=_float_audio_to_pcm16_bytes(audio),
|
||||
text=args.text,
|
||||
sample_rate=sample_rate,
|
||||
config=aligner_config,
|
||||
language=args.language,
|
||||
)
|
||||
if timestamps is None:
|
||||
return None
|
||||
return [
|
||||
{
|
||||
"word": item.word,
|
||||
"start_ms": item.start_ms,
|
||||
"end_ms": item.end_ms,
|
||||
}
|
||||
for item in timestamps
|
||||
]
|
||||
|
||||
|
||||
def main(args: Any) -> None:
|
||||
output_dir = Path(args.output_dir)
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
omni_kwargs = vars(args).copy()
|
||||
for key in (
|
||||
"forced_aligner",
|
||||
"text",
|
||||
"language",
|
||||
"speaker",
|
||||
"instructions",
|
||||
"max_new_tokens",
|
||||
"output_dir",
|
||||
):
|
||||
omni_kwargs.pop(key, None)
|
||||
omni_kwargs["stage_configs_path"] = args.stage_configs_path or _default_stage_config()
|
||||
omni_kwargs["log_stats"] = args.log_stats
|
||||
omni = Omni(**omni_kwargs)
|
||||
|
||||
prompt = _build_custom_voice_input(args)
|
||||
final_output = None
|
||||
for stage_outputs in omni.generate([prompt]):
|
||||
final_output = stage_outputs.request_output
|
||||
if final_output is None:
|
||||
raise RuntimeError("Qwen3-TTS did not produce an output.")
|
||||
|
||||
mm = final_output.outputs[0].multimodal_output
|
||||
audio, sample_rate = _audio_tensor_and_sample_rate(mm)
|
||||
|
||||
wav_path = output_dir / "qwen3_tts_word_timestamps.wav"
|
||||
# Align the in-memory tensor, not the saved WAV. The WAV is only a demo
|
||||
# artifact so users can listen to the same audio referenced by the sidecar.
|
||||
timestamps = asyncio.run(_run_alignment(args, audio, sample_rate))
|
||||
sf.write(wav_path, audio.numpy(), samplerate=sample_rate, format="WAV")
|
||||
|
||||
sidecar = {
|
||||
"text": args.text,
|
||||
"sample_rate": sample_rate,
|
||||
"audio_path": str(wav_path),
|
||||
"timestamps": timestamps,
|
||||
}
|
||||
json_path = output_dir / "qwen3_tts_word_timestamps.json"
|
||||
json_path.write_text(json.dumps(sidecar, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
|
||||
|
||||
print(f"Saved audio: {wav_path}")
|
||||
print(f"Saved timestamps: {json_path}")
|
||||
print(json.dumps(sidecar["timestamps"], ensure_ascii=False, indent=2))
|
||||
|
||||
|
||||
def parse_args() -> Any:
|
||||
parser = FlexibleArgumentParser(description="Offline Qwen3-TTS word timestamp example")
|
||||
parser.add_argument("--model", default="Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice", help="Qwen3-TTS model path/name")
|
||||
parser.add_argument(
|
||||
"--forced-aligner",
|
||||
default=None,
|
||||
help="Qwen3 forced aligner model path/name",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--forced-aligner-config",
|
||||
default=None,
|
||||
help="Optional YAML file for forced aligner settings (incl. gpu_memory_utilization)",
|
||||
)
|
||||
parser.add_argument("--stage-configs-path", default=None, help="Qwen3-TTS deploy YAML")
|
||||
parser.add_argument("--text", default="Hello world.", help="Text to synthesize and align")
|
||||
parser.add_argument("--language", default="English", help="Qwen3-TTS language field")
|
||||
parser.add_argument("--speaker", default="Vivian", help="CustomVoice speaker name")
|
||||
parser.add_argument("--instructions", default="", help="Optional speaking style instruction")
|
||||
parser.add_argument("--max-new-tokens", type=int, default=2048, help="TTS max_new_tokens")
|
||||
parser.add_argument("--output-dir", default="output_audio", help="Directory for WAV and JSON sidecar")
|
||||
parser.add_argument("--log-stats", action="store_true", default=False, help="Enable vLLM-Omni stats logging")
|
||||
nullify_stage_engine_defaults(parser)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,162 @@
|
||||
"""Repeatable SoulX-Singer offline benchmark (preprocess + DiT, RTF + stage timings).
|
||||
|
||||
Usage:
|
||||
python benchmark.py --model /path/to/SoulX-Singer --svs \\
|
||||
--prompt-audio ../SoulX-Singer/example/audio/zh_prompt.mp3 \\
|
||||
--target-audio ../SoulX-Singer/example/audio/music.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o benchmark.wav
|
||||
|
||||
python benchmark.py ... --enable-diffusion-pipeline-profiler
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import statistics
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import soundfile as sf
|
||||
from end2end import (
|
||||
SVC_DEPLOY_CONFIG,
|
||||
SVS_DEPLOY_CONFIG,
|
||||
add_inference_args,
|
||||
build_sampling,
|
||||
extract_audio,
|
||||
resolve_preprocess_weights_dir,
|
||||
)
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
|
||||
|
||||
def _audio_duration_sec(outputs) -> float:
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
return audio_np.size / sr
|
||||
|
||||
|
||||
def _stage_durations_from_output(outputs) -> dict[str, float]:
|
||||
omni_out = outputs[0]
|
||||
durations = getattr(omni_out, "stage_durations", None) or {}
|
||||
if durations:
|
||||
return dict(durations)
|
||||
ro = getattr(omni_out, "request_output", None)
|
||||
if ro is not None:
|
||||
return dict(getattr(ro, "stage_durations", None) or {})
|
||||
return {}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer offline benchmark")
|
||||
add_inference_args(parser)
|
||||
parser.add_argument("--warmup", type=int, default=1, help="Warmup runs (excluded from stats)")
|
||||
parser.add_argument("--runs", type=int, default=3, help="Measured runs")
|
||||
parser.add_argument(
|
||||
"--enable-diffusion-pipeline-profiler",
|
||||
action="store_true",
|
||||
help="Enable low-overhead stage timing in DiffusionOutput.stage_durations",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--enforce-eager",
|
||||
action="store_true",
|
||||
help="Disable torch.compile on diff_estimator (eager baseline for A/B)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
type=str,
|
||||
default=None,
|
||||
choices=["bfloat16", "bf16", "float16", "fp16", "half"],
|
||||
help="DiT trunk dtype (default from deploy YAML). Use float16 to match upstream --fp16.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output",
|
||||
"-o",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Save WAV from the last measured run (after warmup)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--save-each-run",
|
||||
action="store_true",
|
||||
help="With --output, also write measured runs as <stem>_run1.wav, _run2.wav, ...",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
|
||||
deploy_config = args.deploy_config or str(SVS_DEPLOY_CONFIG if args.svs else SVC_DEPLOY_CONFIG)
|
||||
mode = "SVS" if args.svs else "SVC"
|
||||
|
||||
omni_kwargs: dict = {
|
||||
"model": args.model,
|
||||
"deploy_config": deploy_config,
|
||||
"async_chunk": False, # SoulX-Singer currently supports only batch mode (pseudo-streaming was stashed)
|
||||
}
|
||||
if args.enable_diffusion_pipeline_profiler:
|
||||
omni_kwargs["enable_diffusion_pipeline_profiler"] = True
|
||||
if args.enforce_eager:
|
||||
omni_kwargs["enforce_eager"] = True
|
||||
if args.dtype is not None:
|
||||
omni_kwargs["dtype"] = args.dtype
|
||||
|
||||
compile_mode = "eager" if args.enforce_eager else "torch.compile"
|
||||
dtype_label = args.dtype or "from deploy YAML"
|
||||
print(f"Loading SoulX-Singer {mode} from {args.model} [{compile_mode}, dtype={dtype_label}]")
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
kind = "svs" if args.svs else "svc"
|
||||
sampling = build_sampling(args, preprocess_weights_dir=preprocess_dir, kind=kind)
|
||||
prompt = {"prompt_token_ids": [0]}
|
||||
|
||||
latencies_ms: list[float] = []
|
||||
rtfs: list[float] = []
|
||||
last_stages: dict[str, float] = {}
|
||||
last_measured_outputs = None
|
||||
|
||||
total_runs = args.warmup + args.runs
|
||||
for run_idx in range(total_runs):
|
||||
is_warmup = run_idx < args.warmup
|
||||
label = "warmup" if is_warmup else f"run {run_idx - args.warmup + 1}/{args.runs}"
|
||||
t0 = time.perf_counter()
|
||||
outputs = list(omni.generate([prompt], sampling))
|
||||
elapsed_ms = (time.perf_counter() - t0) * 1000.0
|
||||
audio_sec = _audio_duration_sec(outputs)
|
||||
rtf = (elapsed_ms / 1000.0) / audio_sec if audio_sec > 0 else float("inf")
|
||||
print(f"[{label}] client={elapsed_ms:.1f} ms, audio={audio_sec:.2f}s, RTF={rtf:.3f}")
|
||||
if not is_warmup:
|
||||
latencies_ms.append(elapsed_ms)
|
||||
rtfs.append(rtf)
|
||||
last_stages = _stage_durations_from_output(outputs)
|
||||
last_measured_outputs = outputs
|
||||
if args.output and args.save_each_run:
|
||||
measured_idx = run_idx - args.warmup + 1
|
||||
out_path = Path(args.output)
|
||||
run_path = out_path.with_name(f"{out_path.stem}_run{measured_idx}{out_path.suffix}")
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
sf.write(str(run_path), audio_np, sr)
|
||||
print(f" saved {run_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
|
||||
|
||||
omni.close()
|
||||
|
||||
if args.output and last_measured_outputs is not None:
|
||||
audio_np, sr = extract_audio(last_measured_outputs)
|
||||
out_path = Path(args.output)
|
||||
out_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
sf.write(str(out_path), audio_np, sr)
|
||||
print(f"\nSaved last measured run → {out_path} ({sr} Hz, {audio_np.size / sr:.2f}s)")
|
||||
|
||||
if latencies_ms:
|
||||
print("\n=== Summary (measured runs) ===")
|
||||
print(f"client_ms: mean={statistics.mean(latencies_ms):.1f}, stdev={statistics.pstdev(latencies_ms):.1f}")
|
||||
print(f"RTF: mean={statistics.mean(rtfs):.3f}, stdev={statistics.pstdev(rtfs):.3f}")
|
||||
if last_stages:
|
||||
print("\nStage durations (last measured run):")
|
||||
for name, value in sorted(last_stages.items()):
|
||||
if name.endswith("_ms"):
|
||||
print(f" {name}: {value:.1f} ms")
|
||||
else:
|
||||
print(f" {name}: {value * 1000.0:.1f} ms")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,285 @@
|
||||
"""Offline SoulX-Singer SVS / SVC: single-stage DiT (preprocess inline).
|
||||
|
||||
Pass raw prompt/target audio; vLLM-Omni runs integrated preprocess then inference.
|
||||
|
||||
Usage:
|
||||
# SVS (default)
|
||||
python end2end.py --model /path/to/SoulX-Singer \\
|
||||
--prompt-audio path/to/prompt/audio.mp3 \\
|
||||
--target-audio path/to/target/audio.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o output.wav
|
||||
|
||||
# SVC (config.json must declare SoulXSingerSVCPipeline; --model points at SVC view)
|
||||
python end2end.py --model /path/to/SoulX-Singer-svc \\
|
||||
--prompt-audio path/to/prompt/audio.mp3 \\
|
||||
--target-audio path/to/target/audio.mp3 \\
|
||||
--preprocess-weights-dir /path/to/SoulX-Singer-Preprocess \\
|
||||
-o output_svc.wav
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import os
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
|
||||
from vllm_omni.diffusion.models.soulx_singer.utils import resolve_soulx_kind, validate_soulx_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
REPO_ROOT = Path(__file__).resolve().parents[4]
|
||||
SVS_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svs.yaml"
|
||||
SVC_DEPLOY_CONFIG = REPO_ROOT / "vllm_omni" / "deploy" / "soulxsinger_svc.yaml"
|
||||
DEFAULT_PROMPT_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "zh_prompt.mp3"
|
||||
DEFAULT_TARGET_AUDIO = REPO_ROOT / "tests" / "assets" / "soulxsinger" / "music.mp3"
|
||||
_SAMPLE_RATE = 24000
|
||||
|
||||
|
||||
def _require_paths(paths: dict[str, str | None]) -> None:
|
||||
missing = [name for name, p in paths.items() if not p or not os.path.isfile(p)]
|
||||
if missing:
|
||||
raise FileNotFoundError("Missing file(s): " + ", ".join(f"{k}={paths[k]!r}" for k in missing))
|
||||
|
||||
|
||||
def resolve_preprocess_weights_dir(explicit: str | None) -> str:
|
||||
if explicit:
|
||||
path = Path(explicit).expanduser().resolve()
|
||||
if not (path / "rmvpe" / "rmvpe.pt").is_file():
|
||||
raise FileNotFoundError(f"--preprocess-weights-dir must contain rmvpe/rmvpe.pt: {path}")
|
||||
return str(path)
|
||||
env = os.environ.get("SOULX_PREPROCESS_WEIGHTS_DIR")
|
||||
if env:
|
||||
path = Path(env).expanduser().resolve()
|
||||
if (path / "rmvpe" / "rmvpe.pt").is_file():
|
||||
return str(path)
|
||||
raise FileNotFoundError(
|
||||
"SoulX preprocess weights not found. Pass --preprocess-weights-dir or set "
|
||||
"SOULX_PREPROCESS_WEIGHTS_DIR to a directory containing rmvpe/rmvpe.pt "
|
||||
"(huggingface-cli download Soul-AILab/SoulX-Singer-Preprocess)."
|
||||
)
|
||||
|
||||
|
||||
def build_sampling(
|
||||
args: argparse.Namespace,
|
||||
*,
|
||||
preprocess_weights_dir: str,
|
||||
kind: str,
|
||||
) -> OmniDiffusionSamplingParams:
|
||||
extra: dict = {
|
||||
"prompt_audio": os.path.abspath(args.prompt_audio),
|
||||
"target_audio": os.path.abspath(args.target_audio),
|
||||
"preprocess_weights_dir": preprocess_weights_dir,
|
||||
"vocal_sep": args.vocal_sep,
|
||||
"auto_shift": args.auto_shift,
|
||||
"pitch_shift": args.pitch_shift,
|
||||
}
|
||||
if kind == "svs":
|
||||
extra["language"] = args.language
|
||||
extra["control"] = args.control
|
||||
if args.prompt_metadata_path:
|
||||
extra["prompt_metadata_path"] = os.path.abspath(args.prompt_metadata_path)
|
||||
if args.target_metadata_path:
|
||||
extra["target_metadata_path"] = os.path.abspath(args.target_metadata_path)
|
||||
if args.prompt_metadata_path or args.target_metadata_path:
|
||||
extra["audio_path"] = os.path.abspath(args.prompt_wav_path or args.prompt_audio)
|
||||
|
||||
extra = validate_soulx_extra_args(kind, extra)
|
||||
|
||||
return OmniDiffusionSamplingParams(
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
seed=args.seed,
|
||||
extra_args=extra,
|
||||
)
|
||||
|
||||
|
||||
def extract_audio(outputs) -> tuple[np.ndarray, int]:
|
||||
out = outputs[0]
|
||||
mm = getattr(out, "multimodal_output", None)
|
||||
if mm is None:
|
||||
ro = getattr(out, "request_output", None)
|
||||
if ro is not None:
|
||||
mm = getattr(ro, "multimodal_output", None)
|
||||
if mm is None and ro.outputs:
|
||||
mm = getattr(ro.outputs[0], "multimodal_output", None)
|
||||
if not mm or "audio" not in mm:
|
||||
raise RuntimeError("No audio in multimodal_output")
|
||||
|
||||
audio_val = mm["audio"]
|
||||
if isinstance(audio_val, list):
|
||||
chunks = [
|
||||
np.asarray(chunk.detach().cpu().float().numpy() if hasattr(chunk, "detach") else chunk).reshape(-1)
|
||||
for chunk in audio_val
|
||||
if chunk is not None
|
||||
]
|
||||
audio_np = np.concatenate(chunks, axis=0) if chunks else np.array([], dtype=np.float32)
|
||||
elif hasattr(audio_val, "cpu"):
|
||||
audio_np = audio_val.detach().cpu().numpy().squeeze()
|
||||
else:
|
||||
audio_np = np.asarray(audio_val, dtype=np.float32).squeeze()
|
||||
|
||||
sr_val = mm.get("sr", mm.get("audio_sample_rate", _SAMPLE_RATE))
|
||||
if isinstance(sr_val, list) and sr_val:
|
||||
sr_val = sr_val[-1]
|
||||
if hasattr(sr_val, "item"):
|
||||
sr_val = sr_val.item()
|
||||
return audio_np.astype(np.float32, copy=False), int(sr_val)
|
||||
|
||||
|
||||
def add_inference_args(parser: argparse.ArgumentParser) -> None:
|
||||
mode = parser.add_mutually_exclusive_group()
|
||||
mode.add_argument(
|
||||
"--svc",
|
||||
action="store_true",
|
||||
help="Assert SVC mode (must match config.json architectures)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
type=str,
|
||||
required=True,
|
||||
help="SoulX-Singer weight directory ($BASE for SVS, $SVC_DIR for SVC)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--deploy-config",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Override deploy YAML (default: soulxsinger_svs.yaml or soulxsinger_svc.yaml)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-audio",
|
||||
type=str,
|
||||
default=str(DEFAULT_PROMPT_AUDIO),
|
||||
help="Prompt / reference audio (mp3 or wav)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-audio",
|
||||
type=str,
|
||||
default=str(DEFAULT_TARGET_AUDIO),
|
||||
help="Target content audio (mp3 or wav)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-metadata-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS precomputed prompt metadata.json (skips online ASR/ROSVOT for prompt)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-metadata-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS precomputed target metadata.json (skips online ASR/ROSVOT for target)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--prompt-wav-path",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SVS prompt vocal wav for metadata processor when using --prompt-metadata-path (default: --prompt-audio)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preprocess-weights-dir",
|
||||
type=str,
|
||||
default=None,
|
||||
help="SoulX-Singer-Preprocess weights (or set SOULX_PREPROCESS_WEIGHTS_DIR)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--language",
|
||||
type=str,
|
||||
default="Mandarin",
|
||||
help="SVS lyrics language (Mandarin, Cantonese, English)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--control",
|
||||
type=str,
|
||||
default="melody",
|
||||
choices=["score", "melody"],
|
||||
help="SVS control mode",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--vocal-sep",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=False,
|
||||
help="Run vocal separation during preprocess",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--auto-shift",
|
||||
action=argparse.BooleanOptionalAction,
|
||||
default=True,
|
||||
help="Auto pitch shift to align target pitch with prompt",
|
||||
)
|
||||
parser.add_argument("--pitch-shift", type=int, default=0)
|
||||
parser.add_argument("--num-inference-steps", type=int, default=32)
|
||||
parser.add_argument("--guidance-scale", type=float, default=3.0)
|
||||
parser.add_argument(
|
||||
"--seed",
|
||||
type=int,
|
||||
default=42,
|
||||
help="Optional CFM RNG seed. Omit for non-deterministic sampling.",
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="SoulX-Singer offline SVS / SVC (preprocess + DiT)")
|
||||
add_inference_args(parser)
|
||||
parser.add_argument("--output", "-o", type=str, default="output.wav")
|
||||
args = parser.parse_args()
|
||||
return args
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
kind = resolve_soulx_kind(args.model)
|
||||
|
||||
if not (args.svc ^ (kind == "svs")):
|
||||
raise ValueError(f"mode mismatch: {args.model}/config.json (declares {kind!r}), --svc={args.svc}")
|
||||
|
||||
required = {
|
||||
"prompt_audio": args.prompt_audio,
|
||||
"target_audio": args.target_audio,
|
||||
}
|
||||
if kind == "svs" and bool(args.prompt_metadata_path) ^ bool(args.target_metadata_path):
|
||||
raise ValueError(
|
||||
"SVS precomputed metadata requires both --prompt-metadata-path and "
|
||||
"--target-metadata-path (or neither, for online preprocess)."
|
||||
)
|
||||
if args.prompt_metadata_path:
|
||||
required["prompt_metadata_path"] = args.prompt_metadata_path
|
||||
required["target_metadata_path"] = args.target_metadata_path
|
||||
if args.prompt_wav_path:
|
||||
required["prompt_wav_path"] = args.prompt_wav_path
|
||||
_require_paths(required)
|
||||
preprocess_dir = resolve_preprocess_weights_dir(args.preprocess_weights_dir)
|
||||
deploy_config = args.deploy_config or str(SVC_DEPLOY_CONFIG if kind == "svc" else SVS_DEPLOY_CONFIG)
|
||||
mode = kind.upper()
|
||||
|
||||
print(f"Loading SoulX-Singer {mode} from {args.model}")
|
||||
print(f" deploy: {deploy_config}")
|
||||
print(f" prompt: {args.prompt_audio}")
|
||||
print(f" target: {args.target_audio}")
|
||||
|
||||
omni = Omni(model=args.model, deploy_config=deploy_config, async_chunk=False) # SoulX-Singer is batch-only
|
||||
sampling = build_sampling(
|
||||
args,
|
||||
preprocess_weights_dir=preprocess_dir,
|
||||
kind=kind,
|
||||
)
|
||||
prompt = {"prompt_token_ids": [0], "multi_modal_data": {}}
|
||||
|
||||
print("Running generate (inline preprocess + DiT)...")
|
||||
outputs = list(omni.generate([prompt], sampling))
|
||||
if not outputs:
|
||||
raise RuntimeError("No output from omni.generate")
|
||||
|
||||
err = getattr(outputs[0], "error", None)
|
||||
if err is not None:
|
||||
raise RuntimeError(f"Pipeline error: {err}")
|
||||
|
||||
audio_np, sr = extract_audio(outputs)
|
||||
sf.write(args.output, audio_np, sr)
|
||||
print(f"Saved {args.output} ({sr} Hz, {len(audio_np) / sr:.2f}s)")
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -6,9 +6,12 @@ A unified script for text-to-video generation. Supports multiple models with mod
|
||||
|
||||
| Model | Default Resolution | Default Frames | Default Steps | Guidance | VRAM (BF16) |
|
||||
|---|---|---|---|---|---|
|
||||
| `Wan-AI/Wan2.1-VACE-1.3B-diffusers` | 480x832 | 81 | 30 | 5.0 | ~20 GiB (RTX 5090, VAE tiling) |
|
||||
| `Wan-AI/Wan2.2-T2V-A14B-Diffusers` | 720x1280 | 81 | 40 | 4.0 | ~60 GiB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-480p_t2v` | 480x832 | 121 | 50 | 6.0 | 1×A100 80GB |
|
||||
| `hunyuanvideo-community/HunyuanVideo-1.5-Diffusers-720p_t2v` | 720x1280 | 121 | 50 | 6.0 | FP8 + VAE tiling required |
|
||||
| `nvidia/Cosmos3-Nano` | 720x1280 | 189 | 35 | 6.0 | ~46 GiB (peak, 720p) |
|
||||
| `BestWishYsh/Helios-Base` / `Helios-Mid` / `Helios-Distilled` | 384x640 | 99 | 50 | 5.0 / 5.0 / 1.0 | — |
|
||||
|
||||
## Local CLI Usage
|
||||
|
||||
@@ -29,6 +32,28 @@ python text_to_video.py \
|
||||
--output t2v_out.mp4
|
||||
```
|
||||
|
||||
### Wan2.1 VACE (T2V)
|
||||
|
||||
VACE text-to-video uses this shared entrypoint. Conditional VACE tasks use
|
||||
the shared [`image_to_video.py`](../image_to_video/README.md#wan21-vace-conditional-tasks)
|
||||
entrypoint, which constructs the pipeline-native conditioning data from the
|
||||
provided media inputs. No explicit mode parameter is required.
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model Wan-AI/Wan2.1-VACE-1.3B-diffusers \
|
||||
--prompt "A sleek, humanoid robot stands in a vast warehouse filled with neatly stacked cardboard boxes on industrial shelves." \
|
||||
--seed 0 \
|
||||
--height 480 \
|
||||
--width 832 \
|
||||
--num-frames 81 \
|
||||
--num-inference-steps 30 \
|
||||
--guidance-scale 5.0 \
|
||||
--flow-shift 5.0 \
|
||||
--vae-use-tiling \
|
||||
--output vace_t2v_output.mp4
|
||||
```
|
||||
|
||||
LTX2 example:
|
||||
|
||||
```bash
|
||||
@@ -100,6 +125,63 @@ python text_to_video.py \
|
||||
--output quick_test.mp4
|
||||
```
|
||||
|
||||
### Cosmos3
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model nvidia/Cosmos3-Nano \
|
||||
--prompt "A robot arm is cleaning a plate in the kitchen." \
|
||||
--negative-prompt "blurry, distorted, low quality, jittery, deformed" \
|
||||
--height 720 --width 1280 --num-frames 189 --fps 24 \
|
||||
--num-inference-steps 35 --guidance-scale 6.0 \
|
||||
--extra-body '{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false,
|
||||
"use_resolution_template": false, "use_duration_template": false}' \
|
||||
--output cosmos3_t2v.mp4
|
||||
```
|
||||
|
||||
### Helios (T2V)
|
||||
|
||||
Helios ships three variants. Model-specific knobs (declared in
|
||||
`vllm_omni/model_extras/helios.py`) are passed via the generic `--extra-body`
|
||||
JSON flag rather than bespoke per-model flags.
|
||||
|
||||
**Helios-Base** (Stage 1 only):
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model BestWishYsh/Helios-Base \
|
||||
--prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--output helios_t2v_base.mp4
|
||||
```
|
||||
|
||||
**Helios-Mid** (Stage 2 pyramid + CFG-Zero*):
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model BestWishYsh/Helios-Mid \
|
||||
--prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
|
||||
--guidance-scale 5.0 \
|
||||
--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [20, 20, 20], "use_cfg_zero_star": true, "use_zero_init": true, "zero_steps": 1}' \
|
||||
--output helios_t2v_mid.mp4
|
||||
```
|
||||
|
||||
**Helios-Distilled** (Stage 2 pyramid + DMD, few-step):
|
||||
|
||||
```bash
|
||||
python text_to_video.py \
|
||||
--model BestWishYsh/Helios-Distilled \
|
||||
--prompt "A dynamic time-lapse of scenery rushing past the window of a speeding train." \
|
||||
--num-frames 240 \
|
||||
--guidance-scale 1.0 \
|
||||
--extra-body '{"is_enable_stage2": true, "pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}' \
|
||||
--output helios_t2v_distilled.mp4
|
||||
```
|
||||
|
||||
> Helios image-to-video (I2V) and video-to-video (V2V) require image/video
|
||||
> conditioning tensors that cannot be passed through the JSON `--extra-body`
|
||||
> flag; they are out of scope for this text-to-video example.
|
||||
|
||||
## Key Arguments
|
||||
|
||||
### Common
|
||||
@@ -112,6 +194,7 @@ python text_to_video.py \
|
||||
- `--num-inference-steps`: sampling steps. Default depends on model.
|
||||
- `--fps`: frames per second for the saved MP4.
|
||||
- `--output`: path to save the generated video.
|
||||
- `--extra-body`: JSON dict of model-specific knobs (declared in `vllm_omni/model_extras/`), merged into sampling `extra_args`. See the Helios recipes above.
|
||||
- `--vae-use-slicing`: enable VAE slicing for memory optimization.
|
||||
- `--vae-use-tiling`: enable VAE tiling for memory optimization.
|
||||
- `--cfg-parallel-size`: set it to 2 to enable CFG Parallel. See more examples in [`user_guide`](../../../docs/user_guide/diffusion/parallelism_acceleration.md#cfg-parallel).
|
||||
@@ -122,6 +205,7 @@ python text_to_video.py \
|
||||
- `--audio-sample-rate`: audio sample rate for embedded audio (when the pipeline returns audio).
|
||||
- `--quantization`: quantization method (`fp8` for FP8, `gguf` for GGUF).
|
||||
- `--flow-shift`: scheduler flow_shift parameter.
|
||||
- `--extra-body`: JSON object of model-specific generation params, filtered against the model's declared `extra_body_params` (see [`vllm_omni/model_extras`](../../../vllm_omni/model_extras)). Used by Cosmos3 (see above).
|
||||
|
||||
### Wan2.2-specific
|
||||
|
||||
|
||||
@@ -11,12 +11,23 @@ import numpy as np
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.diffusion.utils.param_utils import apply_declared_extra_args
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.model_extras import get_extra_body_params, get_model_class_name
|
||||
from vllm_omni.outputs import OmniRequestOutput
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
_MODEL_PRESETS = {
|
||||
"vace": {
|
||||
"height": 480,
|
||||
"width": 832,
|
||||
"num_frames": 81,
|
||||
"num_inference_steps": 30,
|
||||
"guidance_scale": 5.0,
|
||||
"fps": 16,
|
||||
"output": "vace_t2v_output.mp4",
|
||||
},
|
||||
"wan": {
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
@@ -35,13 +46,38 @@ _MODEL_PRESETS = {
|
||||
"fps": 24,
|
||||
"output": "hunyuan_video_15_output.mp4",
|
||||
},
|
||||
"cosmos": {
|
||||
"height": 720,
|
||||
"width": 1280,
|
||||
"num_frames": 189,
|
||||
"num_inference_steps": 35,
|
||||
"guidance_scale": 6.0,
|
||||
"fps": 24,
|
||||
"flow_shift": 10.0,
|
||||
"output": "cosmos3_t2v_output.mp4",
|
||||
},
|
||||
"helios": {
|
||||
"height": 384,
|
||||
"width": 640,
|
||||
"num_frames": 99,
|
||||
"num_inference_steps": 50,
|
||||
"guidance_scale": 5.0,
|
||||
"fps": 16,
|
||||
"output": "helios_output.mp4",
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _detect_preset(model: str) -> dict:
|
||||
model_lower = model.lower()
|
||||
if "vace" in model_lower:
|
||||
return _MODEL_PRESETS["vace"]
|
||||
if "cosmos" in model_lower:
|
||||
return _MODEL_PRESETS["cosmos"]
|
||||
if "hunyuan" in model_lower:
|
||||
return _MODEL_PRESETS["hunyuan"]
|
||||
if "helios" in model_lower:
|
||||
return _MODEL_PRESETS["helios"]
|
||||
return _MODEL_PRESETS["wan"]
|
||||
|
||||
|
||||
@@ -55,10 +91,26 @@ def parse_profiler_config(value: str) -> dict[str, Any]:
|
||||
return config
|
||||
|
||||
|
||||
def parse_extra_body(value: str) -> dict[str, Any]:
|
||||
"""Parse a JSON object of model-specific extra_body params.
|
||||
|
||||
Pipeline-declared knobs (see vllm_omni/model_extras/) are merged into
|
||||
OmniDiffusionSamplingParams.extra_args, so a single generic example can
|
||||
drive model-specific behaviour without bespoke per-model flags.
|
||||
"""
|
||||
try:
|
||||
body = json.loads(value)
|
||||
except json.JSONDecodeError as e:
|
||||
raise argparse.ArgumentTypeError(f"--extra-body must be valid JSON: {e}") from e
|
||||
if not isinstance(body, dict):
|
||||
raise argparse.ArgumentTypeError("--extra-body must be a JSON object")
|
||||
return body
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Generate a video from a text prompt. "
|
||||
"Supports Wan2.2, HunyuanVideo-1.5, and other text-to-video models."
|
||||
"Supports Wan2.2, HunyuanVideo-1.5, Helios, and other text-to-video models."
|
||||
)
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
@@ -74,6 +126,18 @@ def parse_args() -> argparse.Namespace:
|
||||
)
|
||||
parser.add_argument("--prompt", default="A serene lakeside sunrise with mist over the water.", help="Text prompt.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=parse_extra_body,
|
||||
default=None,
|
||||
help="JSON dict of model-specific extra_body params (declared in vllm_omni/model_extras/), "
|
||||
"merged into sampling extra_args. Unknown keys for the chosen model are dropped. "
|
||||
'Cosmos3 example: \'{"flow_shift": 10.0, "max_sequence_length": 4096, "guardrails": false, '
|
||||
'"use_resolution_template": false, "use_duration_template": false}\' '
|
||||
'(add "generate_sound": true and optional "sound_duration" for synchronized audio). '
|
||||
'Helios-Distilled example: \'{"is_enable_stage2": true, '
|
||||
'"pyramid_num_inference_steps_list": [2, 2, 2], "is_amplify_first_chunk": true}\'.',
|
||||
)
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=None, help="CFG scale. Default: model-specific.")
|
||||
parser.add_argument(
|
||||
@@ -322,7 +386,14 @@ def main():
|
||||
omni_kwargs["cache_config"] = cache_config
|
||||
omni_kwargs["enable_cache_dit_summary"] = args.enable_cache_dit_summary
|
||||
|
||||
# Cosmos3 loads its (gated) guardrail models at build time, so the guardrails
|
||||
# gate is an engine-level config (offline analog of the server's --no-guardrails).
|
||||
if args.extra_body and "guardrails" in args.extra_body:
|
||||
omni_kwargs["model_config"] = {"guardrails": bool(args.extra_body["guardrails"])}
|
||||
|
||||
omni = Omni(**omni_kwargs)
|
||||
model_class_name = get_model_class_name(omni) or model_class_name
|
||||
declared_extra_body_params = get_extra_body_params(model_class_name)
|
||||
|
||||
if profiler_enabled:
|
||||
print("[Profiler] Starting profiling...")
|
||||
@@ -358,10 +429,20 @@ def main():
|
||||
if args.guidance_scale_high is not None:
|
||||
sampling_kwargs["guidance_scale_2"] = args.guidance_scale_high
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(**sampling_kwargs)
|
||||
# Route model-specific knobs through extra_body, filtered against the model's
|
||||
# declared extra_body_params. Models without a declaration only forward explicit
|
||||
# --extra-body JSON (preserving the generic flags' legacy behavior).
|
||||
extra_body = dict(args.extra_body or {})
|
||||
if declared_extra_body_params:
|
||||
apply_declared_extra_args(sampling_params, declared_extra_body_params, extra_body)
|
||||
elif extra_body:
|
||||
sampling_params.extra_args.update({k: v for k, v in extra_body.items() if v is not None})
|
||||
|
||||
generation_start = time.perf_counter()
|
||||
frames = omni.generate(
|
||||
prompt_dict,
|
||||
OmniDiffusionSamplingParams(**sampling_kwargs),
|
||||
sampling_params,
|
||||
)
|
||||
generation_end = time.perf_counter()
|
||||
generation_time = generation_end - generation_start
|
||||
|
||||
@@ -1,148 +0,0 @@
|
||||
# VACE Video Generation
|
||||
|
||||
Generate videos from text prompts, images, or video conditions using vLLM-Omni's VACE diffusion pipeline.
|
||||
|
||||
- `vace_video_generation.py`: command-line script for multi-mode video generation with advanced options.
|
||||
|
||||
## Table of Contents
|
||||
|
||||
- [Overview](#overview)
|
||||
- [Quick Start](#quick-start)
|
||||
- [Key Arguments](#key-arguments)
|
||||
- [More CLI Examples](#more-cli-examples)
|
||||
|
||||
## Overview
|
||||
|
||||
[VACE](https://github.com/ali-vilab/VACE) (Video All-in-one Creation Engine) supports multiple video tasks through a single unified model, including text-to-video, image-to-video, first-last-frame interpolation, inpainting, and reference image-guided generation.
|
||||
|
||||
### Supported Models
|
||||
|
||||
| Model | Architecture | Peak VRAM (GiB) * | Model Weights (GiB) | HuggingFace |
|
||||
|-------|-------------|-------------------|---------------------|-------------|
|
||||
| Wan2.1-VACE (1.3B) | Wan2.1 | TBD | ~10 | [Wan-AI/Wan2.1-VACE-1.3B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-1.3B-diffusers) |
|
||||
| Wan2.1-VACE (14B) | Wan2.1 | TBD | ~38 | [Wan-AI/Wan2.1-VACE-14B-diffusers](https://huggingface.co/Wan-AI/Wan2.1-VACE-14B-diffusers) |
|
||||
|
||||
!!! info
|
||||
*Peak VRAM: based on basic single-card usage, 480x832 resolution, 81 frames, 30 inference steps, without any acceleration/optimization features.
|
||||
|
||||
Default model: `Wan-AI/Wan2.1-VACE-14B-diffusers`
|
||||
|
||||
## Quick Start
|
||||
|
||||
### Python API
|
||||
|
||||
Text-to-video generation:
|
||||
|
||||
```python
|
||||
import torch
|
||||
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
|
||||
if __name__ == "__main__":
|
||||
omni = Omni(model="Wan-AI/Wan2.1-VACE-1.3B-diffusers")
|
||||
prompt = "A sleek robot stands in a vast warehouse filled with boxes"
|
||||
outputs = omni.generate(
|
||||
prompt,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=480,
|
||||
width=832,
|
||||
num_frames=81,
|
||||
num_inference_steps=30,
|
||||
guidance_scale=5.0,
|
||||
),
|
||||
)
|
||||
video = outputs[0].images
|
||||
|
||||
from diffusers.utils import export_to_video
|
||||
export_to_video(list(video[0]), "t2v_output.mp4", fps=16)
|
||||
omni.close()
|
||||
```
|
||||
|
||||
### Local CLI Usage
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode t2v \
|
||||
--prompt "A sleek robot stands in a vast warehouse filled with boxes" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output t2v_output.mp4
|
||||
```
|
||||
|
||||
## Key Arguments
|
||||
|
||||
| Argument | Type | Default | Description |
|
||||
| -------- | ---- | ------- | ----------- |
|
||||
| `--mode` | str | `"t2v"` | VACE task mode: `t2v`, `i2v`, `v2lf`, `flf2v`, `inpaint`, `r2v` |
|
||||
| `--model` | str | `"Wan-AI/Wan2.1-VACE-14B-diffusers"` | Model ID or local path |
|
||||
| `--prompt` | str | `"A cat walking in a garden"` | Text description of desired video |
|
||||
| `--negative-prompt` | str | `""` | Negative prompt for classifier-free guidance |
|
||||
| `--image` | str | `None` | Input image path (for I2V, R2V, FLF2V, inpaint modes) |
|
||||
| `--last-image` | str | `None` | Last frame image path (for FLF2V mode) |
|
||||
| `--height` | int | `480` | Output video height in pixels (should be a multiple of 16) |
|
||||
| `--width` | int | `832` | Output video width in pixels (should be a multiple of 16) |
|
||||
| `--num-frames` | int | `81` | Number of video frames to generate |
|
||||
| `--num-inference-steps` | int | `30` | Number of denoising steps (more steps = higher quality, slower) |
|
||||
| `--guidance-scale` | float | `5.0` | Classifier-free guidance scale |
|
||||
| `--flow-shift` | float | `5.0` | Scheduler flow shift parameter |
|
||||
| `--seed` | int | `42` | Random seed for deterministic sampling |
|
||||
| `--fps` | int | `16` | Frames per second for the saved MP4 |
|
||||
| `--output` | str | `"vace_output.mp4"` | Path to save the generated video |
|
||||
| `--vae-use-tiling` | flag | on | Enable VAE tiling for memory optimization |
|
||||
| `--cfg-parallel-size` | int | `1` | Set to `2` to enable CFG Parallel |
|
||||
| `--ulysses-degree` | int | `1` | Ulysses sequence parallel degree for multi-GPU inference |
|
||||
| `--ring-degree` | int | `1` | Ring sequence parallel degree for hybrid Ulysses + Ring inference |
|
||||
| `--tensor-parallel-size` | int | — | Tensor parallel size |
|
||||
| `--enforce-eager` | flag | off | Disable torch.compile |
|
||||
|
||||
> If you encounter OOM errors, try `--vae-use-tiling` or multi-GPU parallelism options (`--ulysses-degree`, `--cfg-parallel-size`).
|
||||
|
||||
## More CLI Examples
|
||||
|
||||
### Image-to-Video (I2V)
|
||||
|
||||
First frame is kept, remaining frames are generated:
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode i2v \
|
||||
--image astronaut.jpg \
|
||||
--prompt "An astronaut emerging from a cracked egg on the moon" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output i2v_output.mp4
|
||||
```
|
||||
|
||||
### First-Last-Frame Interpolation (FLF2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode flf2v \
|
||||
--image first_frame.jpg --last-image last_frame.jpg \
|
||||
--prompt "A bird takes off from a branch and lands on another" \
|
||||
--height 512 --width 512 --num-frames 81 \
|
||||
--output flf2v_output.mp4
|
||||
```
|
||||
|
||||
### Inpainting
|
||||
|
||||
Center vertical stripe is masked and regenerated:
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode inpaint \
|
||||
--image scene.jpg \
|
||||
--prompt "Shrek walks out of a building" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output inpaint_output.mp4
|
||||
```
|
||||
|
||||
### Reference Image-guided (R2V)
|
||||
|
||||
```bash
|
||||
python vace_video_generation.py \
|
||||
--mode r2v \
|
||||
--image reference.jpg \
|
||||
--prompt "Camera slowly zooms out from the character" \
|
||||
--height 480 --width 832 --num-frames 81 \
|
||||
--output r2v_output.mp4
|
||||
```
|
||||
@@ -1,209 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
"""VACE video generation example.
|
||||
|
||||
VACE (Video All-in-one Creation Engine) supports multiple video tasks:
|
||||
- T2V: Text-to-Video
|
||||
- I2V: Image-to-Video (first frame conditioning)
|
||||
- V2LF: Video-to-Last-Frame
|
||||
- FLF2V: First-Last-Frame interpolation
|
||||
- Inpainting: Masked region generation
|
||||
- R2V: Reference image-guided generation
|
||||
|
||||
Usage examples:
|
||||
# T2V (text-to-video)
|
||||
python vace_video_generation.py --mode t2v --prompt "A robot in a warehouse"
|
||||
|
||||
# I2V (image-to-video, first frame kept)
|
||||
python vace_video_generation.py --mode i2v --image input.jpg --prompt "..."
|
||||
|
||||
# FLF2V (first-last frame interpolation)
|
||||
python vace_video_generation.py --mode flf2v --image first.jpg --last-image last.jpg
|
||||
|
||||
# R2V (reference image guided)
|
||||
python vace_video_generation.py --mode r2v --image ref.jpg --prompt "..."
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
import torch
|
||||
|
||||
from vllm_omni.diffusion.data import DiffusionParallelConfig
|
||||
from vllm_omni.entrypoints.omni import Omni
|
||||
from vllm_omni.inputs.data import OmniDiffusionSamplingParams
|
||||
from vllm_omni.platforms import current_omni_platform
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="VACE video generation.")
|
||||
parser.add_argument(
|
||||
"--model",
|
||||
default="Wan-AI/Wan2.1-VACE-14B-diffusers",
|
||||
help="VACE model ID or local path.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
default="t2v",
|
||||
choices=["t2v", "i2v", "v2lf", "flf2v", "inpaint", "r2v"],
|
||||
help="Generation mode.",
|
||||
)
|
||||
parser.add_argument("--prompt", default="A cat walking in a garden", help="Text prompt.")
|
||||
parser.add_argument("--negative-prompt", default="", help="Negative prompt.")
|
||||
parser.add_argument("--image", type=str, default=None, help="Input image path (for I2V, R2V, FLF2V, inpaint).")
|
||||
parser.add_argument("--last-image", type=str, default=None, help="Last frame image path (for FLF2V).")
|
||||
parser.add_argument("--video-dir", type=str, default=None, help="Directory of video frames (for inpaint).")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed.")
|
||||
parser.add_argument("--guidance-scale", type=float, default=5.0, help="CFG guidance scale.")
|
||||
parser.add_argument("--height", type=int, default=480, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=832, help="Video width.")
|
||||
parser.add_argument("--num-frames", type=int, default=81, help="Number of frames.")
|
||||
parser.add_argument("--num-inference-steps", type=int, default=30, help="Sampling steps.")
|
||||
parser.add_argument("--flow-shift", type=float, default=5.0, help="Scheduler flow_shift.")
|
||||
parser.add_argument("--output", type=str, default="vace_output.mp4", help="Output video path.")
|
||||
parser.add_argument("--fps", type=int, default=16, help="Output video FPS.")
|
||||
parser.add_argument("--vae-use-tiling", action="store_true", default=True, help="Enable VAE tiling.")
|
||||
parser.add_argument("--enforce-eager", action="store_true", help="Disable torch.compile.")
|
||||
parser.add_argument("--ulysses-degree", type=int, default=1, help="Ulysses SP degree.")
|
||||
parser.add_argument("--ring-degree", type=int, default=1, help="Ring attention degree.")
|
||||
parser.add_argument("--cfg-parallel-size", type=int, default=1, choices=[1, 2], help="CFG parallel size.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def build_prompts(args):
|
||||
"""Build prompt dict with multi_modal_data based on mode."""
|
||||
h, w, nf = args.height, args.width, args.num_frames
|
||||
|
||||
gray = PIL.Image.new("RGB", (w, h), (128, 128, 128))
|
||||
mask_black = PIL.Image.new("L", (w, h), 0)
|
||||
mask_white = PIL.Image.new("L", (w, h), 255)
|
||||
|
||||
prompt_data = {
|
||||
"prompt": args.prompt,
|
||||
"negative_prompt": args.negative_prompt,
|
||||
}
|
||||
|
||||
if args.mode == "t2v":
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "r2v":
|
||||
assert args.image, "--image required for R2V mode"
|
||||
ref_img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {"reference_images": [ref_img]}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "i2v":
|
||||
assert args.image, "--image required for I2V mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [img] + [gray] * (nf - 1),
|
||||
"mask": [mask_black] + [mask_white] * (nf - 1),
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "v2lf":
|
||||
assert args.image, "--image required for V2LF mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [gray] * (nf - 1) + [img],
|
||||
"mask": [mask_white] * (nf - 1) + [mask_black],
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "flf2v":
|
||||
assert args.image and args.last_image, "--image and --last-image required for FLF2V"
|
||||
first = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
last = PIL.Image.open(args.last_image).convert("RGB").resize((w, h))
|
||||
prompt_data["multi_modal_data"] = {
|
||||
"video": [first] + [gray] * (nf - 2) + [last],
|
||||
"mask": [mask_black] + [mask_white] * (nf - 2) + [mask_black],
|
||||
}
|
||||
return prompt_data
|
||||
|
||||
if args.mode == "inpaint":
|
||||
assert args.image, "--image required for inpaint mode"
|
||||
img = PIL.Image.open(args.image).convert("RGB").resize((w, h))
|
||||
d = 80
|
||||
frames, masks = [], []
|
||||
for _ in range(nf):
|
||||
base = np.array(img).copy()
|
||||
mask = PIL.Image.new("L", (w, h), 0)
|
||||
stripe = PIL.Image.new("L", (2 * d, h), 255)
|
||||
mask.paste(stripe, (w // 2 - d, 0))
|
||||
base[np.array(mask) > 128] = 128
|
||||
frames.append(PIL.Image.fromarray(base))
|
||||
masks.append(mask)
|
||||
prompt_data["multi_modal_data"] = {"video": frames, "mask": masks}
|
||||
return prompt_data
|
||||
|
||||
raise ValueError(f"Unknown mode: {args.mode}")
|
||||
|
||||
|
||||
def main():
|
||||
args = parse_args()
|
||||
generator = torch.Generator(device=current_omni_platform.device_type).manual_seed(args.seed)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
ulysses_degree=args.ulysses_degree,
|
||||
ring_degree=args.ring_degree,
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
)
|
||||
|
||||
omni = Omni(
|
||||
model=args.model,
|
||||
vae_use_tiling=args.vae_use_tiling,
|
||||
flow_shift=args.flow_shift,
|
||||
enforce_eager=args.enforce_eager,
|
||||
parallel_config=parallel_config,
|
||||
)
|
||||
|
||||
prompt_data = build_prompts(args)
|
||||
|
||||
print(f"\n{'=' * 60}")
|
||||
print(f"VACE {args.mode.upper()} Generation")
|
||||
print(f" Model: {args.model}")
|
||||
print(f" Size: {args.width}x{args.height}, {args.num_frames} frames, {args.num_inference_steps} steps")
|
||||
print(f"{'=' * 60}\n")
|
||||
|
||||
start = time.perf_counter()
|
||||
outputs = omni.generate(
|
||||
prompt_data,
|
||||
OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_frames=args.num_frames,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
guidance_scale=args.guidance_scale,
|
||||
generator=generator,
|
||||
),
|
||||
)
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
video = outputs[0].images
|
||||
if isinstance(video, list):
|
||||
video = video[0]
|
||||
if isinstance(video, torch.Tensor):
|
||||
video = video.cpu().numpy()
|
||||
if video.ndim == 5:
|
||||
video = video[0]
|
||||
print(f"Output shape: {video.shape}, Time: {elapsed:.1f}s")
|
||||
|
||||
output_path = Path(args.output)
|
||||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
from diffusers.utils import export_to_video
|
||||
|
||||
if np.issubdtype(video.dtype, np.integer):
|
||||
video = video.astype(np.float32) / 255.0
|
||||
export_to_video(list(video), str(output_path), fps=args.fps)
|
||||
print(f"Saved to {output_path}")
|
||||
|
||||
omni.close()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -87,6 +87,63 @@ python x_to_video_audio.py \
|
||||
```
|
||||
|
||||
|
||||
## MagiHuman (`--model-type magi-human`)
|
||||
|
||||
MagiHuman is a text → video+audio model with a DiT MoE backbone and a ~9B-param
|
||||
T5Gemma text encoder. A detailed text prompt is the only required input; an optional
|
||||
image and/or audio file may be supplied for conditioning. Natively supports Tensor
|
||||
Parallelism. For an 80GB node, `--tensor-parallel-size 4` is recommended to shard
|
||||
the MoE weights and the text encoder.
|
||||
|
||||
> Install [MagiCompiler](https://github.com/SandAI-org/MagiCompiler) for correct
|
||||
> attention-kernel behaviour (the pipeline otherwise falls back to stubs).
|
||||
|
||||
Text-only generation:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman with long, wavy golden blonde hair... <dialogue and background sound>" \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"seconds": 5, "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
With optional image and audio conditioning:
|
||||
|
||||
```bash
|
||||
python x_to_video_audio.py \
|
||||
--model-type magi-human \
|
||||
--model /path/to/daVinci-MagiHuman \
|
||||
--prompt "A young woman..." \
|
||||
--tensor-parallel-size 4 \
|
||||
--height 256 --width 448 \
|
||||
--num-inference-steps 8 \
|
||||
--seed 52 \
|
||||
--extra-body '{"image_path": "/path/to/ref.jpg", "audio_path": "/path/to/ref.wav", "sr_height": 1080, "sr_width": 1920, "sr_num_inference_steps": 5}' \
|
||||
--output output_magihuman.mp4
|
||||
```
|
||||
|
||||
MagiHuman-specific arguments are passed as a JSON dict via `--extra-body` (declared
|
||||
in `vllm_omni/model_extras/magi_human.py`, routed via `extra_args`):
|
||||
|
||||
- `seconds`: output duration in seconds (default 10; ignored when `audio_path` is set,
|
||||
because the audio length then determines the number of frames).
|
||||
- `audio_path`: path to an audio file for audio-to-video conditioning. When provided,
|
||||
the audio drives the frame count and is encoded as a latent condition into the DiT
|
||||
(`is_a2v` mode). Omit for pure text-to-video+audio generation.
|
||||
- `image_path`: path to an image file for visual conditioning. Applied at both the
|
||||
base-resolution (BR) and super-resolution (SR) stages when SR is enabled.
|
||||
- `sr_height` / `sr_width`: super-resolution output resolution. SR stage is skipped
|
||||
when these are omitted.
|
||||
- `sr_num_inference_steps`: denoising steps for the SR stage.
|
||||
|
||||
## DreamID-Omni arguments
|
||||
|
||||
Key arguments:
|
||||
- `--prompt`: text description (string).
|
||||
- `--model`: path to the model local directory.
|
||||
|
||||
@@ -2,6 +2,7 @@
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import re
|
||||
import time
|
||||
from typing import Any
|
||||
@@ -20,14 +21,26 @@ AUDIO_EXTENSIONS = {".wav", ".mp3", ".flac", ".m4a", ".aac", ".ogg"}
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Offline inference for DreamID-Omni (video + audio).")
|
||||
parser.add_argument("--model", required=True, help="DreamID ckpt root directory.")
|
||||
parser.add_argument("--model-type", default="dreamid-omni", help="Model type.")
|
||||
parser = argparse.ArgumentParser(description="Offline inference for X -> video+audio models.")
|
||||
parser.add_argument("--model", required=True, help="Model ckpt root directory.")
|
||||
parser.add_argument(
|
||||
"--model-type",
|
||||
default="dreamid-omni",
|
||||
choices=["dreamid-omni", "magi-human"],
|
||||
help="Model type.",
|
||||
)
|
||||
parser.add_argument("--prompt", default=None, help="Text prompt.")
|
||||
|
||||
parser.add_argument("--image-path", type=str, nargs="+", help="list of image-path")
|
||||
parser.add_argument("--audio-path", type=str, nargs="+", help="list of audio-path")
|
||||
parser.add_argument("--prompt-file", type=str, default=None, help="Text prompt in json format.")
|
||||
parser.add_argument(
|
||||
"--extra-body",
|
||||
type=str,
|
||||
default=None,
|
||||
help="[magi-human] JSON dict of model-specific extra params (declared in vllm_omni/model_extras/), "
|
||||
'merged into sampling extra_args. Example: \'{"image_path": "/path/to/img.jpg", "seconds": 5}\'.',
|
||||
)
|
||||
|
||||
parser.add_argument("--height", type=int, default=704, help="Video height.")
|
||||
parser.add_argument("--width", type=int, default=1280, help="Video width.")
|
||||
@@ -142,8 +155,6 @@ def main() -> None:
|
||||
|
||||
text_prompt = args.prompt
|
||||
if args.prompt_file:
|
||||
import json
|
||||
|
||||
with open(args.prompt_file) as f:
|
||||
text_prompt = json.load(f)
|
||||
text_prompt = re.sub(
|
||||
@@ -155,25 +166,36 @@ def main() -> None:
|
||||
text_prompt = re.sub(r"\[[A-Z_]+\]", "", text_prompt)
|
||||
text_prompt = re.sub(r"\n\s*\n", "\n", text_prompt).strip()
|
||||
|
||||
image, audio = load_image_and_audio(args.image_path, args.audio_path)
|
||||
if args.model_type == "magi-human":
|
||||
prompt = text_prompt
|
||||
extra_args = json.loads(args.extra_body) if args.extra_body else {}
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args=extra_args,
|
||||
)
|
||||
else:
|
||||
image, audio = load_image_and_audio(args.image_path, args.audio_path)
|
||||
|
||||
prompt = {
|
||||
"prompt": text_prompt,
|
||||
"video_negative_prompt": args.video_negative_prompt,
|
||||
"audio_negative_prompt": args.audio_negative_prompt,
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
}
|
||||
prompt = {
|
||||
"prompt": text_prompt,
|
||||
"video_negative_prompt": args.video_negative_prompt,
|
||||
"audio_negative_prompt": args.audio_negative_prompt,
|
||||
"multi_modal_data": {"image": image, "audio": audio},
|
||||
}
|
||||
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"solver_name": args.solver_name,
|
||||
"shift": args.shift,
|
||||
},
|
||||
)
|
||||
sampling_params = OmniDiffusionSamplingParams(
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
num_inference_steps=args.num_inference_steps,
|
||||
seed=args.seed,
|
||||
extra_args={
|
||||
"solver_name": args.solver_name,
|
||||
"shift": args.shift,
|
||||
},
|
||||
)
|
||||
|
||||
parallel_config = DiffusionParallelConfig(
|
||||
cfg_parallel_size=args.cfg_parallel_size,
|
||||
@@ -215,22 +237,28 @@ def main() -> None:
|
||||
elapsed = time.perf_counter() - start
|
||||
|
||||
if not outputs:
|
||||
raise RuntimeError("No output returned from DreamID-Omni.")
|
||||
raise RuntimeError("No output returned from the model.")
|
||||
result = outputs[0]
|
||||
if not result.images:
|
||||
raise RuntimeError("No video frames found in DreamID-Omni output.")
|
||||
raise RuntimeError("No video frames found in model output.")
|
||||
generated_video = result.images[0]
|
||||
mm = result.multimodal_output or {}
|
||||
generated_audio = mm.get("audio")
|
||||
fps = int(mm.get("fps", 24))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 16000))
|
||||
|
||||
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
|
||||
# mux_video_audio_bytes expects (F, H, W, C) uint8.
|
||||
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
|
||||
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
|
||||
frames = generated_video.transpose(1, 2, 3, 0)
|
||||
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
if args.model_type == "magi-human":
|
||||
# MagiHuman returns frames already as (F, H, W, C) uint8, ready to mux.
|
||||
fps = float(mm.get("fps", 25))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 24000))
|
||||
frames = generated_video
|
||||
else:
|
||||
fps = float(mm.get("fps", 24))
|
||||
sample_rate = int(mm.get("audio_sample_rate", 16000))
|
||||
# DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1].
|
||||
# mux_video_audio_bytes expects (F, H, W, C) uint8.
|
||||
if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4:
|
||||
raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}")
|
||||
frames = generated_video.transpose(1, 2, 3, 0)
|
||||
frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8)
|
||||
|
||||
audio_np = None
|
||||
if generated_audio is not None:
|
||||
|
||||
@@ -0,0 +1,146 @@
|
||||
# AURA Omni Native Pipeline
|
||||
|
||||
`aura_omni` serves AURA as a native multi-stage vLLM-Omni pipeline:
|
||||
|
||||
```text
|
||||
Qwen3-ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Qwen3-TTS Code2Wav
|
||||
```
|
||||
|
||||
The pipeline has three semantic modules, but four engine stages because the
|
||||
existing Qwen3-TTS implementation is natively split into Talker and Code2Wav.
|
||||
|
||||
Start the server with the deploy profile:
|
||||
|
||||
```bash
|
||||
vllm serve aurateam/AURA \
|
||||
--omni \
|
||||
--port 8091 \
|
||||
--deploy-config vllm_omni/deploy/aura_omni.yaml \
|
||||
--served-model-name aurateam/AURA \
|
||||
--trust-remote-code
|
||||
```
|
||||
|
||||
The deploy file sets per-stage model repos:
|
||||
|
||||
- Stage 0 ASR: `Qwen/Qwen3-ASR-1.7B`
|
||||
- Stage 1 AURA: `aurateam/AURA`
|
||||
- Stage 2/3 TTS: `Qwen/Qwen3-TTS-12Hz-1.7B-Base`
|
||||
|
||||
For local weights, edit the `model` value on each stage in
|
||||
`vllm_omni/deploy/aura_omni.yaml`. The deploy profile includes
|
||||
`pipeline: aura_omni`, so the server uses this four-stage topology even when
|
||||
the command-line model path points at one component checkpoint.
|
||||
|
||||
Expected request shape:
|
||||
|
||||
- Send microphone audio as the Stage 0 multimodal audio input.
|
||||
- Include video frames in the original request `multi_modal_data`; the
|
||||
`asr2aura` processor carries them forward to AURA.
|
||||
- Optional `additional_information` keys:
|
||||
- `aura_system_prompt`
|
||||
- `tts_task_type`
|
||||
- `tts_language`
|
||||
- `tts_speaker`
|
||||
- `tts_instruct`
|
||||
- `tts_ref_audio`
|
||||
- `tts_ref_text`
|
||||
- `tts_x_vector_only_mode`
|
||||
- `tts_pass_token_ids`
|
||||
|
||||
If AURA emits `<|silent|>`, the `aura2tts` processor returns no TTS request, so
|
||||
the TTS stages are skipped for that turn.
|
||||
|
||||
## GPU Utilization Recommendation
|
||||
|
||||
Tune `gpu_memory_utilization` per stage in `vllm_omni/deploy/aura_omni.yaml`.
|
||||
Recommended baseline on one GPU for H200
|
||||
|
||||
- Stage 0 (ASR): `0.10`
|
||||
- Stage 1 (AURA): `0.4`
|
||||
- Stage 2 (Qwen3-TTS Talker): `0.20`
|
||||
- Stage 3 (Qwen3-TTS Code2Wav): `0.20`
|
||||
|
||||
## Python Client
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--host localhost \
|
||||
--port 8091 \
|
||||
--model aurateam/AURA \
|
||||
--modalities text,audio
|
||||
```
|
||||
|
||||
Use local media:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--audio-path /path/to/input.wav \
|
||||
--video-path /path/to/video.mp4 \
|
||||
--output-dir output_aura_omni_online
|
||||
```
|
||||
|
||||
Base voice clone mode (default, recommended as x-vector while debugging ICL):
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-task-type Base \
|
||||
--tts-ref-audio vllm-omni/tests/assets/qwen3_tts/clone_2.wav \
|
||||
--tts-ref-text "Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
```
|
||||
|
||||
Enable AURA token-id passthrough explicitly:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-pass-token-ids
|
||||
```
|
||||
|
||||
CustomVoice mode requires stages 2 and 3 in `aura_omni.yaml` to point at a
|
||||
Qwen3-TTS CustomVoice checkpoint:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/openai_chat_completion_client.py \
|
||||
--tts-task-type CustomVoice \
|
||||
--tts-speaker Vivian
|
||||
```
|
||||
|
||||
By default, AURA responses are passed to Qwen3-TTS as text. Set
|
||||
`tts_pass_token_ids=true` to pass AURA-generated assistant token ids directly
|
||||
to Qwen3-TTS instead. The processor still uses AURA token ids, when available,
|
||||
to estimate the Talker prompt length in the default text path.
|
||||
|
||||
## Curl
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/aura_omni
|
||||
bash run_curl_multimodal_generation.sh
|
||||
```
|
||||
|
||||
Set `PORT`, `MODEL`, or `OUTPUT_DIR` to override defaults:
|
||||
|
||||
```bash
|
||||
PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
|
||||
TTS_PASS_TOKEN_IDS=true PORT=8666 MODEL=aurateam/AURA bash run_curl_multimodal_generation.sh
|
||||
```
|
||||
|
||||
## Gradio
|
||||
|
||||
Launch the server and Gradio UI together:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/aura_omni
|
||||
bash run_gradio_demo.sh
|
||||
```
|
||||
|
||||
If the server is already running:
|
||||
|
||||
```bash
|
||||
python examples/online_serving/aura_omni/gradio_demo.py \
|
||||
--model aurateam/AURA \
|
||||
--api-base http://localhost:8091/v1
|
||||
```
|
||||
|
||||
## Offline
|
||||
|
||||
For offline inference, see
|
||||
[`examples/offline_inference/aura_omni`](../../offline_inference/aura_omni/).
|
||||
@@ -0,0 +1,305 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Gradio demo for AURA Omni online serving."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
import io
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
try:
|
||||
import gradio as gr
|
||||
except ImportError:
|
||||
raise ImportError("gradio is required to run this demo. Install it with: pip install 'vllm-omni[demo]'") from None
|
||||
|
||||
import numpy as np
|
||||
import soundfile as sf
|
||||
from openai import OpenAI
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
default_qwen3_tts_ref_audio_path,
|
||||
)
|
||||
|
||||
SEED = 42
|
||||
DEFAULT_TTS_REF_AUDIO = default_qwen3_tts_ref_audio_path()
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = argparse.ArgumentParser(description="Gradio demo for AURA Omni online inference.")
|
||||
parser.add_argument("--model", default="aurateam/AURA", help="Model name served by vLLM.")
|
||||
parser.add_argument("--api-base", default="http://localhost:8091/v1", help="OpenAI-compatible API base.")
|
||||
parser.add_argument("--ip", default="127.0.0.1", help="Gradio host.")
|
||||
parser.add_argument("--port", type=int, default=7862, help="Gradio port.")
|
||||
parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _audio_to_data_url(audio_file: Any | None) -> str | None:
|
||||
if audio_file is None:
|
||||
return None
|
||||
|
||||
def _path_to_data_url(path_str: str) -> str | None:
|
||||
if not path_str:
|
||||
return None
|
||||
path = Path(path_str)
|
||||
if not path.exists():
|
||||
return None
|
||||
with open(path, "rb") as f:
|
||||
return f"data:audio/wav;base64,{base64.b64encode(f.read()).decode('utf-8')}"
|
||||
|
||||
if isinstance(audio_file, str):
|
||||
return _path_to_data_url(audio_file)
|
||||
|
||||
# Gradio >=5 may return dict payloads with local temp path.
|
||||
if isinstance(audio_file, dict):
|
||||
path_val = audio_file.get("path") or audio_file.get("name")
|
||||
if isinstance(path_val, str):
|
||||
return _path_to_data_url(path_val)
|
||||
return None
|
||||
|
||||
if isinstance(audio_file, tuple):
|
||||
# Common shape: (sample_rate, np.ndarray)
|
||||
if len(audio_file) == 2 and isinstance(audio_file[0], (int, float)):
|
||||
sample_rate, audio_np = audio_file
|
||||
audio_np = np.asarray(audio_np)
|
||||
if audio_np.ndim > 1:
|
||||
audio_np = audio_np[:, 0]
|
||||
if audio_np.dtype != np.int16:
|
||||
audio_np = np.clip(audio_np.astype(np.float32), -1.0, 1.0)
|
||||
buf = io.BytesIO()
|
||||
sf.write(buf, audio_np, int(sample_rate), format="WAV")
|
||||
return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode('utf-8')}"
|
||||
# Compatibility shape from some gradio builds: (filepath, (sr, np.ndarray))
|
||||
if len(audio_file) >= 1 and isinstance(audio_file[0], str):
|
||||
path_url = _path_to_data_url(audio_file[0])
|
||||
if path_url is not None:
|
||||
return path_url
|
||||
return None
|
||||
|
||||
|
||||
def _video_to_data_url(video_file: str | None) -> str | None:
|
||||
if not video_file:
|
||||
return None
|
||||
path = Path(video_file)
|
||||
mime = {
|
||||
".mp4": "video/mp4",
|
||||
".webm": "video/webm",
|
||||
".mov": "video/quicktime",
|
||||
".avi": "video/x-msvideo",
|
||||
".mkv": "video/x-matroska",
|
||||
}.get(path.suffix.lower(), "video/mp4")
|
||||
with open(path, "rb") as f:
|
||||
return f"data:{mime};base64,{base64.b64encode(f.read()).decode('utf-8')}"
|
||||
|
||||
|
||||
def _sampling_params_list() -> list[dict[str, Any]]:
|
||||
return [
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED},
|
||||
{
|
||||
"temperature": 0.5,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 256,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
{
|
||||
"temperature": 0.9,
|
||||
"top_k": 50,
|
||||
"max_tokens": 4096,
|
||||
"seed": SEED,
|
||||
"detokenize": False,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [2150],
|
||||
},
|
||||
{
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 65536,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def _decode_audio_choice(audio_data: str | None) -> tuple[int, np.ndarray] | None:
|
||||
if not audio_data:
|
||||
return None
|
||||
wav_bytes = base64.b64decode(audio_data)
|
||||
audio_np, sample_rate = sf.read(io.BytesIO(wav_bytes))
|
||||
if audio_np.ndim > 1:
|
||||
audio_np = audio_np[:, 0]
|
||||
return int(sample_rate), audio_np.astype(np.float32)
|
||||
|
||||
|
||||
def build_interface(client: OpenAI, model: str):
|
||||
def run(
|
||||
audio_file,
|
||||
video_file,
|
||||
prompt: str,
|
||||
aura_system_prompt: str,
|
||||
tts_task_type: str,
|
||||
tts_language: str,
|
||||
tts_speaker: str,
|
||||
tts_instruct: str,
|
||||
tts_ref_audio: str,
|
||||
tts_ref_text: str,
|
||||
tts_x_vector_only_mode: bool,
|
||||
tts_pass_token_ids: bool,
|
||||
):
|
||||
audio_url = _audio_to_data_url(audio_file)
|
||||
video_url = _video_to_data_url(video_file)
|
||||
if not audio_url:
|
||||
return "Please provide audio input for the ASR stage.", None
|
||||
if tts_task_type == "Base" and (not (tts_ref_audio or "").strip() or not (tts_ref_text or "").strip()):
|
||||
return "Base TTS requires both reference audio and reference transcript.", None
|
||||
content: list[dict[str, Any]] = [{"type": "audio_url", "audio_url": {"url": audio_url}}]
|
||||
if video_url:
|
||||
content.append({"type": "video_url", "video_url": {"url": video_url}})
|
||||
content.append(
|
||||
{"type": "text", "text": prompt or "Use the audio and video together to decide whether a reply is needed."}
|
||||
)
|
||||
|
||||
additional_information: dict[str, Any] = {
|
||||
"aura_system_prompt": aura_system_prompt,
|
||||
"tts_task_type": tts_task_type,
|
||||
"tts_instruct": tts_instruct,
|
||||
"tts_pass_token_ids": bool(tts_pass_token_ids),
|
||||
}
|
||||
if tts_task_type == "CustomVoice":
|
||||
additional_information.update(
|
||||
{
|
||||
"tts_language": tts_language or "English",
|
||||
"tts_speaker": tts_speaker or "Vivian",
|
||||
}
|
||||
)
|
||||
else:
|
||||
additional_information.update(
|
||||
{
|
||||
"tts_ref_audio": tts_ref_audio.strip() if tts_ref_audio else None,
|
||||
"tts_ref_text": tts_ref_text.strip() if tts_ref_text else None,
|
||||
"tts_x_vector_only_mode": bool(tts_x_vector_only_mode),
|
||||
}
|
||||
)
|
||||
|
||||
response = client.chat.completions.create(
|
||||
model=model,
|
||||
messages=[{"role": "user", "content": content}],
|
||||
modalities=["text", "audio"],
|
||||
extra_body={
|
||||
"sampling_params_list": _sampling_params_list(),
|
||||
"additional_information": additional_information,
|
||||
},
|
||||
timeout=600.0,
|
||||
)
|
||||
|
||||
text_parts: list[str] = []
|
||||
audio_output = None
|
||||
for choice in response.choices:
|
||||
message = choice.message
|
||||
if message.content:
|
||||
text_parts.append(str(message.content))
|
||||
if getattr(message, "audio", None):
|
||||
audio_output = _decode_audio_choice(message.audio.data)
|
||||
return "\n\n".join(text_parts) or "No text output.", audio_output
|
||||
|
||||
default_system = (
|
||||
"You are receiving a live video stream where the final frame is the present moment. "
|
||||
"Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English."
|
||||
)
|
||||
|
||||
def _toggle_tts_advanced(task_type: str):
|
||||
is_base = task_type == "Base"
|
||||
return (
|
||||
gr.update(visible=not is_base),
|
||||
gr.update(visible=is_base),
|
||||
)
|
||||
|
||||
with gr.Blocks() as demo:
|
||||
gr.Markdown("# AURA Omni")
|
||||
gr.Markdown("ASR -> AURA/Qwen3-VL -> Qwen3-TTS Talker -> Code2Wav")
|
||||
with gr.Row():
|
||||
audio_input = gr.Audio(label="Audio input", sources=["upload", "microphone"], type="numpy")
|
||||
video_input = gr.Video(label="Video input (optional)", sources=["upload"])
|
||||
prompt = gr.Textbox(
|
||||
label="Instruction",
|
||||
value="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.",
|
||||
lines=2,
|
||||
)
|
||||
with gr.Accordion("Advanced", open=False):
|
||||
aura_system_prompt = gr.Textbox(label="AURA system prompt", value=default_system, lines=4)
|
||||
tts_task_type = gr.Radio(choices=["Base", "CustomVoice"], value="Base", label="TTS task type")
|
||||
tts_instruct = gr.Textbox(label="TTS instruction", value="")
|
||||
tts_pass_token_ids = gr.Checkbox(
|
||||
label="Pass AURA token ids directly to TTS",
|
||||
value=False,
|
||||
)
|
||||
|
||||
with gr.Group(visible=False) as customvoice_config_group:
|
||||
tts_language = gr.Dropdown(
|
||||
choices=["Chinese", "English", "Japanese", "Korean", "Cantonese"],
|
||||
value="Chinese",
|
||||
label="CustomVoice language",
|
||||
)
|
||||
tts_speaker = gr.Textbox(label="CustomVoice speaker", value="Vivian")
|
||||
|
||||
with gr.Group(visible=True) as base_config_group:
|
||||
tts_ref_audio = gr.Textbox(
|
||||
label="Base reference audio path/URL",
|
||||
value=DEFAULT_TTS_REF_AUDIO,
|
||||
)
|
||||
tts_ref_text = gr.Textbox(
|
||||
label="Base reference transcript",
|
||||
value=DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
lines=3,
|
||||
)
|
||||
tts_x_vector_only_mode = gr.Checkbox(
|
||||
label="Base x-vector only mode (disable ICL)",
|
||||
value=False,
|
||||
)
|
||||
|
||||
tts_task_type.change(
|
||||
_toggle_tts_advanced,
|
||||
inputs=[tts_task_type],
|
||||
outputs=[customvoice_config_group, base_config_group],
|
||||
)
|
||||
button = gr.Button("Generate", variant="primary")
|
||||
with gr.Row():
|
||||
text_output = gr.Textbox(label="Text output", lines=8)
|
||||
audio_output = gr.Audio(label="Audio output", interactive=False)
|
||||
button.click(
|
||||
run,
|
||||
inputs=[
|
||||
audio_input,
|
||||
video_input,
|
||||
prompt,
|
||||
aura_system_prompt,
|
||||
tts_task_type,
|
||||
tts_language,
|
||||
tts_speaker,
|
||||
tts_instruct,
|
||||
tts_ref_audio,
|
||||
tts_ref_text,
|
||||
tts_x_vector_only_mode,
|
||||
tts_pass_token_ids,
|
||||
],
|
||||
outputs=[text_output, audio_output],
|
||||
)
|
||||
demo.queue()
|
||||
return demo
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
client = OpenAI(api_key="EMPTY", base_url=args.api_base)
|
||||
demo = build_interface(client, args.model)
|
||||
demo.launch(server_name=args.ip, server_port=args.port, share=args.share)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -0,0 +1,199 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""OpenAI-compatible client for the AURA Omni pipeline."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import base64
|
||||
import io
|
||||
import os
|
||||
|
||||
import soundfile as sf
|
||||
from openai import OpenAI
|
||||
from vllm.assets.audio import AudioAsset
|
||||
from vllm.utils.argparse_utils import FlexibleArgumentParser
|
||||
|
||||
from vllm_omni.model_executor.stage_input_processors.aura_omni import (
|
||||
DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
default_qwen3_tts_ref_audio_path,
|
||||
)
|
||||
|
||||
SEED = 42
|
||||
DEFAULT_MODEL = "aurateam/AURA"
|
||||
DEFAULT_VIDEO_URL = "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
|
||||
|
||||
def _encode_file(path: str) -> str:
|
||||
with open(path, "rb") as f:
|
||||
return base64.b64encode(f.read()).decode("utf-8")
|
||||
|
||||
|
||||
def _data_url(path: str, default_mime: str) -> str:
|
||||
suffix = os.path.splitext(path)[1].lower()
|
||||
mime_by_suffix = {
|
||||
".wav": "audio/wav",
|
||||
".mp3": "audio/mpeg",
|
||||
".ogg": "audio/ogg",
|
||||
".flac": "audio/flac",
|
||||
".m4a": "audio/mp4",
|
||||
".mp4": "video/mp4",
|
||||
".webm": "video/webm",
|
||||
".mov": "video/quicktime",
|
||||
".avi": "video/x-msvideo",
|
||||
".mkv": "video/x-matroska",
|
||||
}
|
||||
return f"data:{mime_by_suffix.get(suffix, default_mime)};base64,{_encode_file(path)}"
|
||||
|
||||
|
||||
def media_url(path_or_url: str | None, *, kind: str) -> str:
|
||||
if path_or_url:
|
||||
if path_or_url.startswith(("http://", "https://", "data:")):
|
||||
return path_or_url
|
||||
if not os.path.exists(path_or_url):
|
||||
raise FileNotFoundError(f"{kind} file not found: {path_or_url}")
|
||||
return _data_url(path_or_url, "audio/wav" if kind == "audio" else "video/mp4")
|
||||
if kind == "audio":
|
||||
return AudioAsset("mary_had_lamb").url
|
||||
return DEFAULT_VIDEO_URL
|
||||
|
||||
|
||||
def sampling_params_list() -> list[dict]:
|
||||
return [
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": SEED},
|
||||
{
|
||||
"temperature": 0.5,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 256,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
{
|
||||
"temperature": 0.9,
|
||||
"top_k": 50,
|
||||
"max_tokens": 4096,
|
||||
"seed": SEED,
|
||||
"detokenize": False,
|
||||
"repetition_penalty": 1.05,
|
||||
"stop_token_ids": [2150],
|
||||
},
|
||||
{
|
||||
"temperature": 0.0,
|
||||
"top_p": 1.0,
|
||||
"top_k": -1,
|
||||
"max_tokens": 65536,
|
||||
"seed": SEED,
|
||||
"repetition_penalty": 1.0,
|
||||
},
|
||||
]
|
||||
|
||||
|
||||
def parse_modalities(value: str | None) -> list[str] | None:
|
||||
if not value:
|
||||
return None
|
||||
return [item.strip() for item in value.split(",") if item.strip()]
|
||||
|
||||
|
||||
def save_response(response, output_dir: str) -> None:
|
||||
os.makedirs(output_dir, exist_ok=True)
|
||||
for idx, choice in enumerate(response.choices):
|
||||
message = choice.message
|
||||
if message.content:
|
||||
out_txt = os.path.join(output_dir, f"choice_{idx}.txt")
|
||||
with open(out_txt, "w", encoding="utf-8") as f:
|
||||
f.write(str(message.content).strip() + "\n")
|
||||
print(f"Text saved to {out_txt}")
|
||||
print(message.content)
|
||||
if getattr(message, "audio", None):
|
||||
audio_bytes = base64.b64decode(message.audio.data)
|
||||
audio_np, sample_rate = sf.read(io.BytesIO(audio_bytes))
|
||||
out_wav = os.path.join(output_dir, f"choice_{idx}.wav")
|
||||
sf.write(out_wav, audio_np, int(sample_rate), format="WAV")
|
||||
print(f"Audio saved to {out_wav}")
|
||||
|
||||
|
||||
def main(args) -> None:
|
||||
client = OpenAI(base_url=f"http://{args.host}:{args.port}/v1", api_key="EMPTY")
|
||||
messages = [
|
||||
{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": media_url(args.audio_path, kind="audio")}},
|
||||
{"type": "video_url", "video_url": {"url": media_url(args.video_path, kind="video")}},
|
||||
{"type": "text", "text": args.prompt},
|
||||
],
|
||||
}
|
||||
]
|
||||
response = client.chat.completions.create(
|
||||
model=args.model,
|
||||
messages=messages,
|
||||
modalities=parse_modalities(args.modalities),
|
||||
extra_body={
|
||||
"sampling_params_list": sampling_params_list(),
|
||||
"additional_information": {
|
||||
"aura_system_prompt": args.aura_system_prompt,
|
||||
"tts_task_type": args.tts_task_type,
|
||||
"tts_language": args.tts_language,
|
||||
"tts_speaker": args.tts_speaker,
|
||||
"tts_instruct": args.tts_instruct,
|
||||
"tts_ref_audio": args.tts_ref_audio,
|
||||
"tts_ref_text": args.tts_ref_text,
|
||||
"tts_x_vector_only_mode": args.tts_x_vector_only_mode,
|
||||
"tts_pass_token_ids": args.tts_pass_token_ids,
|
||||
},
|
||||
},
|
||||
timeout=args.timeout,
|
||||
)
|
||||
save_response(response, args.output_dir)
|
||||
|
||||
|
||||
def parse_args():
|
||||
parser = FlexibleArgumentParser(description="AURA Omni online serving client")
|
||||
parser.add_argument("--host", default="localhost")
|
||||
parser.add_argument("--port", type=int, default=8091)
|
||||
parser.add_argument("--model", default=DEFAULT_MODEL)
|
||||
parser.add_argument("--audio-path", default=None, help="Audio file, URL, or data URL.")
|
||||
parser.add_argument("--video-path", default=None, help="Video file, URL, or data URL.")
|
||||
parser.add_argument(
|
||||
"--prompt",
|
||||
default="Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English.",
|
||||
)
|
||||
parser.add_argument("--modalities", default="text,audio")
|
||||
parser.add_argument("--output-dir", default="output_aura_omni_online")
|
||||
parser.add_argument(
|
||||
"--aura-system-prompt",
|
||||
default=(
|
||||
"You are receiving a live video stream where the final frame is the present moment. "
|
||||
"Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--tts-task-type", default="Base", choices=["Base", "CustomVoice"])
|
||||
parser.add_argument("--tts-language", default="English")
|
||||
parser.add_argument("--tts-speaker", default="Vivian")
|
||||
parser.add_argument("--tts-instruct", default="")
|
||||
parser.add_argument(
|
||||
"--tts-ref-audio",
|
||||
default=default_qwen3_tts_ref_audio_path(),
|
||||
help="Base-mode reference audio path/URL visible to server.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-ref-text",
|
||||
default=DEFAULT_QWEN3_TTS_REF_TEXT,
|
||||
help="Base-mode reference audio transcript.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-x-vector-only-mode",
|
||||
action="store_true",
|
||||
help="Use speaker embedding only for Base mode (disable ICL ref_text conditioning).",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--tts-pass-token-ids",
|
||||
action="store_true",
|
||||
help="Pass AURA-generated assistant token ids directly to Qwen3-TTS. Defaults to sending text.",
|
||||
)
|
||||
parser.add_argument("--timeout", type=float, default=600.0)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main(parse_args())
|
||||
@@ -0,0 +1,60 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
PORT="${PORT:-8091}"
|
||||
MODEL="${MODEL:-aurateam/AURA}"
|
||||
OUTPUT_DIR="${OUTPUT_DIR:-output_aura_omni_online}"
|
||||
TTS_PASS_TOKEN_IDS="${TTS_PASS_TOKEN_IDS:-false}"
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
VLLM_OMNI_ROOT="$(cd "${SCRIPT_DIR}/../../.." && pwd)"
|
||||
CLONE_REF_AUDIO="${VLLM_OMNI_ROOT}/tests/assets/qwen3_tts/clone_2.wav"
|
||||
CLONE_REF_TEXT="Okay. Yeah. I resent you. I love you. I respect you. But you know what? You blew it! And thanks to you."
|
||||
mkdir -p "$OUTPUT_DIR"
|
||||
|
||||
MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg"
|
||||
SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4"
|
||||
|
||||
request_body=$(cat <<EOF
|
||||
{
|
||||
"model": "$MODEL",
|
||||
"modalities": ["text", "audio"],
|
||||
"sampling_params_list": [
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42},
|
||||
{"temperature": 0.5, "top_p": 1.0, "top_k": -1, "max_tokens": 256, "seed": 42, "repetition_penalty": 1.0},
|
||||
{"temperature": 0.9, "top_k": 50, "max_tokens": 4096, "seed": 42, "detokenize": false, "repetition_penalty": 1.05, "stop_token_ids": [2150]},
|
||||
{"temperature": 0.0, "top_p": 1.0, "top_k": -1, "max_tokens": 65536, "seed": 42, "repetition_penalty": 1.0}
|
||||
],
|
||||
"additional_information": {
|
||||
"aura_system_prompt": "You are receiving a live video stream where the final frame is the present moment. Respond only when a response is needed. Otherwise output '<|silent|>'. Respond in English.",
|
||||
"tts_task_type": "Base",
|
||||
"tts_ref_audio": "file://${CLONE_REF_AUDIO}",
|
||||
"tts_ref_text": "${CLONE_REF_TEXT}",
|
||||
"tts_language": "English",
|
||||
"tts_speaker": "Vivian",
|
||||
"tts_instruct": "",
|
||||
"tts_pass_token_ids": ${TTS_PASS_TOKEN_IDS}
|
||||
},
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "audio_url", "audio_url": {"url": "$MARY_HAD_LAMB_AUDIO_URL"}},
|
||||
{"type": "video_url", "video_url": {"url": "$SAMPLE_VIDEO_URL"}},
|
||||
{"type": "text", "text": "Use the audio and video together to decide whether a reply is needed. If needed, respond briefly in English."}
|
||||
]
|
||||
}]
|
||||
}
|
||||
EOF
|
||||
)
|
||||
|
||||
response=$(curl -sS --retry 3 --retry-delay 3 --retry-connrefused \
|
||||
-X POST "http://localhost:${PORT}/v1/chat/completions" \
|
||||
-H "Content-Type: application/json" \
|
||||
-d "$request_body")
|
||||
|
||||
echo "$response" | jq '.choices[].message.content'
|
||||
|
||||
audio_b64=$(echo "$response" | jq -r '.choices[]?.message.audio.data // empty' | head -n 1)
|
||||
if [[ -n "$audio_b64" ]]; then
|
||||
echo "$audio_b64" | base64 -d > "${OUTPUT_DIR}/aura_omni_output.wav"
|
||||
echo "Audio saved to ${OUTPUT_DIR}/aura_omni_output.wav"
|
||||
fi
|
||||
+79
@@ -0,0 +1,79 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
MODEL="aurateam/AURA"
|
||||
SERVER_MODEL="aurateam/AURA"
|
||||
DEPLOY_CONFIG="/data/yrr/vllm-omni/vllm_omni/deploy/aura_omni.yaml"
|
||||
SERVER_PORT=8091
|
||||
GRADIO_PORT=7862
|
||||
SERVER_HOST="0.0.0.0"
|
||||
GRADIO_IP="127.0.0.1"
|
||||
GRADIO_SHARE=false
|
||||
|
||||
|
||||
while [[ $# -gt 0 ]]; do
|
||||
case "$1" in
|
||||
--model) MODEL="$2"; shift 2 ;;
|
||||
--server-model) SERVER_MODEL="$2"; shift 2 ;;
|
||||
--deploy-config) DEPLOY_CONFIG="$2"; shift 2 ;;
|
||||
--server-port) SERVER_PORT="$2"; shift 2 ;;
|
||||
--gradio-port) GRADIO_PORT="$2"; shift 2 ;;
|
||||
--server-host) SERVER_HOST="$2"; shift 2 ;;
|
||||
--gradio-ip) GRADIO_IP="$2"; shift 2 ;;
|
||||
--share) GRADIO_SHARE=true; shift ;;
|
||||
--help)
|
||||
echo "Usage: $0 [--model SERVED_MODEL_NAME] [--server-model MODEL_PATH] [--deploy-config YAML] [--server-port PORT] [--gradio-port PORT] [--share]"
|
||||
exit 0
|
||||
;;
|
||||
*) echo "Unknown option: $1"; exit 1 ;;
|
||||
esac
|
||||
done
|
||||
|
||||
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
|
||||
API_BASE="http://localhost:${SERVER_PORT}/v1"
|
||||
LOG_FILE="/tmp/aura_omni_vllm_${SERVER_PORT}.log"
|
||||
|
||||
cleanup() {
|
||||
echo "Shutting down..."
|
||||
[[ -n "${SERVER_PID:-}" ]] && kill "$SERVER_PID" 2>/dev/null || true
|
||||
[[ -n "${GRADIO_PID:-}" ]] && kill "$GRADIO_PID" 2>/dev/null || true
|
||||
}
|
||||
trap cleanup SIGINT SIGTERM EXIT
|
||||
|
||||
vllm serve "$SERVER_MODEL" \
|
||||
--omni \
|
||||
--host "$SERVER_HOST" \
|
||||
--port "$SERVER_PORT" \
|
||||
--deploy-config "$DEPLOY_CONFIG" \
|
||||
--served-model-name "$MODEL" \
|
||||
--trust-remote-code 2>&1 | tee "$LOG_FILE" &
|
||||
SERVER_PID=$!
|
||||
|
||||
echo "Waiting for server startup..."
|
||||
for _ in $(seq 1 600); do
|
||||
if grep -q "Application startup complete" "$LOG_FILE" 2>/dev/null; then
|
||||
break
|
||||
fi
|
||||
if ! kill -0 "$SERVER_PID" 2>/dev/null; then
|
||||
echo "vLLM server exited before startup completed"
|
||||
wait "$SERVER_PID" || true
|
||||
exit 1
|
||||
fi
|
||||
sleep 1
|
||||
done
|
||||
|
||||
# cd "$SCRIPT_DIR"
|
||||
GRADIO_CMD=(python gradio_demo.py --model "$MODEL" --api-base "$API_BASE" --ip "$GRADIO_IP" --port "$GRADIO_PORT")
|
||||
if [[ "$GRADIO_SHARE" == "true" ]]; then
|
||||
GRADIO_CMD+=(--share)
|
||||
fi
|
||||
"${GRADIO_CMD[@]}" &
|
||||
GRADIO_PID=$!
|
||||
|
||||
echo "vLLM server: http://${SERVER_HOST}:${SERVER_PORT}"
|
||||
echo "Gradio demo: http://${GRADIO_IP}:${GRADIO_PORT}"
|
||||
if [[ -n "${SERVER_PID:-}" ]]; then
|
||||
wait "$SERVER_PID" "$GRADIO_PID"
|
||||
else
|
||||
wait "$GRADIO_PID"
|
||||
fi
|
||||
@@ -1,404 +0,0 @@
|
||||
# BAGEL-7B-MoT
|
||||
|
||||
## Installation
|
||||
|
||||
Please refer to [README.md](../../../README.md)
|
||||
|
||||
## Architecture
|
||||
|
||||
BAGEL-7B-MoT is a Mixture-of-Transformers (MoT) model supporting both image generation and understanding. It offers two deployment topologies:
|
||||
|
||||
| Topology | Stages | Description |
|
||||
| :------- | :----- | :---------- |
|
||||
| **Two-stage** (default) | Stage 0 (Thinker, AR) + Stage 1 (DiT, Diffusion) | Thinker handles text/understanding via vLLM AR engine; DiT handles image generation. KV cache is transferred between stages. |
|
||||
| **Single-stage** | Stage 0 (DiT, Diffusion) only | The DiT stage contains a full LLM, ViT, VAE, and tokenizer internally. All modalities are handled within a single diffusion process. |
|
||||
|
||||
Both topologies support all four modalities: `text2img`, `img2img`, `img2text`, `text2text`.
|
||||
|
||||
> **Note**: These examples work with the default configuration on an **NVIDIA A100 (80GB)**. We also tested on dual **NVIDIA RTX 5000 Ada (32GB each)**. For dual-GPU setups, modify the deploy YAML to distribute stages across devices.
|
||||
|
||||
## Launch the Server
|
||||
|
||||
### Two-Stage (Default)
|
||||
|
||||
The default pipeline is auto-detected from the model. No extra flags needed:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/bagel
|
||||
bash run_server.sh
|
||||
|
||||
# Initialize each stage in a discrete isolated process terminal
|
||||
bash run_server_stage_cli.sh --stage 0
|
||||
bash run_server_stage_cli.sh --stage 1
|
||||
```
|
||||
|
||||
To use a custom deploy YAML, pass it via `--deploy-config`:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config /path/to/deploy_config.yaml
|
||||
```
|
||||
|
||||
See [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) for the default two-stage deploy configuration.
|
||||
|
||||
### Single-Stage
|
||||
|
||||
The DiT stage contains a full LLM, ViT, VAE, and tokenizer, so it can handle all modalities (text2img, img2img, img2text, text2text, think) without a separate Thinker stage:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml
|
||||
```
|
||||
|
||||
See [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) for configuration. The `pipeline: bagel_single_stage` field selects the single-stage topology from the pipeline registry.
|
||||
|
||||
### Tensor Parallelism (TP)
|
||||
|
||||
For larger models or multi-GPU environments, enable TP via CLI:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 --tensor-parallel-size 2
|
||||
```
|
||||
|
||||
Or set `tensor_parallel_size` per stage in a custom deploy YAML.
|
||||
|
||||
### VAE Patch Parallelism
|
||||
|
||||
[VAE Patch Parallelism](https://docs.vllm.ai/projects/vllm-omni/en/latest/user_guide/diffusion/parallelism/vae_patch_parallel.html) distributes Bagel VAE **decode/encode** across multiple GPUs by splitting latent tiles. It lowers **per-GPU peak memory during VAE decode**, which helps high-resolution `text2img` / `img2img` when VAE becomes a bottleneck.
|
||||
|
||||
**Scope for Bagel:**
|
||||
|
||||
| Topology | VAE patch parallel |
|
||||
| :------- | :----------------- |
|
||||
| **Single-stage** (DiT only) | Supported on stage 0 (`BagelPipeline` + `DistributedAutoEncoder`) |
|
||||
| **Two-stage** | Supported on **stage 1 (DiT)** only; stage 0 (Thinker) uses encoder-only VAE and is unrelated |
|
||||
|
||||
**Requirements:**
|
||||
|
||||
- `vae_patch_parallel_size > 1` and a distributed VAE (`DistributedAutoEncoder` on the DiT pipeline).
|
||||
- The DiT process group must have at least `vae_patch_parallel_size` ranks. In practice this means the diffusion stage `world_size` must be ≥ 2 (commonly `tensor_parallel_size=2` on that stage).
|
||||
- `vae_use_tiling` must be enabled. If you set `vae_patch_parallel_size > 1` and omit tiling, the registry auto-enables `vae_use_tiling` at startup.
|
||||
|
||||
VAE patch parallel **reuses the DiT process group** (`dit_group`); it does not create a separate VAE-only worker pool. It is not a substitute for single-GPU VAE tiling (`vae_pp=1`).
|
||||
|
||||
**Online serving (single-stage, 2 GPUs):**
|
||||
|
||||
```bash
|
||||
CUDA_VISIBLE_DEVICES=0,1 vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config vllm_omni/deploy/bagel_single_stage.yaml \
|
||||
--tensor-parallel-size 2 \
|
||||
--vae-patch-parallel-size 2 \
|
||||
--vae-use-tiling
|
||||
```
|
||||
|
||||
**Online serving (two-stage, VAE PP on DiT stage 1):** use a custom deploy YAML, for example:
|
||||
|
||||
```yaml
|
||||
stages:
|
||||
- stage_id: 0
|
||||
devices: "0"
|
||||
# Thinker (AR) — no VAE patch parallel here
|
||||
|
||||
- stage_id: 1
|
||||
devices: "0,1"
|
||||
vae_use_tiling: true
|
||||
parallel_config:
|
||||
tensor_parallel_size: 2
|
||||
vae_patch_parallel_size: 2
|
||||
```
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni --port 8091 \
|
||||
--deploy-config /path/to/bagel_vae_pp.yaml
|
||||
```
|
||||
|
||||
**Verify it is active** (check server logs at startup):
|
||||
|
||||
```text
|
||||
INFO ... vae_patch_parallel_size=2 requires vae_use_tiling; automatically enabling it.
|
||||
```
|
||||
|
||||
| CLI flag | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--vae-patch-parallel-size` | `1` | Number of DiT ranks used for VAE tile parallelism. Set to `2` or higher to enable. Should be ≤ DiT process group size (typically match `--tensor-parallel-size` on the diffusion stage). |
|
||||
| `--vae-use-tiling` | off | Enable VAE spatial tiling. Required for VAE patch parallel (auto-enabled when `vae_patch_parallel_size > 1`). |
|
||||
|
||||
#### Hybrid Sharded Data Parallel (HSDP)
|
||||
|
||||
For larger Bagel deployments on multiple GPUs, you can enable HSDP (Hybrid Sharded Data Parallel) by modifying the stage configuration (for example, [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml)). HSDP shards transformer weights across GPUs to reduce per-GPU memory usage.
|
||||
|
||||
1. **Enable HSDP**: Set `use_hsdp: true`.
|
||||
2. **Set shard size**: Set `hsdp_shard_size` to the number of GPUs used for sharding (for example, `4`).
|
||||
3. **Set replicate size**: Usually keep `hsdp_replicate_size: 1` unless you want replicated HSDP groups.
|
||||
4. **Set devices**: Specify the comma-separated GPU IDs used by the diffusion stage (for example, `"0,1,2,3"`).
|
||||
|
||||
Example configuration for HSDP across 4 GPUs:
|
||||
```yaml
|
||||
- stage_id: 1
|
||||
devices: "0,1,2,3"
|
||||
parallel_config:
|
||||
use_hsdp: true
|
||||
hsdp_shard_size: 4
|
||||
hsdp_replicate_size: 1
|
||||
```
|
||||
|
||||
### Multi-Node Deployment
|
||||
|
||||
Deploy each stage on a **separate node** for better resource utilization. Replace `<ORCHESTRATOR_IP>` with the actual IP address of your orchestrator node.
|
||||
|
||||
**1. Launch Stage 0 (Thinker / Orchestrator)** on the orchestrator node:
|
||||
|
||||
```bash
|
||||
# API server port for client requests: 8000
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
|
||||
--port 8000 \
|
||||
--stage-id 0 \
|
||||
--omni-master-address <ORCHESTRATOR_IP> \
|
||||
--omni-master-port 8091
|
||||
```
|
||||
|
||||
**2. Launch Stage 1 (DiT)** on the remote node in headless mode:
|
||||
|
||||
```bash
|
||||
vllm serve ByteDance-Seed/BAGEL-7B-MoT --omni \
|
||||
--stage-id 1 \
|
||||
--headless \
|
||||
--omni-master-address <ORCHESTRATOR_IP> \
|
||||
--omni-master-port 8091
|
||||
```
|
||||
|
||||
Or use the convenience script:
|
||||
|
||||
```bash
|
||||
# Terminal 1: Stage 0
|
||||
bash run_server_stage_cli.sh --stage 0
|
||||
|
||||
# Terminal 2: Stage 1
|
||||
bash run_server_stage_cli.sh --stage 1
|
||||
|
||||
# With extra args
|
||||
bash run_server_stage_cli.sh --stage 0 -- --tensor-parallel-size 2
|
||||
bash run_server_stage_cli.sh --stage 1 -- --gpu-memory-utilization 0.9
|
||||
```
|
||||
|
||||
**vllm serve arguments:**
|
||||
|
||||
| Argument | Description |
|
||||
| :------- | :---------- |
|
||||
| `--stage-id` | Which stage this process runs (0 = Thinker, 1 = DiT) |
|
||||
| `--headless` | Run without the API server (worker-only mode) |
|
||||
| `-oma` / `--omni-master-address` | Orchestrator master address |
|
||||
| `-omp` / `--omni-master-port` | Orchestrator master port |
|
||||
|
||||
> [!IMPORTANT]
|
||||
> **Startup Order**: Stage 0 (orchestrator) must be launched **before** Stage 1 (headless).
|
||||
> Stage 0 will appear to hang on startup until Stage 1 (worker) connects — this is expected behavior.
|
||||
|
||||
### Inter-Stage Connectors
|
||||
|
||||
When deploying stages across nodes, configure the connector type in the deploy YAML:
|
||||
|
||||
- **SharedMemoryConnector** (default): Used for single-node deployments. No explicit configuration needed.
|
||||
- **MooncakeTransferEngineConnector**: For multi-node setups with RDMA hardware. Defined in [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) under `connectors.rdma_connector`.
|
||||
|
||||
To use Mooncake, create a custom deploy YAML that binds `output_connectors` / `input_connectors` on each stage to the `rdma_connector` defined in the `connectors` section.
|
||||
|
||||
## Send Requests
|
||||
|
||||
```bash
|
||||
cd examples/online_serving/bagel
|
||||
```
|
||||
|
||||
### Text to Image (text2img)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A beautiful sunset over mountains" \
|
||||
--modality text2img \
|
||||
--output sunset.png \
|
||||
--steps 50
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>A beautiful sunset over mountains<|im_end|>"}]}],
|
||||
"modalities": ["image"],
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}'
|
||||
```
|
||||
|
||||
### Image to Image (img2img)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Make the cat stand up" \
|
||||
--modality img2img \
|
||||
--image-url /path/to/input.jpg \
|
||||
--output transformed.png
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
|
||||
|
||||
cat <<EOF > payload.json
|
||||
{
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "<|im_start|>Make the cat stand up<|im_end|>"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
|
||||
]
|
||||
}],
|
||||
"modalities": ["image"],
|
||||
"height": 512,
|
||||
"width": 512,
|
||||
"num_inference_steps": 50,
|
||||
"seed": 42
|
||||
}
|
||||
EOF
|
||||
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @payload.json
|
||||
```
|
||||
|
||||
### Image to Text (img2text)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "Describe this image in detail" \
|
||||
--modality img2text \
|
||||
--image-url /path/to/image.jpg
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
IMAGE_BASE64=$(base64 -w 0 cat.jpg)
|
||||
|
||||
cat <<EOF > payload.json
|
||||
{
|
||||
"messages": [{
|
||||
"role": "user",
|
||||
"content": [
|
||||
{"type": "text", "text": "<|im_start|>user\n<|image_pad|>\nDescribe this image in detail<|im_end|>\n<|im_start|>assistant\n"},
|
||||
{"type": "image_url", "image_url": {"url": "data:image/jpeg;base64,${IMAGE_BASE64}"}}
|
||||
]
|
||||
}],
|
||||
"modalities": ["text"]
|
||||
}
|
||||
EOF
|
||||
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d @payload.json
|
||||
```
|
||||
|
||||
### Text to Text (text2text)
|
||||
|
||||
**Python client:**
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "What is the capital of France?" \
|
||||
--modality text2text
|
||||
```
|
||||
|
||||
**curl:**
|
||||
|
||||
```bash
|
||||
curl http://localhost:8091/v1/chat/completions \
|
||||
-H "Content-Type: application/json" \
|
||||
-d '{
|
||||
"messages": [{"role": "user", "content": [{"type": "text", "text": "<|im_start|>user\nWhat is the capital of France?<|im_end|>\n<|im_start|>assistant\n"}]}],
|
||||
"modalities": ["text"]
|
||||
}'
|
||||
```
|
||||
|
||||
### Python Client Arguments
|
||||
|
||||
| Argument | Default | Description |
|
||||
| :------- | :------ | :---------- |
|
||||
| `--prompt` / `-p` | `A cute cat` | Text prompt |
|
||||
| `--output` / `-o` | `bagel_output.png` | Output file path |
|
||||
| `--server` / `-s` | `http://localhost:8091` | Server URL |
|
||||
| `--image-url` / `-i` | `None` | Input image URL or local path (img2img/img2text) |
|
||||
| `--modality` / `-m` | `text2img` | `text2img`, `img2img`, `img2text`, `text2text` |
|
||||
| `--height` | `512` | Image height in pixels |
|
||||
| `--width` | `512` | Image width in pixels |
|
||||
| `--steps` | `25` | Number of inference steps |
|
||||
| `--seed` | `42` | Random seed |
|
||||
| `--negative` | `None` | Negative prompt for CFG |
|
||||
|
||||
Example with custom parameters:
|
||||
|
||||
```bash
|
||||
python openai_chat_client.py \
|
||||
--prompt "A futuristic city" \
|
||||
--modality text2img \
|
||||
--height 768 \
|
||||
--width 768 \
|
||||
--steps 50 \
|
||||
--seed 42 \
|
||||
--negative "blurry, low quality"
|
||||
```
|
||||
|
||||
## Configuration Reference
|
||||
|
||||
### Deploy YAML Files
|
||||
|
||||
| File | Description |
|
||||
| :--- | :---------- |
|
||||
| [`bagel.yaml`](../../../vllm_omni/deploy/bagel.yaml) | Two-stage default (Thinker + DiT on GPU 0) |
|
||||
| [`bagel_single_stage.yaml`](../../../vllm_omni/deploy/bagel_single_stage.yaml) | Single-stage (DiT only) |
|
||||
|
||||
### Key Deploy YAML Fields
|
||||
|
||||
| Field | Scope | Description |
|
||||
| :---- | :---- | :---------- |
|
||||
| `pipeline` | top-level | Override auto-detected pipeline (e.g. `bagel_single_stage`) |
|
||||
| `stages[].stage_id` | per-stage | Stage identifier (0, 1, ...) |
|
||||
| `stages[].devices` | per-stage | GPU device IDs (e.g. `"0"`, `"0,1"`) |
|
||||
| `stages[].max_num_seqs` | per-stage | Maximum concurrent sequences |
|
||||
| `stages[].gpu_memory_utilization` | per-stage | Fraction of GPU memory to use |
|
||||
| `stages[].enforce_eager` | per-stage | Disable CUDA graphs |
|
||||
| `stages[].tensor_parallel_size` | per-stage | TP degree for this stage |
|
||||
| `stages[].parallel_config.vae_patch_parallel_size` | per-stage (DiT) | VAE tile parallelism degree (DiT stage only) |
|
||||
| `stages[].vae_use_tiling` | per-stage (DiT) | Enable VAE tiling (required for VAE patch parallel) |
|
||||
| `connectors` | top-level | Define available connector instances (SHM, Mooncake) |
|
||||
| `platforms` | top-level | Platform-specific overrides (e.g. `xpu`) |
|
||||
|
||||
## FAQ
|
||||
|
||||
- If you encounter OOM errors, try decreasing `max_model_len` or `gpu_memory_utilization` in the deploy YAML.
|
||||
|
||||
**Two-stage VRAM usage:**
|
||||
|
||||
| Stage | VRAM |
|
||||
| :---- | :--- |
|
||||
| Stage 0 (Thinker) | **15.04 GiB + KV Cache** |
|
||||
| Stage 1 (DiT) | **26.50 GiB** |
|
||||
| Total | **~42 GiB + KV Cache** |
|
||||
|
||||
**Single-stage VRAM usage:** The DiT loads the full model (~42 GiB) in one process.
|
||||
@@ -1,184 +0,0 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Bagel OpenAI-compatible chat client for image generation and multimodal tasks.
|
||||
|
||||
Usage:
|
||||
python openai_chat_client.py --prompt "A cute cat" --output output.png
|
||||
python openai_chat_client.py --prompt "Describe this image" --image-url https://example.com/image.png
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import base64
|
||||
from pathlib import Path
|
||||
|
||||
import requests
|
||||
|
||||
|
||||
def generate_image(
|
||||
prompt: str,
|
||||
server_url: str = "http://localhost:8091",
|
||||
image_url: str | None = None,
|
||||
height: int | None = None,
|
||||
width: int | None = None,
|
||||
steps: int | None = None,
|
||||
seed: int | None = None,
|
||||
negative_prompt: str | None = None,
|
||||
modality: str = "text2img", # "text2img" (default), "img2img", "img2text", "text2text"
|
||||
) -> bytes | str | None:
|
||||
"""Generate an image or text using the chat completions API.
|
||||
|
||||
Args:
|
||||
prompt: Text description or prompt
|
||||
server_url: Server URL
|
||||
image_url: URL or path to input image (for img2img/img2text)
|
||||
height: Image height in pixels
|
||||
width: Image width in pixels
|
||||
steps: Number of inference steps
|
||||
seed: Random seed
|
||||
negative_prompt: Negative prompt
|
||||
modality: Task modality hint
|
||||
|
||||
Returns:
|
||||
Image bytes (for image outputs) or Text string (for text outputs) or None if failed
|
||||
"""
|
||||
|
||||
# Construct Message Content
|
||||
content = [{"type": "text", "text": f"<|im_start|>{prompt}<|im_end|>"}]
|
||||
|
||||
if image_url:
|
||||
# Check if local file
|
||||
if Path(image_url).exists():
|
||||
with open(image_url, "rb") as f:
|
||||
b64_data = base64.b64encode(f.read()).decode("utf-8")
|
||||
final_image_url = f"data:image/jpeg;base64,{b64_data}"
|
||||
else:
|
||||
final_image_url = image_url
|
||||
|
||||
content.append({"type": "image_url", "image_url": {"url": final_image_url}})
|
||||
|
||||
messages = [{"role": "user", "content": content}]
|
||||
|
||||
# Build request payload with all parameters at top level
|
||||
# Note: vLLM ignores "extra_body", so we put parameters directly in the payload
|
||||
payload = {"messages": messages}
|
||||
|
||||
# Set output modalities at top level
|
||||
if modality == "text2img" or modality == "img2img":
|
||||
payload["modalities"] = ["image"]
|
||||
elif modality == "img2text" or modality == "text2text":
|
||||
payload["modalities"] = ["text"]
|
||||
|
||||
# Add generation parameters directly to payload
|
||||
if height is not None:
|
||||
payload["height"] = height
|
||||
if width is not None:
|
||||
payload["width"] = width
|
||||
if steps is not None:
|
||||
payload["num_inference_steps"] = steps
|
||||
if seed is not None:
|
||||
payload["seed"] = seed
|
||||
if negative_prompt:
|
||||
payload["negative_prompt"] = negative_prompt
|
||||
|
||||
# Send request
|
||||
try:
|
||||
print(f"Sending request to {server_url} with modality {modality}...")
|
||||
response = requests.post(
|
||||
f"{server_url}/v1/chat/completions",
|
||||
headers={"Content-Type": "application/json"},
|
||||
json=payload,
|
||||
timeout=300,
|
||||
)
|
||||
response.raise_for_status()
|
||||
data = response.json()
|
||||
|
||||
# Extract content - check ALL choices since server may return multiple
|
||||
# (e.g., text in choices[0], image in choices[1])
|
||||
choices = data.get("choices", [])
|
||||
|
||||
# First pass: look for image output in any choice
|
||||
for choice in choices:
|
||||
choice_content = choice.get("message", {}).get("content")
|
||||
|
||||
# Handle Image Output
|
||||
if isinstance(choice_content, list) and len(choice_content) > 0:
|
||||
first_item = choice_content[0]
|
||||
if isinstance(first_item, dict) and "image_url" in first_item:
|
||||
img_url_str = first_item["image_url"].get("url", "")
|
||||
if img_url_str.startswith("data:image"):
|
||||
_, b64_data = img_url_str.split(",", 1)
|
||||
return base64.b64decode(b64_data)
|
||||
|
||||
# Second pass: look for text output if no image found
|
||||
for choice in choices:
|
||||
choice_content = choice.get("message", {}).get("content")
|
||||
if isinstance(choice_content, str) and choice_content:
|
||||
return choice_content
|
||||
|
||||
print(f"Unexpected response format: {choices}")
|
||||
return None
|
||||
|
||||
except Exception as e:
|
||||
print(f"Error: {e}")
|
||||
return None
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description="Bagel multimodal chat client")
|
||||
parser.add_argument("--prompt", "-p", default="A cute cat", help="Text prompt")
|
||||
parser.add_argument("--output", "-o", default="bagel_output.png", help="Output file (for image results)")
|
||||
parser.add_argument("--server", "-s", default="http://localhost:8091", help="Server URL")
|
||||
|
||||
# Modality Control
|
||||
parser.add_argument("--image-url", "-i", type=str, help="Input image URL or local path")
|
||||
parser.add_argument(
|
||||
"--modality",
|
||||
"-m",
|
||||
default="text2img",
|
||||
choices=["text2img", "img2img", "img2text", "text2text"],
|
||||
help="Task modality",
|
||||
)
|
||||
|
||||
# Generation Params
|
||||
parser.add_argument("--height", type=int, default=512, help="Image height")
|
||||
parser.add_argument("--width", type=int, default=512, help="Image width")
|
||||
parser.add_argument("--steps", type=int, default=25, help="Inference steps")
|
||||
parser.add_argument("--seed", type=int, default=42, help="Random seed")
|
||||
parser.add_argument("--negative", help="Negative prompt")
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
print(f"Mode: {args.modality}")
|
||||
if args.image_url:
|
||||
print(f"Input Image: {args.image_url}")
|
||||
|
||||
result = generate_image(
|
||||
prompt=args.prompt,
|
||||
server_url=args.server,
|
||||
image_url=args.image_url,
|
||||
height=args.height,
|
||||
width=args.width,
|
||||
steps=args.steps,
|
||||
seed=args.seed,
|
||||
negative_prompt=args.negative,
|
||||
modality=args.modality,
|
||||
)
|
||||
|
||||
if result:
|
||||
if isinstance(result, bytes):
|
||||
# It's an image
|
||||
output_path = Path(args.output)
|
||||
output_path.write_bytes(result)
|
||||
print(f"Image saved to: {output_path}")
|
||||
print(f"Size: {len(result) / 1024:.1f} KB")
|
||||
elif isinstance(result, str):
|
||||
# It's text
|
||||
print("Response:")
|
||||
print(result)
|
||||
else:
|
||||
print("Failed to generate response")
|
||||
exit(1)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Loaded 100 of 782 files, more files were not shown because too many files have changed in this diff.
Show more
Reference in new issue
Block a user